import numpy as np
import qiskit
from qiskit.circuit.library import RXGate
from qiskit.providers.aer.noise import NoiseModel, ReadoutError, depolarizing_error, coherent_unitary_error
from scipy.stats import norm
from scipy.optimize import curve_fit
import matplotlib.pylab as plt
import qc_grader



# Create circuit to generate an equal superposition of quantum states
# Set the number of qubits
qubits_count = 1
# Set the number of bits
cbits_count = 1
# Instantiate a quantum register
quantum_register = qiskit.QuantumRegister(qubits_count)
# Instantiate a classical register
classical_register = qiskit.ClassicalRegister(cbits_count)
# Instantiate a quantum circuit
circuit = qiskit.QuantumCircuit(quantum_register, classical_register)

# Add a Hadamard gate and measurement
###INSERT CODE BELOW THIS LINE
circuit.h(0)
circuit.measure(0,0)

###DO NOT EDIT BELOW THIS LINE



from qc_grader.challenges.qgss_2022 import grade_lab3_ex1

grade_lab3_ex1(circuit)

# Run the quantum circuit on a simulator backend which returns measured shots in the Z basis
shots_backend = qiskit.Aer.get_backend('qasm_simulator')

# Fix number of times to repeat the experiment
num_shots = 10

# Create a job, run on the simulator, extract the results and get the counts dictionary
###INSERT CODE BELOW THIS LINE
job = shots_backend.run(circuit, shots = num_shots).result()
counts = job.get_counts()

###DO NOT EDIT BELOW THIS LINE

print(counts)




from qc_grader.challenges.qgss_2022 import  grade_lab3_ex2

grade_lab3_ex2(counts)

# Set different number of shots
num_shots_list = [10, 100, 1000, 10000]

# For each number of shots, extract the probability of |0> from the experimental counts
p0_list = []
for num_shots in num_shots_list:
    
    # Create a job, run on the simulator, extract the results and get the counts dictionary
    ###INSERT CODE BELOW THIS LINE
    job = shots_backend.run(circuit, shots = num_shots).result()
    counts = job.get_counts()
    counts_0 = counts['0']
    prob = counts_0/num_shots
    p0_list.append(prob)

    ###DO NOT EDIT BELOW THIS LINE   


for num_shots, p0 in zip(num_shots_list, p0_list):
    print(f"The probability of observing the 0 state with {num_shots} shots is: {p0}")


from qc_grader.challenges.qgss_2022 import  grade_lab3_ex3

grade_lab3_ex3(p0_list)

# Set the number of shots
num_shots = 10
# Set the number of times to repeat the experiment
iterations = 10000
# Set the number of bins for the histogram
bin_count = 100

# For each iteration, create job, run the experiment, 
# collect results and calculate probability of |0>
p0_list = []
for _ in range(iterations):

    job = shots_backend.run(circuit, shots=num_shots)

    result = job.result()

    if '0' in result.get_counts().keys():
        p0 = result.get_counts()['0']/num_shots
    else:
        p0 = 1 - result.get_counts()['1']/num_shots

    p0_list.append(p0)

# Calculate a histogram from the results
yhist, xhist = np.histogram(p0_list, bins=bin_count)
# Restrict horizontal axis to 0-1
x_01 = np.array([xhist[n] + (xhist[n + 1] - xhist[n])/2 for n in range(len(xhist)-1)])

# Find x-points for which the histogram has more than zero counts
xh = np.where(yhist > 0)[0]
# Get the y-points of the corresponding x ones
yh = yhist[xh]
# Get the x-points
x_01_h = x_01[xh]

# Plot
plt.scatter(x_01_h, yh, color = 'red')
plt.title(f"Using {num_shots} shots to estimate probability")
plt.xlim((0, 1))
plt.xlabel(r'probability of $\vert 0 \rangle$')
plt.ylabel('counts')
plt.show()

# Define a Gaussian function for the fit
def gaussian(x, a, mean, sigma):
    return a * (1/(sigma*np.sqrt(2*np.pi)))*np.exp(-((x - mean)**2 / (2 * sigma**2)))

# Set initial estimates for the fit: the maximum height of the histogram, the theoretical
# average and the theoretical standard deviation
height_estimate = np.max(yh)
mu_estimate = 0.5
sigma_estimate = np.sqrt(mu_estimate*(1-mu_estimate)/num_shots)

# Define list to store the estimated values of the average (mu_sigma[0]) and standard deviation (mu_sigma[1])
mu_sigma = []

# Use the curve_fit function to fit the experimental data using the definition of the Gaussian function
# which will return the estimate of the parameters.
###INSERT CODE BELOW THIS LINE
popt, pcov = curve_fit(gaussian, x_01_h, yh, [height_estimate, mu_estimate, sigma_estimate])
# (x_01_h, 1, mu_estimate, sigma_estimate)
mu_sigma.append(popt[1])
mu_sigma.append(popt[2])

# ###DO NOT EDIT BELOW THIS LINE   

print("The mean is ", mu_sigma[0])
print("The standard deviation is ", mu_sigma[1])

# Plot experimental data and the fit
plt.scatter(x_01_h, yh, label = 'data', color = 'red')
plt.plot(x_01, gaussian(x_01, *popt), label = 'fit', linestyle = '--')
plt.title(f"Using {num_shots} shots to estimate probability")
plt.xlim((0, 1))
plt.xlabel(r'probability of $\vert 0 \rangle$')
plt.ylabel('counts')
plt.legend()
plt.show()

from qc_grader.challenges.qgss_2022 import grade_lab3_ex4

grade_lab3_ex4(mu_sigma)

# Set different number of shots
num_shots_list = [10, 100, 1000,10000]
# Set the number of times to repeat the experiment
iterations = 10000
# Set the number of bins for the histogram
bin_count = 100

# Repeat the protocol above for different number of shots to see how the parameters of the
# estimated Gaussian change with the number of shots
popt_list = []
gaussian_curve_list, x_01_list = [], []
for num_shots in num_shots_list:

    p0_list = []
    for _ in range(iterations):

        job = shots_backend.run(circuit, shots=num_shots)

        result = job.result()

        if '0' in result.get_counts().keys():
            p0 = result.get_counts()['0']/num_shots
        else:
            p0 = 1 - result.get_counts()['1']/num_shots

        p0_list.append(p0)



    yhist, xhist = np.histogram(p0_list, bins=bin_count)
    norm_y = np.sum(yhist)
    yhist = yhist/iterations
    x_01 = np.array([xhist[n] + (xhist[n + 1] - xhist[n])/2 for n in range(len(xhist)-1)])


    xh = np.where(yhist > 0)[0]
    yh = yhist[xh]
    x_01_h = x_01[xh]

    height_estimate = np.max(yh)
    mu_estimate = p0
    sigma_estimate = np.sqrt(mu_estimate*(1-mu_estimate)/num_shots)


    popt, pcov = curve_fit(gaussian, x_01_h, yh, [height_estimate, mu_estimate, sigma_estimate])

    gaussian_curve = gaussian(x_01, *popt)
    gaussian_curve_list.append(gaussian_curve)
    x_01_list.append(x_01)
    popt_list.append(popt)


    plt.scatter(x_01_h, yh, label = 'data', color = 'red')
    plt.plot(x_01, gaussian_curve, label = 'fit', linestyle = '--')
    plt.title(f"Using {num_shots} shots to estimate probability")
    plt.xlim((0, 1))
    plt.xlabel(r'probability of $\vert 0 \rangle$')
    plt.ylabel('counts')
    plt.legend()
    plt.show()

# Define list which will contain the experimental standard deviations 
sigma_exp_list = []

# Append the standard deviations from the experiments with different number of shots to sigma_exp_list
###INSERT CODE BELOW THIS LINE
for popt in popt_list:
    sigma_exp_list.append(popt[2])

###DO NOT EDIT BELOW THIS LINE   


# Define theoretical sigma
p=0.5
sigma_th_list = [np.sqrt(p*(1-p)/num_shots) for num_shots in num_shots_list]

