from qiskit_cold_atom.providers import ColdAtomProvider

provider = ColdAtomProvider()
backend = provider.get_backend("collective_spin_simulator")

from qiskit import QuantumCircuit
from qiskit_cold_atom.spins.spins_gate_library import OATGate
from qiskit.visualization import plot_histogram

qc = QuantumCircuit(1)

qc.append(OATGate(chi=0.4, delta=-1., omega=2.), qargs=[0])

# alternatively append the OAT gate directly:
# qc.oat(chi=0.4, delta=-1., omega=2., wire=0)

qc.measure_all()

# simulate the circuit with a spin length L=5
counts = backend.run(qc, spin=5).result().get_counts()
# convert counts to integers for better formatting
plot_histogram({int(k):v for k,v in counts.items()})

from qiskit_cold_atom.spins.spins_gate_library import RLXGate, RLZGate, RLZ2Gate
print("LX generator: \n", RLXGate(0.5).generator)
print("LZ generator: \n", RLZGate(1.).generator)
print("LZ2 generator: \n", RLZ2Gate(1.5).generator)

from qiskit.circuit import Parameter
import matplotlib.pyplot as plt
import numpy as np

# lenght of spin for the simulation
spin = 20

# sample at 50 angles between 0 and 4*pi
vals = 50
angles = np.linspace(0, 4*np.pi, vals)

omega = Parameter("omega")
# create list of circuits
circuit = QuantumCircuit(1, 1)
circuit.rlx(omega, 0)
circuit.measure(0, 0)

rabi_list = [circuit.assign_parameters([angle]) for angle in angles]

job_rabi = backend.run(rabi_list, shots = 10, spin = spin, seed=5462)
result_rabi = job_rabi.result()

outcomes = np.array([np.array(result_rabi.get_memory(i), dtype=float) for i in range(vals)])
plt.errorbar(angles, np.mean(outcomes, axis=1), yerr=np.std(outcomes, axis=1), fmt='o', color="#4589ff")
plt.title(f"Rabi oscillations for a single collective spin of length L = {spin}")
plt.ylabel("measurement outcome")
plt.xlabel("angle of x-rotation")
plt.show()

# circuit with spin squeezing
squeez_circ = QuantumCircuit(1, 1)
squeez_circ.rlx(-np.pi/2, 0)
squeez_circ.rlz2(0.3, 0)
squeez_circ.rlz(-np.pi/2, 0)
squeez_circ.rlx(-0.15, 0)
squeez_circ.measure(0, 0)

job_squeez = backend.run(squeez_circ, shots = 1000, spin=20, seed=14)

squeez_circ.draw(output='mpl', style="clifford")

# for comparison: circuit with a single rotation
circ_x = QuantumCircuit(1, 1)
circ_x.rlx(np.pi/2, 0)
circ_x.measure_all()

job_x = backend.run(circ_x, shots = 1000, spin=20, seed=14)

fig, (ax1, ax2) = plt.subplots(2, figsize=(10,8))
fig.tight_layout(pad=5.0)
ax1.set_title("spin squeezing circuit")
ax2.set_title("single rotation circuit")

plot_histogram(job_squeez.result().get_counts(), ax=ax1, color="#9f1853") #, number_to_keep=10)
plot_histogram(job_x.result().get_counts(), ax=ax2) #, number_to_keep=10)

from qiskit_cold_atom.spins.spins_gate_library import RLZLZGate

qc = QuantumCircuit(2)

qc.rlx(np.pi/2, [0, 1])
qc.rlzlz(np.pi, [0, 1])
qc.rlx(np.pi/2, [0, 1])

qc.measure_all()

plot_histogram(backend.run(qc, spin=2).result().get_counts())

plot_histogram(backend.run(qc, spin=1/2).result().get_counts())

from pprint import pprint

collective_spins_configuration = backend.configuration().to_dict()

pprint(collective_spins_configuration)

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright

