%load_ext autoreload
%autoreload 2
%load_ext jupyter_spaces

import hops
import stocproc as sp
import numpy as np
import matplotlib.pyplot as plt
import scipy
plt.style.use('ggplot')
from IPython.display import set_matplotlib_formats
set_matplotlib_formats('pdf', 'svg')

σ1 = np.matrix([[0,1],[1,0]])
σ2 = np.matrix([[0,-1j],[1j,0]])
σ3 = np.matrix([[1,0],[0,-1]])

γ = 5 # coupling ratio
ω_c = 1 # center of spect. dens
δ = 2 # breadth BCF
W = -ω_c * 1j - δ # exponent BCF
τ_max = 40 # the maximal simulation time
seed = 1 # seed for all pseudo random generators
H_s = σ3
L = 1/2 * (σ1 - 1j * σ2) * γ

def α(τ):
    return np.exp(-1j * ω_c * τ - np.abs(τ) * δ)

def I(ω):
    return 1 / (δ + (ω - ω_c) ** 2 / δ)

%%space plot
t = np.linspace(0, τ_max, 1000)
ω = np.linspace(ω_c - 10, ω_c + 10, 1000)
fig, axs = plt.subplots(2)
axs[0].plot(t, np.real(α(t)))
axs[0].plot(t, np.imag(α(t)))
axs[1].plot(ω, I(ω))

η = sp.StocProc_FFT(
    I, τ_max, α, negative_frequencies=True, seed=seed, intgr_tol=1e-2, intpl_tol=1e-2
)

%%space plot

η.new_process(seed=seed)
plt.plot(η.t, np.real(η(η.t)), label="Re")
plt.plot(η.t, np.imag(η(η.t)), linestyle="--", label="Im")
plt.ylabel("η")
plt.xlabel("τ")
plt.legend()

step = hops.make_hops_step(η, H_s, L, W, 3)
step(0, np.array([1, 1, 0, 0, 0, 0, 0, 0]))

res = hops.integrate_hops_trajectory(η, H_s, L, W, 15, [1, 0], τ_max, seed, atol=1e-8, rtol=1e-8)

%%space plot
t = np.linspace(0, τ_max, 1000)
plt.plot(t, np.abs(res.sol(t)[0]))

ts, ρs = hops.integrate_hops_ensemble(η, H_s, L, W, 10, [1, 0], τ_max, 10)

energy = [np.trace(ρ @ σ3 @ σ3)/np.trace(ρ)-(np.trace(ρ @ σ3)/np.trace(ρ))**2 for ρ in ρs]
plt.plot(ts, energy)