import numpy as np

from mitiq import cdr, Observable, PauliString, Executor

import cirq
from cirq.contrib.qasm_import import circuit_from_qasm
from cirq.circuits.qasm_output import QasmUGate

from qiskit import *
from qiskit import transpile, assemble, IBMQ
from qiskit.visualization import *
from qiskit  import Aer, QuantumCircuit
from qiskit.circuit import ParameterVector

pi = np.pi
feature_dim = 2
reps = 2
duplicates = 1
feature_map_circuit = QuantumCircuit(feature_dim, name='feature_map_circuit')
# x = ParameterVector('x', feature_dim * reps)
x = [0.1, 0.2]

for i in range(reps):
    for i in range(feature_dim):
        feature_map_circuit.h(i)

    feature_map_circuit.rz(2.0*x[0], 0)
    feature_map_circuit.rz(2.0*x[1], 1)
    feature_map_circuit.cnot(0, 1)
    feature_map_circuit.rz(2.0*(pi - x[0])*(pi - x[1]), 1)
    feature_map_circuit.cnot(0, 1)

feature_map_circuit.draw()

qasm_feature_map = feature_map_circuit.qasm(filename='feature_map.qasm')
print(qasm_feature_map)
cirq_feature_map = circuit_from_qasm(qasm_feature_map)
cirq_3_feature_map = cirq_feature_map * 3
print(cirq_3_feature_map)

from mitiq.interface.mitiq_cirq import compute_density_matrix

compute_density_matrix(cirq_feature_map).round(3)


obs = Observable(PauliString("ZZ"))
print(obs)

def simulate(circuit: cirq.Circuit) -> np.ndarray:
    return compute_density_matrix(circuit, noise_level=(0.0,))

simulate(cirq_feature_map).round(3)

ideal_measurement = obs.expectation(cirq_feature_map, simulate).real
print("ideal_measurement = ",ideal_measurement)

unmitigated_measurement = obs.expectation(cirq_feature_map, compute_density_matrix).real
print("unmitigated_measurement = ", unmitigated_measurement)

mitigated_measurement = cdr.execute_with_cdr(
    cirq_feature_map,
    compute_density_matrix,
    observable=obs,
    simulator=simulate,
    seed=0,
).real
print("mitigated_measurement = ", mitigated_measurement)


from mitiq.zne import scaling

vncdr = cdr.execute_with_cdr(
    cirq_feature_map,
    compute_density_matrix,
    observable=obs,
    simulator=simulate,
    seed=0,
    scale_factors=(1.0, 3.0),
).real

print("vncdr = ", vncdr)


error_unmitigated = abs(unmitigated_measurement-ideal_measurement)
error_mitigated = abs(mitigated_measurement-ideal_measurement)
error_vn_mitigated = abs(vncdr-ideal_measurement)

print("Error (unmitigated):", error_unmitigated)
print("Error (mitigated with CDR):", error_mitigated)
print("Error (mitigated with VNCDR):", error_vn_mitigated)

print("Relative error (unmitigated):", (error_unmitigated/ideal_measurement))
print("Relative error (mitigated with CDR):",
      error_mitigated/ideal_measurement)
print("Relative error (mitigated with VNCDR):",
      error_vn_mitigated/ideal_measurement)

print(
    f"Error reduction with CDR: {(error_unmitigated-error_mitigated)/error_unmitigated :.1%}.")
print(
    f"Error reduction with VNCDR: {(error_unmitigated-error_vn_mitigated)/error_unmitigated :.1%}.")


num_qubits = 2
ansatz_circuit = QuantumCircuit(num_qubits, name='ansatz_circuit')
reps = 1
# params = ParameterVector('θ', num_qubits + num_qubits * reps)
params = [0.7, 0.3, 0.5, 1.2]
counter = 0

for i in range(num_qubits):
    ansatz_circuit.ry(params[counter], i)
    counter += 1

for j in range(reps):
    for k in range(num_qubits-1):
        ansatz_circuit.cnot(k, k+1)
    for m in range(num_qubits):
        ansatz_circuit.ry(params[counter], m)
        counter += 1

ansatz_circuit.draw()

qasm_ansatz = ansatz_circuit.qasm(filename='ansatz.qasm')
print(qasm_ansatz)
cirq_ansatz = circuit_from_qasm(qasm_ansatz) * 1
print(cirq_ansatz)


compute_density_matrix(cirq_ansatz).round(3)

simulate(cirq_ansatz).round(3)

ideal_measurement = obs.expectation(cirq_ansatz, simulate).real
print("ideal_measurement = ",ideal_measurement)

unmitigated_measurement = obs.expectation(cirq_ansatz, compute_density_matrix).real
print("unmitigated_measurement = ", unmitigated_measurement)

mitigated_measurement = cdr.execute_with_cdr(
    cirq_ansatz,
    compute_density_matrix,
    observable=obs,
    simulator=simulate,
    seed=0,
).real
print("mitigated_measurement = ", mitigated_measurement)


from mitiq.zne import scaling

vncdr = cdr.execute_with_cdr(
    cirq_ansatz,
    compute_density_matrix,
    observable=obs,
    simulator=simulate,
    seed=0,
    scale_factors=(1.0, 3.0),
).real

print("vncdr = ", vncdr)


error_unmitigated = abs(unmitigated_measurement-ideal_measurement)
error_mitigated = abs(mitigated_measurement-ideal_measurement)
error_vn_mitigated = abs(vncdr-ideal_measurement)

print("Error (unmitigated):", error_unmitigated)
print("Error (mitigated with CDR):", error_mitigated)
print("Error (mitigated with VNCDR):", error_vn_mitigated)

print("Relative error (unmitigated):", (error_unmitigated/ideal_measurement))
print("Relative error (mitigated with CDR):",
      error_mitigated/ideal_measurement)
print("Relative error (mitigated with VNCDR):",
      error_vn_mitigated/ideal_measurement)

print(
    f"Error reduction with CDR: {(error_unmitigated-error_mitigated)/error_unmitigated :.1%}.")
print(
    f"Error reduction with VNCDR: {(error_unmitigated-error_vn_mitigated)/error_unmitigated :.1%}.")
