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

from scipy.linalg import expm
from pyqcs import State, sample

from transfer_matrix import T_time_slice
from hamiltonian import H
from bootstrap import bootstrap


np.random.seed(0xdeadbeef)

matplotlib.rcParams.update(
        {'errorbar.capsize': 2
        , 'figure.figsize': (16, 9)}
        )

nqbits = 6
g = 3
N_trot = 80
t_stop = 9
delta_t = 0.09
qbits = list(range(nqbits))

n_sample = 2200
measure = 0b10

measure_coefficient_mask = [False if (i & measure) else True for i in range(2**nqbits)]


results_qc = []
results_np = []
errors_sampling = []

print()
for t in np.arange(0, t_stop, delta_t):
    # QC simulation
    state = State.new_zero_state(nqbits)

    T_dt = T_time_slice(qbits, t, g, N_trot)
    for _ in range(N_trot):
        state = T_dt * state

    result = sample(state, measure, n_sample)
    results_qc.append(result[0] / n_sample)

    errors_sampling.append(bootstrap(result[0], n_sample, n_sample // 4, n_sample // 10, np.average))

    #amplitude = np.sqrt(np.sum(np.abs(state._qm_state[measure_coefficient_mask])**2))
    #results_qc.append(amplitude)

    # Simulation using matrices
    np_zero_state = np.zeros(2**nqbits)
    np_zero_state[0] = 1

    itH = np.matrix(-0.5j * t * H(nqbits, g))
    T = expm(itH)

    np_state = T.dot(np_zero_state)
    amplitude = (np.sum(np.abs(np_state[measure_coefficient_mask])**2))
    results_np.append(amplitude)

    print(f"simulating... {int(t/t_stop*100)} %  ", end="\r")
print()
print("done.")

results_qc = np.array(results_qc)

errors_trotter = (np.arange(0, t_stop, delta_t) * g)**2 / N_trot**2
errors_sampling = np.array(errors_sampling)


h0 = plt.errorbar(np.arange(0, t_stop, delta_t)
                , results_qc
                , yerr=(errors_trotter + errors_sampling)
                , label=f"Quantum computing ({n_sample} samples, {N_trot} trotterization steps)"
                , )
h1, = plt.plot(np.arange(0, t_stop, delta_t), results_np, label="Classical simulation using explicit transfer matrix")
plt.xlabel("t")
plt.ylabel(r"$|0\rangle$ probability amplitude for second spin")
plt.title(f"{nqbits} site spin chain with g={g} coupling to external field")
plt.legend(handles=[h0, h1])

plt.savefig("time_evo_6spin_g3.png", dpi=400)
plt.show()