for sigma_exp, sigma_th in zip(sigma_exp_list, sigma_th_list):
    print(f"The variance from experiment is {sigma_exp} while the theoretical one is {sigma_th}")

from qc_grader.challenges.qgss_2022 import grade_lab3_ex5

grade_lab3_ex5(sigma_exp_list)

# Set number of shots
num_shots_measurement = 10000

# Create an empty noise model
readout_noise_model = NoiseModel()

# Measurement miss-assignement probabilities
p0given1 = 0.3 # Probability of measuuring |0> given the state is |1>
p1given0 = 0.05 # Probability of measuring |1> given the state is |0>
readout_error = ReadoutError([[1 - p1given0, p1given0], [p0given1, 1 - p0given1]]) # Define readout error
readout_noise_model.add_readout_error(readout_error, [0]) # Add error to noise model


# Calculate probability of seeing |0> given that |0> is prepared
circuit_0 = qiskit.QuantumCircuit(quantum_register, classical_register)

# Add measurement
###INSERT CODE BELOW THIS LINE
circuit_0.measure(0,0)

###DO NOT EDIT BELOW THIS LINE


from qc_grader.challenges.qgss_2022 import grade_lab3_ex6

grade_lab3_ex6(circuit_0)

# Run the experiment
job = shots_backend.run(circuit_0, shots=num_shots_measurement, noise_model=readout_noise_model)
# Retrieve results
result = job.result()

# Calculate probability of |0>
if '0' in result.get_counts().keys():
    p0_0 = result.get_counts()['0']/num_shots_measurement
else:
    p0_0 = 1 - result.get_counts()['1']/num_shots_measurement

p1_0 = 1 - p0_0

print('Probability of measuring |0> when |0> is prepared:', p0_0)
print('Probability of measuring |1> when |0> is prepared:', p1_0)

# Calculate probability of seeing |1> given that |1> is prepared
circuit_1 = qiskit.QuantumCircuit(quantum_register, classical_register)

# Add X gate and measurement
###INSERT CODE BELOW THIS LINE
circuit_1.x(0)
circuit_1.measure(0,0)

###DO NOT EDIT BELOW THIS LINE





from qc_grader.challenges.qgss_2022 import grade_lab3_ex7

grade_lab3_ex7(circuit_1)

# Run the experiment
job = shots_backend.run(circuit_1, shots=num_shots_measurement, noise_model=readout_noise_model)
# Retrieve results
result = job.result()

# Calculate probability of |0>
if '0' in result.get_counts().keys():
    p1_1 = 1 - result.get_counts()['0']/num_shots_measurement
else:
    p1_1 = result.get_counts()['1']/num_shots_measurement

p0_1 = 1 - p1_1


print('Probability of measuring |0> when |1> is prepared:', p0_1)
print('Probability of measuring |1> when |1> is prepared:', p1_1)


# Define the confusion matrix from the probabilities found above
###INSERT CODE BELOW THIS LINE
confusion_matrix = np.array([[p0_0, p0_1], [p1_0, p1_1]])

###DO NOT EDIT BELOW THIS LINE


print("Confusion matrix:")
print(confusion_matrix)



from qc_grader.challenges.qgss_2022 import grade_lab3_ex8

# The grader expects a np.array object (not list)
grade_lab3_ex8(confusion_matrix)

# Run experiments
job = shots_backend.run(circuit, shots=num_shots_measurement, noise_model=readout_noise_model)
# Extract results
result = job.result()

# Calculate probability of |0>
if '0' in result.get_counts().keys():
    p0_noisy = result.get_counts()['0']/num_shots_measurement
    p1_noisy = 1 - p0_noisy
else:
    p0_noisy = 1 - result.get_counts()['1']/num_shots_measurement
    p1_noisy = 1 - p0_noisy

p_vector_noisy = np.array([p0_noisy, p1_noisy])
print("Noisy probability of |0>: ", p0_noisy)
print("Noisy probability of |1>: ", p1_noisy)

# Invert the confusion matrix
inverse_confusion_matrix = np.linalg.inv(confusion_matrix)

