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.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.draw('mpl')

ex2_grader(psi_plus)

psi_minus = QuantumCircuit(2)

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




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

psi_minus.draw('mpl')

ex3_grader(psi_minus)

GHZ = QuantumCircuit(3)

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



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

GHZ.draw('mpl')

ex4_grader(GHZ)

Even = QuantumCircuit(3)

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





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

Even.draw('mpl')

ex5_grader(Even)

Odd = QuantumCircuit(3)

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





##### ==================================
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.

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

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.

# Create bell state circuit


# 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 = []


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

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.
# 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

# 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 = 
exp_Ab = 
exp_aB = 
exp_ab = 

CHSH = 
##### ==================================
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 =  #Write True or False :D
##### ==================================

# There is no grader for this question hehe

##### ==================================
# Write your solution in here.
Where_did_the_bell_pair_go = 'Write where you think they went'
##### ==================================

# There is no grader for this question hehe

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

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

import qiskit.tools.jupyter
%qiskit_version_table