import numpy as np

from qiskit import QuantumCircuit

circuit = QuantumCircuit(2)

# initialization to create a 50 - 50 superposition
circuit.h(0)
circuit.h(1)

# measurement
circuit.measure_all()
circuit.draw("mpl", style='clifford')

from qiskit import Aer
from qiskit.visualization import plot_histogram

def bit_counts_to_Lz(counts):
    """transform bitstring counts to collective counts.

    Args:
        counts: a dict whose keys are the observed integers for the spin
                and whose values is the number of observations

    Returns:
        counts: a dict whose keys are the observed integers for the total observed numbers of up
                and whose values is the number of observations
    """
    lz_counts = {}
    for bitstring, count in counts.items():
        measured_label = bitstring.count("1")
        if str(measured_label) in lz_counts:
            lz_counts[str(measured_label)] += count
        else:
            lz_counts[str(measured_label)] = count

    return lz_counts


qasm_simulator = Aer.get_backend("aer_simulator")
job = qasm_simulator.run(circuit, shots=1000)
counts = job.result().get_counts()
lz_counts_uncorrelated = bit_counts_to_Lz(counts)
plot_histogram(lz_counts_uncorrelated)

circuit = QuantumCircuit(2)

# initialization on x
circuit.h(0)
circuit.x(1)
circuit.cx(0, 1)

# measurement
circuit.measure_all()
circuit.draw("mpl", style='clifford')

job = qasm_simulator.run(circuit, shots=1000)
counts = job.result().get_counts()
counts["00"] = 0
counts["11"] = 0
squeezed_lz_counts = bit_counts_to_Lz(counts)
plot_histogram(squeezed_lz_counts)

from qiskit_cold_atom.providers import ColdAtomProvider
from qiskit_cold_atom.spins.spins_gate_library import OATGate
from qiskit.circuit import Parameter

import numpy as np

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

qc = QuantumCircuit(1)

alpha = Parameter("α")
beta = Parameter("β")
gamma = Parameter("γ")

qc.rly(np.pi / 2, 0)
qc.append(OATGate(chi=gamma, delta=0, omega=beta), qargs=[0])
qc.rlx(np.pi, 0)
qc.append(OATGate(chi=gamma, delta=0, omega=beta), qargs=[0])
qc.rlx(alpha, 0)
qc.measure_all()

qc.draw("mpl", style='clifford')

# Number of atoms
Nat = 100

# The length of the spin
l = Nat / 2

# Strength of the rotation around x 
omegax = 2 * np.pi * 20

# Time of the evolution
time = 15e-3

# Strength of the Lz-squared operator
Lambda = 1.5
chi = Lambda * abs(omegax) / Nat;

alphas = np.array([50, 140]) / 180 * np.pi

circuit = qc.assign_parameters(
    {
        alpha: np.pi - alphas[0],
        beta: omegax * time / 2,
        gamma: chi * time / 2,
    },
    inplace=False,
)

job = backend.run(circuit, shots=1000, spin=Nat / 2)
counts1 = job.result().get_counts()


circuit = qc.assign_parameters(
    {
        alpha: np.pi - alphas[1],
        beta: omegax * time / 2,
        gamma: chi * time / 2,
    },
    inplace=False,
)

job = backend.run(circuit, shots=1000, spin=Nat / 2)
counts2 = job.result().get_counts()

import matplotlib.pyplot as plt
from scipy.special import binom

f, ax = plt.subplots()

obs_num = np.array([int(key) for key in counts1.keys()])
obs_vals = np.array([val for val in counts1.values()])

ax.bar(
    obs_num - Nat / 2,
    obs_vals / obs_vals.sum(),
    align="center",
    alpha=0.5,
    label="squeezed",
)

obs_num = np.array([int(key) for key in counts2.keys()])
obs_vals = np.array([val for val in counts2.values()])

ax.bar(
    obs_num - Nat / 2,
    obs_vals / obs_vals.sum(),
    align="center",
    alpha=0.5,
    label="anti-squeezed",
)

# and compare to the binomial distribution;
ks = np.arange(Nat)
pk = binom(Nat, ks) * 0.5 ** ks * 0.5 ** (Nat - ks)
ax.plot(ks - Nat / 2, pk, label="uncorrelated noise", c="C2")
ax.legend()
ax.set_ylabel("Probability")
ax.set_xlabel(r"$L_z$");

