import matplotlib.pyplot as plt
from hopsflow import gaussflow_two as gf
import utilities as ut
import hops.util.bcf
import numpy as np

t_max = 40
bcf_1 = hops.util.bcf.OhmicBCF_zeroTemp(
    1,
    1,
    1,
)

bcf_2 = hops.util.bcf.OhmicBCF_zeroTemp(
    1,
    1,
    1,
)

α_0_1 = gf.BCF(
    t_max,
    bcf_1,
    num_terms=4,
    resolution=0.01,
)

α_1 = gf.BCF(
    t_max,
    hops.util.bcf.OhmicBCF_nonZeroTemp(s=1, eta=1, w_c=1, beta=1 / .8),
    num_terms=8,
    resolution=0.01,
)

α_0_2 = gf.BCF(
    t_max,
    bcf_2,
    num_terms=4,
    resolution=0.01,
)

params = gf.SystemParams(Ω=1, Λ=1, η=[1, 1], γ=0, α_0=[α_0_1, α_0_2])
t_points = np.linspace(0, t_max, 500)

G = gf.Propagator(params)

initial_state = np.array([1,0,2,0])
traj = G(t_points) @ initial_state
with ut.hiro_style():
    plt.plot(t_points, traj[:, [0,2]])

C = gf.CorrelationMatrix(params,gf.initial_correlation_pure_osci(1 , 1), [α_1, None])

energy = C.system_energy(t_points)

with ut.hiro_style():
    plt.plot(t_points, energy)

flow = [C.flow(t_points, i) for i in range(2)]

with ut.hiro_style():
    plt.plot(t_points, flow[0])
    plt.plot(t_points, flow[1])

flow[0][-4]/flow[1][-4]