from qiskit_alice_bob_provider import AliceBobLocalProvider
from qiskit_aer import AerSimulator
import tempfile
from qiskit import transpile, execute, QuantumCircuit
from qiskit.circuit import Qubit, Clbit
import numpy as np
from tqdm.notebook import tqdm
from matplotlib import pyplot as plt

provider = AliceBobLocalProvider()
print(provider.backends())

backend = provider.get_backend('EMU:15Q:LOGICAL_EARLY')

def build_teleportation_routine() -> QuantumCircuit:
    teleportation = QuantumCircuit(3, 2, name='teleport')
    teleportation.initialize('+', [1])
    teleportation.initialize('0', [2])
    teleportation.cx(1, 2)
    teleportation.cx(0, 1)
    teleportation.measure_x(0, 0)
    teleportation.measure(1, 1)
    teleportation.x(2).c_if(1, 1)
    teleportation.z(2).c_if(0, 1)
    return teleportation

build_teleportation_routine().draw('mpl')

def build_teleportation_tomography(teleported_state: str) -> QuantumCircuit:
    assert teleported_state in {'+', '-', '0', '1'}
    circuit = QuantumCircuit(3, 3)
    circuit.initialize(teleported_state, 0)
    circuit.append(build_teleportation_routine().to_instruction(), [0, 1, 2], [0, 1])
    if teleported_state in {'+', '-'}:
        circuit.measure_x(2, 2)
    else:
        circuit.measure(2, 2)
    return circuit

plus_circuit = build_teleportation_tomography('+')
plus_circuit.draw('mpl')

transpiled = transpile(plus_circuit, backend)
transpiled.draw('mpl')

execute(plus_circuit, backend).result().get_counts()

zero_circuit = build_teleportation_tomography('0')
execute(zero_circuit, backend).result().get_counts()

k1 = 100
k2 = 100_000
average_nb_photonss = np.linspace(4, 24, 11)
distances = 2 * np.arange(1, 10) + 1
shots = 10_000

bit_flips = np.zeros((max(average_nb_photonss.shape), max(distances.shape)))
phase_flips = np.zeros((max(average_nb_photonss.shape), max(distances.shape)))

plus_circuit = build_teleportation_tomography('+')
zero_circuit = build_teleportation_tomography('0')

for j, distance in enumerate(tqdm(distances)):
    for i, average_nb_photons in enumerate(tqdm(average_nb_photonss)):
        backend = provider.get_backend(
            'EMU:40Q:LOGICAL_TARGET', 
            average_nb_photons=average_nb_photons, distance=distance, 
            kappa_2=k2, kappa_1=k1
        )
        counts = execute(zero_circuit, backend, shots=shots).result().get_counts()
        for word, count in counts.items():
            if word[0] == '1':
                bit_flips[i][j] += count
        counts = execute(plus_circuit, backend, shots=shots).result().get_counts()
        for word, count in counts.items():
            if word[0] == '1':
                phase_flips[i][j] += count

plt.figure()
plt.title(f'Bit flip rate (k1/k2={k1/k2:.0e})')
colors = plt.cm.viridis(np.linspace(0, 1, max(distances.shape)))
for j, distance in enumerate(distances):
    plt.plot(average_nb_photonss, bit_flips[:, j] / shots, label=f'Distance {distance}', c=colors[j])
plt.xlabel('Average number of photons')
plt.ylabel('Bit flip rate')
plt.legend()
plt.show()

plt.figure()
plt.title(f'Phase flip rate (k1/k2={k1/k2:.0e})')
colors = plt.cm.viridis(np.linspace(0, 1, max(distances.shape)))
for j, distance in enumerate(distances):
    plt.plot(average_nb_photonss, phase_flips[:, j] / shots, label=f'Distance {distance}', c=colors[j])
plt.xlabel('Average number of photons')
plt.ylabel('Phase flip rate')
plt.legend()
plt.show()