t1 = 15e-3
t2 = 25e-3

def mean_var_from_counts(counts):
    """calculate mean and variance from the counts.

    Args:
        counts: a dict whose keys are the observed integers for the spin
                and whose values is the number of observations

    Returns:
        mean, var: mean and variance of the observed counts
    """
    mean, var, count_sum = 0, 0, 0

    for k, v in counts.items():
        mean += int(k) * v
        count_sum += v

    mean = mean / count_sum
    for k, v in counts.items():
        var += (int(k) - mean) ** 2 * v / count_sum
    return mean, var

alphas = np.linspace(0, np.pi, 15)
variances_1 = np.zeros(len(alphas))
variances_2 = np.zeros(len(alphas))

for ii in range(len(alphas)):

    # first time step
    circuit = qc.assign_parameters(
        {
            alpha: alphas[ii],
            beta: omegax * t1 / 2,
            gamma: chi * t1 / 2,
        },
        inplace=False,
    )

    job = backend.run(circuit, shots=1000, spin=Nat / 2, seed=14)
    counts = job.result().get_counts()
    mean, var = mean_var_from_counts(counts)
    variances_1[ii] = var

    # and the second time step
    circuit = qc.assign_parameters(
        {
            alpha: alphas[ii],
            beta: omegax * t2 / 2,
            gamma: chi * t2 / 2,
        },
        inplace=False,
    )

    job = backend.run(circuit, shots=1000, spin=Nat / 2, seed=14)
    counts = job.result().get_counts()
    mean, var = mean_var_from_counts(counts)
    variances_2[ii] = var;

# Import the data from the paper by Strobel et al.
strobel_15_dB = np.array(
    [
        0.51993068,
        10.39861352,
        20.5372617,
        30.41594454,
        40.81455806,
        45.49393414,
        50.6932409,
        56.152513,
        60.05199307,
        69.93067591,
        79.80935875,
        90.98786828,
        100.08665511,
        109.96533795,
        120.36395147,
        130.24263432,
        139.86135182,
        150.51993068,
        160.13864818,
        170.01733102,
        180.1559792,
    ]
)

strobel_15_alpha = np.array(
    [
        9.27455357,
        8.03571429,
        6.19419643,
        3.54910714,
        -0.46875,
        -2.27678571,
        -4.11830357,
        -4.01785714,
        -2.94642857,
        1.30580357,
        4.55357143,
        6.99776786,
        8.50446429,
        9.140625,
        10.41294643,
        10.546875,
        10.64732143,
        10.27901786,
        9.74330357,
        9.40848214,
        8.27008929,
    ]
)

# and the values for 25ms
strobel_25_dB = np.array(
    [
        0.51993068,
        10.39861352,
        20.5372617,
        30.41594454,
        40.81455806,
        45.49393414,
        50.6932409,
        56.152513,
        60.05199307,
        69.93067591,
        79.80935875,
        90.98786828,
        100.08665511,
        109.96533795,
        120.36395147,
        130.24263432,
        139.86135182,
        150.51993068,
        160.13864818,
        170.01733102,
        180.1559792,
    ]
)

strobel_25_alpha = np.array(
    [
        16.07142857,
        15.46875,
        14.02901786,
        11.71875,
        8.40401786,
        4.72098214,
        2.20982143,
        -0.46875,
        0.0,
        2.74553571,
        7.734375,
        10.98214286,
        13.39285714,
        14.89955357,
        15.73660714,
        16.50669643,
        17.00892857,
        17.31026786,
        17.24330357,
        16.77455357,
        15.80357143,
    ]
)

def number_squeezing_factor_to_db(var_CSS, var):
    """Calcululate the squeezing parameter from the observed variance."""
    return 10 * np.log10(var / var_CSS)

f, ax = plt.subplots()
ax.set_title("Number Squeezing")

