import numpy as np
from qiskit import QuantumCircuit, Aer, execute
from qiskit.visualization import plot_bloch_multivector, array_to_latex
from grader import *
language(False)

phi_plus = QuantumCircuit(2)
phi_plus.h(0)
phi_plus.cnot(0, 1)
phi_plus.draw('mpl')

# Let's see the result
backend = Aer.get_backend('statevector_simulator')
result = execute(phi_plus, backend).result().get_statevector()
array_to_latex(result)

phi_minus = QuantumCircuit(2)

##### ==================================
# Write your solution in here.
phi_minus.h(0)
phi_minus.z(0)
phi_minus.cnot(0, 1)


##### ==================================

phi_minus.draw('mpl')

# This cell will tell you if what you did gives the correct result.
# Just take it as an indication, the tests do not take into consideration every possible option
# If you think that your answer is (in)correct and the grader says different, just ask me please
# Don't try to tamper with the code, please :)

ex1_grader(phi_minus)

psi_plus = QuantumCircuit(2)

##### ==================================
# Write your solution in here.
psi_plus.h(0)
psi_plus.x(1)
psi_plus.cnot(0, 1)


##### ==================================

psi_plus.draw('mpl')

ex2_grader(psi_plus)

psi_minus = QuantumCircuit(2)

##### ==================================
# Write your solution in here.
psi_minus.h(0)
psi_minus.x(1)
psi_minus.z(0)
psi_minus.cnot(0, 1)
##### ==================================

psi_minus.draw('mpl')

ex3_grader(psi_minus)

GHZ = QuantumCircuit(3)

##### ==================================
# Write your solution in here.
GHZ.h(0)
GHZ.cnot(0, 1)
GHZ.cnot(0, 2)
##### ==================================

GHZ.draw('mpl')

ex4_grader(GHZ)

Even = QuantumCircuit(3)

##### ==================================
# Write your solution in here.
Even.h(0)
Even.h(1)
Even.cnot(0, 2)
Even.cnot(1, 2)
##### ==================================

Even.draw('mpl')

ex5_grader(Even)

Odd = QuantumCircuit(3)

##### ==================================
# Write your solution in here.
Odd.h(0)
Odd.h(1)
Odd.x(2)
Odd.cnot(0, 2)
Odd.cnot(1, 2)
##### ==================================
Odd.draw('mpl')

ex6_grader(Odd)

def measure_z_axis(qc, qubit, cbit): #Measurements have to be saved in classical bits
    qc.measure(qubit, cbit)

def measure_x_axis(qc, qubit, cbit):
    ##### ==================================
    # Write your solution in here.
    qc.h(qubit)
    qc.measure(qubit, cbit)
    ##### ==================================

ex7_grader(measure_x_axis)

def expectation_value_single_qubit(qc, nshots=8092):
    # Run the given circuit that is already measuring the selected basis
    backend_shots = Aer.get_backend('qasm_simulator')
    counts = execute(qc, backend_shots, shots=nshots).result().get_counts() #Counts is a dictionary that contains the number of zeros and the number of ones measured
    
    # This is just to make things easier in your solution
    if '1' not in counts.keys():
        counts['1'] = 0
    if '0' not in counts.keys():
        counts['0'] = 0
        
    # Get the number of times that 0 and 1 were measured
    n_zeros = counts['0']
    n_ones  = counts['1']
    
    # Compute the probabilities
    p_0 = n_zeros/nshots
    p_1 = n_ones/nshots #OR 1-p_0

    expectation_value = 1*p_0+(-1)*p_1
    return expectation_value

# Measure <Z> over a state that you like
qc_z = QuantumCircuit(1,1)
measure_z_axis(qc_z, 0, 0)

print('<Z>=', expectation_value_single_qubit(qc_z))
qc_z.draw('mpl')

# Measure <X> over a state that you like
qc_x = QuantumCircuit(1,1)
measure_x_axis(qc_x, 0, 0)

print('<X>=', expectation_value_single_qubit(qc_x))
qc_x.draw('mpl')

##### ==================================
# Write your solution in here.
measure_strings = ['ZZ', 'ZX', 'XZ', 'XX']
# Create bell state circuit
qc = QuantumCircuit(2,2)
qc.h(0)
qc.cnot(0, 1)

# For each of the operators that we want to measure (ZZ, ZX, XZ and XX) create a new circuit (maybe using qc.copy())
# that measures along that axis. 
chsh_circuits = []
for string in measure_strings:
    qc1 = qc.copy()
    if string[0] == 'Z':
        measure_z_axis(qc1, 0, 0)
        
    if string[0] == 'X':
        measure_x_axis(qc1, 0, 0)
        
    if string[1] == 'Z':
        measure_z_axis(qc1, 1, 1)
        
    if string[1] == 'X':
        measure_x_axis(qc1, 1, 1)
        
    chsh_circuits.append(qc1)


##### ==================================

chsh_circuits[0].draw('mpl')

chsh_circuits[1].draw('mpl')

chsh_circuits[2].draw('mpl')

chsh_circuits[3].draw('mpl')

# I couldn't come up with tests that would not give the answer away if you looked at the code so you are on your own here!

# We can run all the circuits at the same time and then post process the data
nshots = 8092 # The more shots, the less error
backend_shots = Aer.get_backend('qasm_simulator')
counts_list = execute(chsh_circuits, backend_shots, shots=nshots).result().get_counts()
counts_list

##### ==================================
# Write your solution in here.
exp_dict = {}
# For each of the circuits, take the counts that you got from that circuit, compute the probability
# of each state and combine them together being very careful of which sign goes with which term
for i in range(len(chsh_circuits)):
    string = measure_strings[i] #ZZ, ZX, XZ and XX
    counts = counts_list[i]
    
    exp = 0.
    for element in counts:
        parity = (-1)**(int(element[0])+int(element[1]))
        #print(element, parity)
        exp += counts[element]*parity
    exp_dict[string] = exp/nshots
# You might want to take some inspiration from the function expectation_value_single_qubit
# If you are completely lost and don't know how to solve this, just ask

##### ==================================

# If you don't know if your answer is correct you can compute the results by hand
# (it's probably easier than doing this) and check if everything is working

##### ==================================
# Write your solution in here.
exp_AB = np.sqrt(1/2)*(exp_dict['ZZ']-exp_dict['XZ'])
exp_Ab = np.sqrt(1/2)*(exp_dict['ZX']-exp_dict['XX'])
exp_aB = np.sqrt(1/2)*(exp_dict['ZZ']+exp_dict['XZ'])
exp_ab = np.sqrt(1/2)*(exp_dict['ZX']+exp_dict['XX'])

CHSH = exp_AB - exp_Ab + exp_aB + exp_ab
##### ==================================
print('Your result is <CHSH>=', CHSH)

print('The correct result is <CHSH>=', 2*np.sqrt(2))

##### ==================================
# Write your solution in here.
Can_entanglement_be_teleported = True #Write True or False :D
##### ==================================

# There is no grader for this question hehe

##### ==================================
# Write your solution in here.
Where_did_the_bell_pair_go = 'To the last two qubits'
##### ==================================

# There is no grader for this question hehe

##### ==================================
# Write your solution in here.

##### ==================================

import qiskit.tools.jupyter
%qiskit_version_table