import numpy as np
import matplotlib.pyplot as plt
import time

from qiskit_cold_atom.fermions import FermionSimulator, FfsimBackend

from qiskit_cold_atom.fermions.fermion_gate_library import FermiHubbard
from qiskit_nature.operators.second_quantization import FermionicOp
from qiskit.circuit import Parameter

backend = FermionSimulator()
backend_ffsim = FfsimBackend()

def prep_init_states(dmax: int, with_hole: bool) -> tuple[list, list]:
    """
    Prepares a list of initial states for the Fermi-Hubbard model.

    Args:
        dmax: the maximum distance from the center
        with_hole: whether to include a hole in the initial state

    Returns:
        two lists of initial states starting with spin-up and spin-down, respectively
    """
    Nsites = 2 * dmax + 1

    # the up state
    state_up = [i % 2 for i in range(Nsites)]
    state_down = [(i + 1) % 2 for i in range(Nsites)]

    if with_hole:
        state_up[dmax] = 0
        state_down[dmax] = 0

    init_state_1 = [state_up, state_down]
    init_state_2 = [state_down, state_up]

    return init_state_1, init_state_2

dmax = 3
Nsites = 2 * dmax + 1

init_state_1, init_state_2 = prep_init_states(dmax, with_hole=True)

hop_param = Parameter("hop_param")
int_param = Parameter("int_param")
# initialize one spin-up and one spin-down atom in the left tweezer
qc_1 = backend.initialize_circuit(init_state_1)

# apply a global Fermi-Hubbard gate with interaction
qc_1.fhubbard(j=hop_param * np.ones(Nsites - 1), u=int_param, mu=np.zeros(Nsites), modes=range(2*Nsites))
qc_1.measure_all()

# qc_1.draw(output="mpl")

hop_param_down = Parameter("hop_param")
int_param_down = Parameter("int_param")

# initialize one spin-up and one spin-down atom in the left tweezer
qc_2 = backend.initialize_circuit(init_state_2)

# apply a global Fermi-Hubbard gate with interaction
qc_2.fhubbard(
    j=hop_param_down * np.ones(Nsites - 1),
    u=int_param_down,
    mu=np.zeros(Nsites),
    modes=range(2 * Nsites),
)
qc_2.measure_all()

circuit_up_init = qc_1.bind_parameters({hop_param: 0, int_param: 0})
job_up_init = backend.run(circuit_up_init, shots=5, num_species=2)

circuit_down_init = qc_2.bind_parameters({hop_param_down: 0, int_param_down: 0})
job_down_init = backend.run(circuit_down_init, shots=5, num_species=2)

def nocc_from_counts(counts: dict[str, int], Nsites: int) -> np.ndarray:
    """
    The calculates the occupation from all observed configurations.

    Args:
        counts: the counts of the observed configurations
        Nsites: the number of sites in the system

    Returns:
        1D array The occupation for each site
    """
    nocc_t = np.zeros(Nsites)
    Nobs = 0
    for occ_string, observations in counts.items():
        # get the occupation of each site in the observed configuration
        nocc = np.zeros(Nsites)

        for site_index in range(Nsites):
            nocc[site_index] = int(occ_string[site_index]) + int(occ_string[site_index + Nsites])

        # now add the value to the correct time step and with it by the number of observations
        Nobs = Nobs + observations
        nocc_t = nocc_t + nocc * observations
    return nocc_t / Nobs