ax.plot(
    np.rad2deg(np.pi - alphas),
    number_squeezing_factor_to_db(l / 2, variances_1),
    "r-",
    lw=5,
    label="simulated  15ms",
    alpha=0.5,
)
ax.plot(strobel_15_dB, strobel_15_alpha, "ro", label="1experiment 5ms", markersize=4)
ax.plot(
    np.rad2deg(np.pi - alphas),
    number_squeezing_factor_to_db(l / 2, variances_2),
    "b-",
    lw=5,
    label="simulated  25ms",
    alpha=0.5,
)
ax.plot(strobel_25_dB, strobel_25_alpha, "bo", label="experiment 25ms", markersize=4)
ax.axhline(y=0, color="g", linestyle="--")
ax.set_ylabel(r"Number Squeezing in $dB$")
ax.set_xlabel(r"tomography angle $\alpha$")
ax.legend();

from qiskit_cold_atom.providers import ColdAtomProvider
from pprint import pprint

# save an account to disk
#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",
#    ],
#    username="name",
#    token="token",
#    overwrite=True
#)

# or enable an account in the current session
#provider = ColdAtomProvider.enable_account(
#    url=[
#        "https://qlued.alqor.io/api/v2/singlequdit",
#        "https://qlued.alqor.io/api/v2/multiqudit",
#        "https://qlued.alqor.io/api/v2/fermions",
#    ],
#    username="name",
#    token="token",
#    overwrite=True
#)

provider = ColdAtomProvider.load_account()

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

Nat = 200
l = Nat / 2  # spin length
omegax = 2 * np.pi * 20
t1 = 15e-3
t2 = 25e-3
Lambda = 1.5

chi = Lambda * abs(omegax) / Nat

Ntrott = 15
alphas = np.linspace(0, np.pi, 15)

alpha = Parameter("α")
beta = Parameter("β")
gamma = Parameter("γ")

qc_trott = QuantumCircuit(1, 1)
qc_trott.load_spins(Nat, 0)
qc_trott.rlx(np.pi / 2, 0)
qc_trott.rlz(np.pi / 2, 0)
for ii in range(Ntrott):
    qc_trott.rlx(beta, 0)
    qc_trott.rlz2(gamma, 0)
qc_trott.rlx(alpha, 0)
qc_trott.measure(0, 0)
qc_trott.draw(output="mpl")

circuit1_list = [
    qc_trott.assign_parameters(
        {alpha: -a % (2 * np.pi), gamma: chi * t1 / Ntrott, beta: omegax * t1 / Ntrott},
        inplace=False,
    )
    for a in alphas
]
job1 = spin_device_backend.run(circuit1_list, shots=500)

job1.job_id()

job_retrieved1 = spin_device_backend.retrieve_job(job_id=job1.job_id())

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

circuit2_list = [
    qc_trott.assign_parameters(
        {alpha: -a % (2 * np.pi), gamma: chi * t2 / Ntrott, beta: omegax * t2 / Ntrott},
        inplace=False,
    )
    for a in alphas
]

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())

result1 = job_retrieved1.result()
result2 = job_retrieved2.result()

outcomes1 = [result1.get_memory(i) for i in range(len(circuit1_list))]
outcomes2 = [result2.get_memory(i) for i in range(len(circuit2_list))]

variances_1 = [np.var(np.array(counts1, dtype=float)) for counts1 in outcomes1]
variances_2 = [np.var(np.array(counts2, dtype=float)) for counts2 in outcomes2]

variances_1 = np.array(variances_1)
variances_2 = np.array(variances_2)

f, ax = plt.subplots()
ax.set_title("Number Squeezing")
plt.plot(
    np.rad2deg(alphas),
    number_squeezing_factor_to_db(l / 2, variances_1),
    "r-",
    lw=5,
    label="simulated  15ms",
    alpha=0.5,
)
plt.plot(
    np.rad2deg(alphas),
    number_squeezing_factor_to_db(l / 2, variances_2),
    "b-",
    lw=5,
    label="simulated  25ms",
    alpha=0.5,
)

plt.plot(strobel_15_dB, strobel_15_alpha, "ro", label="15ms", markersize=4)
plt.plot(strobel_25_dB, strobel_25_alpha, "bo", label="25ms", markersize=4)


ax.axhline(y=0, color="g", linestyle="--")
ax.set_ylabel(r"Number Squeezing in $dB$")
ax.set_xlabel(r"tomography angle $\alpha$")
ax.legend()

import qiskit.tools.jupyter

%qiskit_version_table

