from hops.core.hierarchy_parameters import HIParams, HiP, IntP, SysP, ResultType
from stocproc.stocproc import StocProc_TanhSinh
from hops.util.bcf_fits import get_ohm_g_w
from hops.util.truncation_schemes import TruncationScheme_Power_multi
import hops.util.bcf
import numpy as np
import hops.util.matrixLib as ml
from stocproc import StocProc_FFT

np.__config__.blas_opt_info = np.__config__.blas_ilp64_opt_info  # fix for qutip
import qutip

max_HO_level = 6
wc = 2
s = 1
Ω = 1
Λ = 1
γ = 0.1

# The BCF fit
bcf_terms = 4
g, w = get_ohm_g_w(bcf_terms, s, wc)

integration = IntP(t_max=40, t_steps=int(40 // 0.1))

q_proto = qutip.operators.create(max_HO_level) + qutip.operators.destroy(max_HO_level)
p_proto = (
    qutip.operators.destroy(max_HO_level) - qutip.operators.create(max_HO_level)
) / 1j

q_1 = qutip.tensor(q_proto, qutip.qeye(max_HO_level))
p_1 = qutip.tensor(p_proto, qutip.qeye(max_HO_level))
q_2 = qutip.tensor(qutip.qeye(max_HO_level), q_proto)
p_2 = qutip.tensor(qutip.qeye(max_HO_level), p_proto)

system = SysP(
    H_sys=(
        0.25 * (p_1 ** 2 + q_1 ** 2 + p_2 ** 2 + q_2 ** 2)
        + γ / 4 * (q_1 ** 2 - q_2 ** 2)
    ).full(),
    L=[0.5 * q_1.full(), 0.5 * q_2.full()],
    psi0=qutip.tensor(
        qutip.states.fock(max_HO_level, n=0), qutip.states.fock(max_HO_level, n=0)
    )
    .full()
    .flatten(),
    g=[g, g],
    w=[w, w],
    bcf_scale=[0.5, 0.5],
    T=[0.3, 0],
)

hops_bcf = hops.util.bcf.OhmicBCF_zeroTemp(
    s,
    1,
    wc,
)

params = HIParams(
    SysP=system,
    IntP=integration,
    HiP=HiP(
        nonlinear=True,
        result_type=ResultType.ZEROTH_AND_FIRST_ORDER,
        truncation_scheme=TruncationScheme_Power_multi.from_g_w(
            g=system.g, w=system.w, p=[1, 1], q=[0.5, 0.5], kfac=[1.6, 1.6]
        ),
        save_therm_rng_seed=True,
    ),
    Eta=[
        StocProc_FFT(
            spectral_density=hops.util.bcf.OhmicSD_zeroTemp(
                s,
                1,
                wc,
            ),
            alpha=hops_bcf,
            t_max=integration.t_max,
            intgr_tol=1e-5,
            intpl_tol=1e-5,
            negative_frequencies=False,
        ),
        StocProc_FFT(
            spectral_density=hops.util.bcf.OhmicSD_zeroTemp(
                s,
                1,
                wc,
            ),
            alpha=hops_bcf,
            t_max=integration.t_max,
            intgr_tol=1e-5,
            intpl_tol=1e-5,
            negative_frequencies=False,
        ),
    ],
    EtaTherm=[
        StocProc_TanhSinh(
            spectral_density=hops.util.bcf.Ohmic_StochasticPotentialDensity(
                s, 1, wc, beta=1 / system.__non_key__["T"][0]
            ),
            alpha=hops.util.bcf.Ohmic_StochasticPotentialCorrelations(
                s, 1, wc, beta=1 / system.__non_key__["T"][0]
            ),
            t_max=integration.t_max,
            intgr_tol=1e-5,
            intpl_tol=1e-5,
            negative_frequencies=False,
        ),
        None,
    ],
)