def squeezed_spin_corr_from_counts(counts: dict[str, int], Nsites: int) -> np.ndarray:
    """
    The calculates the spin projection from all observed configurations.
    However, we filter out situations with a hole and a doublon in the system.

    Args:
        counts: the counts of the observed configurations
        Nsites: the number of sites in the system

    Returns:
        1D array The occupation for each site
    """
    sz_squeeze = np.zeros(Nsites)
    szvar = np.zeros(Nsites)
    Nobs = 0
    Nobs_var = np.zeros(Nsites)
    Nobs = np.zeros(Nsites)

    for occ_string, observations in counts.items():
        occ_int = [int(occ) for occ in occ_string]
        # get the spin of each site in the observed configuration
        sz = np.zeros(Nsites)
        nocc = np.zeros(Nsites)
        for site_index in range(Nsites):
            sz[site_index] = 1 / 2 * (occ_int[site_index] - occ_int[site_index + Nsites])
            nocc[site_index] = occ_int[site_index] + occ_int[site_index + Nsites]

        # now obtain the symmetrized variance
        # first for the first site
        site_index = 0
        if nocc[site_index] == 1 and nocc[site_index + 1] == 1:
                    szvar[site_index] = (
                        szvar[site_index] + sz[site_index] * sz[site_index + 1] * observations
                    )
                    Nobs_var[site_index] = Nobs_var[site_index] + observations
        
        # then for the last site
        site_index = Nsites - 1
        if nocc[site_index] == 1 and nocc[site_index - 1] == 1:
                    szvar[site_index] = (
                        szvar[site_index] + sz[site_index] * sz[site_index - 1] * observations
                    )
                    Nobs_var[site_index] = Nobs_var[site_index] + observations
        
        # and finally for all the other sites
        for site_index in range(1,Nsites-1):
            if (
                    nocc[site_index] == 1
                    and nocc[site_index + 1] == 1
                    and nocc[site_index - 1] == 1
                ):
                    szvar[site_index] = (
                        szvar[site_index]
                        + 1
                        / 2
                        * sz[site_index]
                        * (sz[site_index + 1] + sz[site_index - 1])
                        * observations
                    )
                    Nobs_var[site_index] = Nobs_var[site_index] + observations
            elif (
                    nocc[site_index] == 1
                    and nocc[site_index + 1] == 1
                    and nocc[site_index - 1] == 0
                ):
                    szvar[site_index] = (
                        szvar[site_index] + sz[site_index] * sz[site_index + 1] * observations
                    )
                    Nobs_var[site_index] = Nobs_var[site_index] + observations
            elif (
                    nocc[site_index] == 1
                    and nocc[site_index + 1] == 0
                    and nocc[site_index - 1] == 1
                ):
                    szvar[site_index] = (
                        szvar[site_index] + sz[site_index] * sz[site_index - 1] * observations
                    )
                    Nobs_var[site_index] = Nobs_var[site_index] + observations

        sz_squeeze += sz * (nocc == 1) * observations
        Nobs += (nocc == 1) * observations

    # regularize the observations
    var_weight = 1 / Nobs_var
    var_weight[np.isinf(var_weight)] = 0

    mean_weight = 1 / Nobs
    mean_weight[np.isinf(mean_weight)] = 0
    # now that we have gone through all observations, we can average
    fluc = np.zeros(Nsites)
    for site_index in range(Nsites):
        if site_index == 0:
            fluc[site_index] = 4 * (
                szvar[site_index] * var_weight[site_index]
                - sz_squeeze[site_index]
                * mean_weight[site_index]
                * sz_squeeze[site_index + 1]
                * mean_weight[site_index + 1]
            )
        elif site_index == Nsites - 1:
            fluc[site_index] = 4 * (
                szvar[site_index] * var_weight[site_index]
                - sz_squeeze[site_index]
                * mean_weight[site_index]
                * sz_squeeze[site_index - 1]
                * mean_weight[site_index - 1]
            )
        else:
            fluc[site_index] = 4 * (
                szvar[site_index] * var_weight[site_index]
                - sz_squeeze[site_index]
                * mean_weight[site_index]
                * (
                    sz_squeeze[site_index + 1] * mean_weight[site_index + 1]
                    + sz_squeeze[site_index - 1] * mean_weight[site_index - 1]
                )
                / 2
            )
    return fluc, Nobs_var, Nobs

noccs_init = np.zeros( Nsites)
corrs_init = np.zeros(Nsites)

counts_1 = job_up_init.result().get_counts()
counts_2 = job_down_init.result().get_counts()
total_counts = counts_1
total_counts.update(counts_2)

nocc = nocc_from_counts(total_counts, Nsites)
corr, _, _ = squeezed_spin_corr_from_counts(total_counts, Nsites)

f, ax = plt.subplots()

ax.plot(np.arange(Nsites) - dmax, nocc, "bo", label="nocc")
ax.plot(np.arange(Nsites) - dmax, corr, "ro", label="corr", alpha = 0.5)
ax.legend()
ax.set_xlabel("site index")

# the experimental parameters
J = 2 * np.pi * 250  # hopping parameter in units of hbar
U = 15 * J  # interaction parameter

Ntimes = 8
tmax = 1.5
times = np.linspace(0, tmax, Ntimes) * 1e-3

Nshots = 500
# create list of circuits with interaction
circuits_1 = [qc_1.bind_parameters({hop_param: J * t, int_param: U * t}) for t in times]
# measure the observable from simulated shots
jobs_1 = backend.run(circuits_1, shots=Nshots, num_species=2)

# create list of circuits with interaction
circuits_2 = [qc_2.bind_parameters({hop_param_down: J * t, int_param_down: U * t}) for t in times]
# measure the observable from simulated shots
jobs_2 = backend.run(circuits_2, shots=Nshots, num_species=2)

noccs = np.zeros((Ntimes, Nsites))
corrs = np.zeros((Ntimes, Nsites))

start_time = time.time()
count_list_1 = jobs_1.result().get_counts()
count_list_2 = jobs_2.result().get_counts()
end_time = time.time()
print(f"Duration for of the two jobs: {end_time - start_time} seconds")

