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

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

# Simulamos para ver su resultado
backend = Aer.get_backend('statevector_simulator')
result = execute(phi_plus, backend).result().get_statevector()
array_to_latex(result)

phi_minus = QuantumCircuit(2)

##### ==================================
# Escribe tu solución aquí.
phi_minus.h(0)
phi_minus.z(0)
phi_minus.cnot(0, 1)
##### ==================================

phi_minus.draw('mpl')

# Esta celda te dirá si lo que has hecho da el resultado correcto.
# Tómalo como un indicador, los tests que hace no consideran todas las opciones posibles.
# Si crees que tu respuesta es (in)correcta y el grader dice algo diferente, simplemente pregúntame.
# No intentes hacer trampas modificando el código del corrector :)

ex1_grader(phi_minus)

psi_plus = QuantumCircuit(2)

##### ==================================
# Escribe tu solución aquí.
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)

##### ==================================
# Escribe tu solución aquí.
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)

##### ==================================
# Escribe tu solución aquí.
GHZ.h(0)
GHZ.cnot(0, 1)
GHZ.cnot(0, 2)
##### ==================================

GHZ.draw('mpl')

ex4_grader(GHZ)

Even = QuantumCircuit(3)

##### ==================================
# Escribe tu solución aquí.
Even.h(0)
Even.h(1)
Even.cnot(0, 2)
Even.cnot(1, 2)
##### ==================================

Even.draw('mpl')

ex5_grader(Even)

Odd = QuantumCircuit(3)

##### ==================================
# Escribe tu solución aquí.
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): #Las medidas deben guardarse en bits clásicos.
    qc.measure(qubit, cbit)

def measure_x_axis(qc, qubit, cbit):
    ##### ==================================
    # Escribe tu solución aquí.
    qc.h(qubit)
    qc.measure(qubit, cbit)
    ##### ==================================

ex7_grader(measure_x_axis)

def expectation_value_single_qubit(qc, nshots=8092):
    # Ejecuta el circuito que ya está rotado a la base correspondiente
    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
    
    # Esto es simplemente para evitar problemas. En caso de que falte una key en el diccionario la añade
    if '1' not in counts.keys():
        counts['1'] = 0
    if '0' not in counts.keys():
        counts['0'] = 0
        
    # Obten el numero de veces que se ha medido 0 y 1
    n_zeros = counts['0']
    n_ones  = counts['1']
    
    # Calcula las probabilidades
    p_0 = n_zeros/nshots
    p_1 = n_ones/nshots # O 1-p_0

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

# Mide <Z> sobre un estado que te guste
qc_z = QuantumCircuit(1,1)
measure_z_axis(qc_z, 0, 0)

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

# Mide <X> sobre un estado que te guste
qc_x = QuantumCircuit(1,1)
measure_x_axis(qc_x, 0, 0)

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

##### ==================================
# Escribe tu solución aquí.
measure_strings = ['ZZ', 'ZX', 'XZ', 'XX']
# Primero crea un estado de Bell
qc = QuantumCircuit(2,2)
qc.h(0)
qc.cnot(0, 1)

# Para cada uno de los operadores que queremos medir, (ZZ, ZX, XZ and XX) crea un nuevo circuito (quizás usando qc.copy())
# que mida en ese eje.
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')

# Hacer tests aquí que no revelasen la respuesta era complicado, así que no hay ayuda :(

# Podemos correr todos los circuitos a la vez y luego postprocesar los datos
nshots = 8092 # Cuantos más shots, menos error
backend_shots = Aer.get_backend('qasm_simulator')
counts_list = execute(chsh_circuits, backend_shots, shots=nshots).result().get_counts()
counts_list

##### ==================================
# Escribe tu solución aquí.
exp_dict = {}
# Para cada uno de los circuitos, coge las counts que tengas para ese circuito, calcula la probabilidad
#de cada estado y combinalas teniendo mucho cuidado con qué signo lleva cada término
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
# Puede que quieras inspirarte en la función expectation_value_single_qubit
# Si estás muy perdido y no sabes cómo resolver esto, puedes preguntar.

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

# Si no sabes si tu respuesta es correcta, puedes calcular los resultados a mano
# (de hecho seguramente sea más fácil que calcularlo así) y comprobar si tu código está funcionando

##### ==================================
# Escribe tu solución aquí.
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('Tu resultado es <CHSH>=', CHSH)

print('El resultado correcto es <CHSH>=', 2*np.sqrt(2))

##### ==================================
# Escribe tu solución aquí.
Can_entanglement_be_teleported = True#Escribe True or False :D
##### ==================================

##### ==================================
# Escribe tu solución aquí.
Where_did_the_bell_pair_go = 'En los dos últimos cúbits'
##### ==================================

##### ==================================
# Escribe tu solución aquí.

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

import qiskit.tools.jupyter
%qiskit_version_table

