import numpy as np
from pyquil import get_qc, Program
from pyquil.gates import RX, RY, S, T, Z, CNOT, MEASURE
from pyquil.paulis import PauliTerm, PauliSum, sZ
from pyquil.noise import pauli_kraus_map, append_kraus_to_gate
from typing import List, Union
from collections import Counter
from matplotlib import pyplot as plt
from scipy import optimize
import mitiq
from mitiq import zne
from mitiq.zne.scaling.folding import fold_gates_at_random

backend = get_qc("2q-qvm")

program = Program()
theta = program.declare("theta", memory_type="REAL")
program += RX(theta, 0)
program += T(0)
program += CNOT(1, 0)
program += S(0)
program += Z(0)

def add_noise_to_circuit(quil_prog):
    """Define pyQuil gates with a custom noise model via Kraus operators:
        1. Generate Kraus operators at given survival probability
        2. Append Kraus operators to the gate matrices
        3. Add custom gates to circuit

    Args:
      quil_prog: the pyQuil quantum program to which the noise model will be added
      
    Returns:
      A quantum program with depolarizing noise on the static gates.
    """
    prob = 0.8
    num_qubits = 1
    d = 4 ** num_qubits
    d_sq = d ** 2
    
    kraus_list = [(1 - prob) / d] * d
    kraus_list[0] += prob
    kraus_ops = pauli_kraus_map(kraus_list)
    
    k_list = [(1 - prob) / d_sq] * d_sq
    k_list[0] += prob
    k_ops = pauli_kraus_map(k_list)
    
    T_gate = np.array([[1, 0], [0, np.exp(1j * np.pi / 4)]])
    CNOT_gate = np.block(
        [[np.eye(2), np.zeros((2, 2))], [np.zeros((2, 2)), np.flip(np.eye(2), 1)]]
    )
    S_gate = np.array([[1, 0], [0, 1j]])
    Z_gate = np.array([[1, 0], [0, -1]])
    
    quil_prog.define_noisy_gate("T", [0], append_kraus_to_gate(kraus_ops, T_gate))
    quil_prog.define_noisy_gate("CNOT", [1, 0], append_kraus_to_gate(k_ops, CNOT_gate))
    quil_prog.define_noisy_gate("S", [0], append_kraus_to_gate(kraus_ops, S_gate))
    quil_prog.define_noisy_gate("Z", [0], append_kraus_to_gate(kraus_ops, Z_gate))
    
    return quil_prog

hamiltonian = sZ(0)
pauli_sum = PauliSum([hamiltonian])

for j, term in enumerate(pauli_sum.terms):
    meas_basis_change = Program()
    marked_qubits = []
    for index, gate in term:
        marked_qubits.append(index)
        if gate == "X":
            meas_basis_change.inst(RY(-np.pi / 2, index))
        elif gate == "Y":
            meas_basis_change.inst(RX(np.pi / 2, index))
        program += meas_basis_change
    readout_qubit = program.declare("ro", "BIT", max(marked_qubits) + 1)

samples = 3000
program.wrap_in_numshots_loop(samples)

def executor(
    theta,
    backend,
    readout_qubit,
    samples: int,
    pauli_sum: Union[PauliSum, PauliTerm, np.ndarray],
    pyquil_prog: Program,
) -> float:
    """
    Compute the expectation value of pauli_sum over the distribution generated from 
    pyquil_prog.
    """
    noisy = pyquil_prog.copy()
    noisy += [
        MEASURE(qubit, r) for qubit, r in zip(list(range(max(marked_qubits) + 1)), readout_qubit)
    ]
    noisy = add_noise_to_circuit(noisy)
    expectation = 0.0
    pauli_sum = PauliSum([pauli_sum])
    for j, term in enumerate(pauli_sum.terms):
        qubits_to_measure = []
        for index, gate in term:
            qubits_to_measure.append(index)
            meas_outcome = expectation_from_sampling(
                theta, noisy, qubits_to_measure, backend, samples
            )
            expectation += term.coefficient * meas_outcome
    return expectation.real

def expectation_from_sampling(
    theta, executable: Program, marked_qubits: List[int], backend, samples: int
) -> float:
    """Calculate the expectation value of the Zi operator where i ranges over all 
    qubits given in marked_qubits.
    """
    bitstring_samples = backend.run(
        executable.write_memory(region_name="theta", value=theta)
    ).readout_data.get("ro")
    bitstring_tuples = list(map(tuple, bitstring_samples))

    freq = Counter(bitstring_tuples)

    
    exp_val = 0
    for bitstring, count in freq.items():
        bitstring_int = int("".join([str(x) for x in bitstring[::-1]]), 2)
        if parity_even_p(bitstring_int, marked_qubits):
            exp_val += float(count) / samples
        else:
            exp_val -= float(count) / samples
    return exp_val

def parity_even_p(state, marked_qubits):
    mask = 0
    for q in marked_qubits:
        mask |= 1 << q
    return bin(mask & state).count("1") % 2 == 0

thetas = np.linspace(0, 2 * np.pi, 51)
results = []

for theta in thetas:
    results.append(executor(theta, backend, readout_qubit, samples, hamiltonian, program))

init_angle = [3.0]

res = optimize.minimize(
    executor,
    init_angle,
    args=(backend, readout_qubit, samples, hamiltonian, program),
    method="Nelder-Mead",
    options={"xatol": 1.0e-3, "fatol": 1.0e-2},
)
print(res)

def mitigated_expectation(
    thetas, backend, readout_qubit, samples, pauli_sum, executable: Program, factory
) -> float:
    """
    This function is the ZNE-wrapped executor, which outputs the error-mitigated 
    expectation value. 

    Args:
      thetas: the input parameter for the optimization
      backend: the quantum computer that runs the quantum program
      readout_qubit: declared memory for the readout
      samples: number of times the experiment (or simulation) will be run
      pauli_sum: the Hamiltonian expressed as 
      executable: the pyQuil quantum program
      factory: factory object containing the type of inference and scaling parameters

    Returns:
      The error-mitigated expectation value as a float.
    """
    mitigated_exp = zne.execute_with_zne(
        executable,
        lambda p: executor(thetas, backend, readout_qubit, samples, pauli_sum, p),
        factory=factory,
        scale_noise=fold_gates_at_random,
    )
    return mitigated_exp

fac = mitiq.zne.inference.LinearFactory(scale_factors=[1.0, 3.0])

results_zne = []
for theta in thetas:
    results_zne.append(
        mitigated_expectation(theta, backend, readout_qubit, samples, hamiltonian, program, fac)
    )

_ = plt.figure()
_ = plt.plot(thetas, np.cos(thetas), "o-", label="Ideal landscape")
_ = plt.plot(thetas, results, "o-", label="Noisy landscape")
_ = plt.plot(thetas, results_zne, "o-", label="Mitigated landscape")
_ = plt.xlabel(r"$\theta$", fontsize=18)
_ = plt.ylabel(r"$\langle \Psi(\theta) | Z | \Psi(\theta) \rangle$", fontsize=18)
_ = plt.legend()
_ = plt.title("Mitigated Energy Landscape")
plt.show()

res_zne = optimize.minimize(
    mitigated_expectation,
    init_angle,
    args=(backend, readout_qubit, samples, hamiltonian, program, fac),
    method="Nelder-Mead",
    options={"xatol": 1.0e-3, "fatol": 1.0e-2},
)
print(res_zne)

mitiq.about()