for ii in range(Ntimes):
    counts_1 = count_list_1[ii]
    counts_2 = count_list_2[ii]

    total_counts = counts_1
    total_counts.update(counts_2)

    nocc = nocc_from_counts(total_counts, Nsites)
    corr, _,_ = squeezed_spin_corr_from_counts(total_counts, Nsites)

    noccs[ii, :] = nocc
    corrs[ii, :] = corr

f, [ax1, ax2] = plt.subplots(1, 2, sharey=True)
# color plot of the occupation in a heatmap
velo = J
ax1.pcolor(np.arange(Nsites) - dmax, times * 1e3, corrs, cmap="Reds")

ax1.set_xlabel("time [ms]")
ax1.set_xlabel("site index")
ax1.set_title("Spin")

ax2.pcolor(np.arange(Nsites) - dmax, times * 1e3, noccs, cmap="Blues")
ax2.set_xlabel("site index")
ax2.set_title("Density")

stds_spin = np.zeros(Ntimes)
for ii, corr in enumerate(corrs):
    mass_func = corr + 1
    mean_position = np.average(np.arange(len(mass_func)), weights=mass_func)
    # Calculate the standard deviation of the position
    std_position = np.sqrt(
        np.average((np.arange(len(mass_func)) - mean_position) ** 2, weights=mass_func)
    )
    stds_spin[ii] = std_position

stds_density = np.zeros(Ntimes)
for ii, nocc in enumerate(noccs):
    # print(nocc)
    mass_func = 1 - nocc
    mean_position = np.average(np.arange(len(mass_func)), weights=mass_func)
    # Calculate the standard deviation of the position
    std_position = np.sqrt(
        np.average((np.arange(len(mass_func)) - mean_position) ** 2, weights=mass_func)
    )
    stds_density[ii] = std_position

f, ax = plt.subplots()
ax.plot(times * 1e3, stds_spin, "ro", label="spin")
ax.plot(times * 1e3, stds_density, "bo", label="density")
ax.set_xlabel("time [ms]")
ax.set_ylabel("standard deviation")
ax.legend()

# prepare the jobs
jobs_1_ffsim = backend_ffsim.run(circuits_1, shots=Nshots, num_species=2)
jobs_2_ffsim = backend_ffsim.run(circuits_2, shots=Nshots, num_species=2)

# run them and time it
start_time = time.time()
count_list_1 = jobs_1_ffsim.result().get_counts()
count_list_2 = jobs_2_ffsim.result().get_counts()
end_time = time.time()
print(f"Duration for of the two jobs: {end_time - start_time} seconds")


hop_long = Parameter("hop_long")
int_long = Parameter("int_long")
# initialize one spin-up and one spin-down atom in the left tweezer
qc_long = backend_ffsim.initialize_circuit(init_state_1)

# apply a global Fermi-Hubbard gate with interaction
qc_long.fhubbard(
    j=hop_long * np.ones(Nsites - 1), u=int_long, mu=np.zeros(Nsites), modes=range(2 * Nsites)
)
qc_long.measure_all()

tmax = 1 / U

circuit_long = qc_long.bind_parameters({hop_long: J * tmax, int_long: U * tmax})
job_long = backend.run(circuit_long, shots=5, num_species=2)
start_time = time.time()
count_long = job_long.result().get_counts()
end_time = time.time()
print(f"Duration for of the two jobs: {end_time - start_time:.2f} seconds")

time_step = 1/(4*U)
print(f"time step = {time_step*1e3:.2f} ms")

Ntrott = int(tmax / time_step)
print(f"Number of trotter steps: {Ntrott}")

trott_hop = Parameter("trott_hop")
trott_int = Parameter("trott_int")
# initialize one spin-up and one spin-down atom in the left tweezer
qc_trott = backend_ffsim.initialize_circuit(init_state_1)

for ii in range(Ntrott):
    # apply a global Fermi-Hubbard gate with interaction
    qc_trott.fhop(
        j=trott_hop * np.ones(Nsites - 1),
        modes=range(2 * Nsites),
    )
    qc_trott.fint(u=trott_int, modes=range(2 * Nsites))
# apply a global Fermi-Hubbard gate with interaction
qc_trott.measure_all()

circuit_trott = qc_trott.bind_parameters({trott_hop: J * time_step, trott_int: U * time_step})
job_trott = backend_ffsim.run(circuit_trott, shots=5, num_species=2)
start_time = time.time()
count_long = job_long.result().get_counts()
end_time = time.time()
print(f"Duration for of the two jobs: {(end_time - start_time)*1e3:.2f} ms")

import qiskit.tools.jupyter

%qiskit_version_table