from qiskit import QuantumRegister
from qiskit.circuit import QuantumCircuit, Parameter

from qiskit_cold_atom.spins.spins_gate_library import RydbergFull
from qiskit_cold_atom.spins import SpinSimulator

import numpy as np
import matplotlib.pyplot as plt

Nwires = 2
backend = SpinSimulator()
qc_rabi = QuantumCircuit(QuantumRegister(Nwires, "spin"))

all_modes=range(Nwires)
omega_t = Parameter("ω")

qc_rabi.rlx(omega_t, [0, 1])

qc_rabi.measure_all()
qc_rabi.draw(output='mpl', style='clifford')

phases = np.linspace(0, 2*np.pi, 15)

rabi_list = [
    qc_rabi.assign_parameters(
        {omega_t: phase},
        inplace=False,
    )
    for phase in phases
]

n_shots = 500
job_rabi = backend.run(rabi_list, shots=n_shots)
result_rabi = job_rabi.result()
counts_rabi = result_rabi.get_counts()

sum_rabi = []
for i, _ in enumerate(phases):
    counts_up = 0
    for outcome, count in counts_rabi[i].items():
        if outcome[0] == "1":
            counts_up += count
        if outcome[-1] == "1":
            counts_up += count        
    sum_rabi.append(counts_up/n_shots)

f, ax = plt.subplots()

ax.plot(phases, sum_rabi, 'o')
ax.set_xlabel('phase')
ax.set_ylabel('sum of spins up')

Nwires = 2
backend = SpinSimulator()
qc_block = QuantumCircuit(QuantumRegister(Nwires, "spin"))

all_modes=range(Nwires)
omega_t = Parameter("ω")
phi_t = Parameter("φ")

qc_block.rydberg_full(omega=omega_t, delta=0, phi=phi_t, modes=all_modes)

qc_block.measure_all()
qc_block.draw(output='mpl', style='clifford')

phases = np.linspace(0., 2*np.pi, 15)

block_list = [
    qc_block.assign_parameters(
        {omega_t: phase, phi_t: phase*10},
        inplace=False,
    )
    for phase in phases
]

job_block = backend.run(block_list, shots=n_shots)
result_block = job_block.result()
counts_block = job_block.result().get_counts()

sum_block = []
for i, _ in enumerate(phases):
    counts_up = 0
    for outcome, count in counts_block[i].items():
        if outcome[0] == "1":
            counts_up += count
        if outcome[-1] == "1":
            counts_up += count        
    sum_block.append(counts_up/n_shots)

f, ax = plt.subplots()

ax.plot(phases, sum_rabi, 'o', label= 'Rabi')
ax.plot(phases, sum_block, 'o', label= 'Blockade')
ax.plot(phases, (1-np.cos(np.sqrt(2)*phases))/2, '-', color = "C1", alpha = 0.5, linewidth = 3, label= 'blocked theory')
ax.legend()
ax.set_xlabel('phase')
ax.set_ylabel('sum of spins up')

Nwires = 2
backend = SpinSimulator()
qc_entangle = QuantumCircuit(QuantumRegister(Nwires, "spin"))

all_modes=range(Nwires)

qc_entangle.rlx(np.pi/2, all_modes)
qc_entangle.rydberg_block(np.pi, modes=all_modes)
qc_entangle.rly(np.pi/2, 0)
qc_entangle.rly(np.pi/2, 1)
qc_entangle.measure_all()
qc_entangle.draw(output='mpl', style='clifford')

job_entangle = backend.run(qc_entangle, shots=500)
result_entangle = job_entangle.result()
print(result_entangle.get_counts())

Nwires = 5
backend = SpinSimulator()
qc_rabi = QuantumCircuit(QuantumRegister(Nwires, "spin"))

all_modes=range(Nwires)
omega_t = Parameter("ω")

qc_rabi.rlx(omega_t, all_modes)

qc_rabi.measure_all()
qc_rabi.draw(output='mpl', style='clifford')

phases = np.linspace(0, 2*np.pi, 15)

rabi_list = [
    qc_rabi.assign_parameters(
        {omega_t: phase},
        inplace=False,
    )
    for phase in phases
]

job_rabi = backend.run(rabi_list, shots=500)
result_rabi = job_rabi.result()

outcomes_rabi = [result_rabi.get_memory(i) for i in range(len(rabi_list))]
for i, outcome in enumerate(outcomes_rabi):
    for j, run in enumerate(outcome):
        outcomes_rabi[i][j] = np.array(run.split(' '), dtype = int)
outcomes_rabi = np.array(outcomes_rabi)

means_rabi = outcomes_rabi.mean(axis=1)

f, ax = plt.subplots()
im = ax.pcolormesh(phases, np.arange(Nwires), means_rabi.T)
f.colorbar(im, ax=ax)
ax.set_ylabel('site')
ax.set_xlabel('time')