# Mitigate the counts using the inverse of the confusion matrix
###INSERT CODE BELOW THIS LINE
p_vector_mitigated = (inverse_confusion_matrix @ p_vector_noisy)

###DO NOT EDIT BELOW THIS LINE


print("Mitigated probability of |0>: ", p_vector_mitigated[0])
print("Mitigated probability of |1>: ", p_vector_mitigated[1])

from qc_grader.challenges.qgss_2022 import grade_lab3_ex9

# The grader expects a np.array object (not list)
grade_lab3_ex9(p_vector_mitigated)

# Run ideal experiments
job = shots_backend.run(circuit, shots=num_shots_measurement)
# Extract results
result = job.result()

# Calculate probability of |0>
if '0' in result.get_counts().keys():
    p0_ideal = result.get_counts()['0']/num_shots_measurement
    p1_ideal = 1 - p0_ideal
else:
    p0_ideal = 1 - result.get_counts()['1']/num_shots_measurement
    p1_ideal = 1 - p0_ideal

print("Ideal probability of |0>: ", p0_ideal)
print("Ideal probability of |1>: ", p1_ideal)


# Set number of shots
num_shots_coherent = 100

# Create an empty noise model
rx_overrotation_noise_model = NoiseModel()

# Construct a 1 qubit over-rotation of the RX gate
epsilon = np.pi/5 # over rotation amount
epsilon_rotation = RXGate(epsilon).to_matrix() # get matrix representation

# Create coherent error
over_rotation = coherent_unitary_error(epsilon_rotation)

# Add error to error model
rx_overrotation_noise_model.add_quantum_error(over_rotation, ['rx'], qubits = [0])

# Set value of rotation
theta = np.pi
# Instantiate a quantum circuit
circuit = qiskit.QuantumCircuit(quantum_register, classical_register)

# Add a parametrized RX rotation and bind the value of the parameter. (By default, parameter binding is not an in-place operation)
# Then measure the qubit and calculate probability of seeing |0> after rx(np.pi) 
###INSERT CODE BELOW THIS LINE
circuit.rx(theta, 0)
circuit.measure(0,0)

###DO NOT EDIT BELOW THIS LINE



from qc_grader.challenges.qgss_2022 import grade_lab3_ex10

grade_lab3_ex10(circuit)

# Run experiments
job = shots_backend.run(circuit, shots=num_shots_coherent, noise_model=rx_overrotation_noise_model)
# Extract results
result = job.result()

# Calculate the probability of |0>
if '0' in result.get_counts().keys():
    p0 = result.get_counts()['0']/num_shots_coherent
else:
    p0 = 1 - result.get_counts()['1']/num_shots_coherent

print("Noisy probability of |0> after coherent error in idling: ", p0)
print("Noisy probability of |1> after coherent error in idling: ", 1-p0)

# Define number of shots
num_shots_incoherent = 10000

# Create an empty noise model
depolarizing_noise_model = NoiseModel()

# Define a depolarizing error on the identity gate for qubit zero and add it to the noise model
###INSERT CODE BELOW THIS LINE
lamb = float(5/100)
depolarizing = depolarizing_error(lamb, 1)
depolarizing_noise_model.add_quantum_error(depolarizing, ['id'], qubits = [0])

###DO NOT EDIT BELOW THIS LINE



from qc_grader.challenges.qgss_2022 import grade_lab3_ex11

grade_lab3_ex11(depolarizing_noise_model)

# Create a circuit with an idle instruction
circuit_idle = qiskit.QuantumCircuit(quantum_register, classical_register)
circuit_idle.id(0)
circuit_idle.measure(0,0)

# Run experiments
job = shots_backend.run(circuit_idle, shots=num_shots_incoherent, noise_model=depolarizing_noise_model)
# Extract results
result = job.result()

# Calculate probability of seeing |0> after idle time
if '0' in result.get_counts().keys():
    p0 = result.get_counts()['0']/num_shots_incoherent
else:
    p0 = 1 - result.get_counts()['1']/num_shots_incoherent

print("Noisy probability of |0> after incoherent error in idling: ", p0)
print("Noisy probability of |1> after incoherent error in idling: ", 1-p0)