Nwires = 7
backend = SpinSimulator()
qc_block = QuantumCircuit(QuantumRegister(Nwires, "spin"))

all_modes=range(Nwires)
omega_t = Parameter("ω")
phi_t = Parameter("φ")

qc_block.rydberg_full(omega=omega_t, delta=0, phi=phi_t, modes=all_modes)

qc_block.measure_all()
qc_block.draw(output='mpl', style='clifford')

phases = np.linspace(0, 2*np.pi, 15)

block_list = [
    qc_block.assign_parameters(
        {omega_t: phase, phi_t: phase*5},
        inplace=False,
    )
    for phase in phases
]

job_block = backend.run(block_list, shots=500)
result_block = job_block.result()

outcomes = [result_block.get_memory(i) for i in range(len(block_list))]
for i, outcome in enumerate(outcomes):
    for j, run in enumerate(outcome):
        outcomes[i][j] = np.array(run.split(' '), dtype = int)
outcomes = np.array(outcomes)

means_block = outcomes.mean(axis=1)

f, ax = plt.subplots()
im = ax.pcolormesh(phases, np.arange(Nwires), means_block.T)
f.colorbar(im, ax=ax)
ax.set_ylabel('site')
ax.set_xlabel('time')

from qiskit_cold_atom.providers import ColdAtomProvider
from pprint import pprint

#provider = ColdAtomProvider.save_account(
#    url=[
#        "https://qlued.alqor.io/api/v2/singlequdit",
#        "https://qlued.alqor.io/api/v2/multiqudit",
#        "https://qlued.alqor.io/api/v2/fermions",
#        "https://qlued.alqor.io/api/v2/rydberg",
#    ],
#    username="name",
#    token="token",
#    overwrite=True
#)

provider = ColdAtomProvider.load_account()

spin_device_backend = provider.get_backend("alqor_rydberg_simulator")
pprint(spin_device_backend.configuration().supported_instructions)

Nwires = 5

qc_rabi = QuantumCircuit(5)

all_modes=range(Nwires)
omega_t = Parameter("ω")

qc_rabi.rlx( omega_t, all_modes)

qc_rabi.measure_all()
qc_rabi.draw(output='mpl')

phases = np.linspace(0, 2*np.pi,15)

remote_rabi_list = [
    qc_rabi.assign_parameters(
        {omega_t: phase},
        inplace=False,
    )
    for phase in phases
]

job_remote_rabi = spin_device_backend.run(remote_rabi_list, shots=500)

job_remote_rabi.job_id()

job_rabi_retrieved = spin_device_backend.retrieve_job(job_id=job_remote_rabi.job_id())

print("job status: ", job_rabi_retrieved.status())

result_remote_rabi = job_rabi_retrieved.result()

outcomes = [result_remote_rabi.get_memory(i) for i in range(len(remote_rabi_list))]
for i, outcome in enumerate(outcomes):
    for j, run in enumerate(outcome):
        outcomes[i][j] = np.array(run.split(' '), dtype = int)
        
outcomes = np.array(outcomes)
means_remote_rabi = outcomes.mean(axis=1)

f, ax = plt.subplots()
im = ax.pcolormesh(phases,np.arange(5), means_remote_rabi.T)
f.colorbar(im, ax=ax)

ax.set_ylabel('site')
ax.set_xlabel('time')

Nwires = 5
qc2 = QuantumCircuit(5)

all_modes=range(Nwires)
alpha = Parameter("α")
beta = Parameter("β")

qc2.rydberg_full(omega = alpha, delta =0, phi = beta, modes=all_modes)

qc2.measure_all()
qc2.draw(output='mpl')

phases = np.linspace(0, 2*np.pi,15)

circuit2_list = [
    qc2.assign_parameters(
        {alpha: phase, beta: phase*2},
        inplace=False,
    )
    for phase in phases
]

job2 = spin_device_backend.run(circuit2_list, shots=500)

job2.job_id()

job_retrieved2 = spin_device_backend.retrieve_job(job_id=job2.job_id())

print("job status: ", job_retrieved2.status())

result2 = job_retrieved2.result()

outcomes = [result2.get_memory(i) for i in range(len(circuit2_list))]
for i, outcome in enumerate(outcomes):
    for j, run in enumerate(outcome):
        outcomes[i][j] = np.array(run.split(' '), dtype = int)
outcomes = np.array(outcomes)

means_1 = outcomes.mean(axis=1)
means_1.shape

f, ax = plt.subplots()
im = ax.pcolormesh(phases,np.arange(5), means_1.T)
f.colorbar(im, ax=ax)

ax.set_ylabel('site')
ax.set_xlabel('time')

import qiskit.tools.jupyter

%qiskit_version_table

