from hops.core.hierarchy_parameters import HIParams, HiP, IntP, SysP, ResultType
from hops.core.hierarchyLib import HI
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 = 15
wc = 2
s = 1

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

integration = IntP(t_max=20, t_steps=int(20 // 0.01))

q = (
    qutip.operators.create(max_HO_level) + qutip.operators.destroy(max_HO_level)
)
p = ((qutip.operators.destroy(max_HO_level) - qutip.operators.create(max_HO_level)) / 1j)

system = SysP(
    H_sys=0.25 * (p ** 2 + q ** 2).full(),
    L=0.5 * q.full(),
    psi0=qutip.states.fock(max_HO_level, n=1).full().flatten(),
    g=g,
    w=w,
    bcf_scale=0.2,
    T=0,
)

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

params = HIParams(
    SysP=system,
    IntP=integration,
    HiP=HiP(
        nonlinear=True,
        normalized_by_hand=True,
        result_type=ResultType.ZEROTH_AND_FIRST_ORDER,
        truncation_scheme=TruncationScheme_Power_multi.from_g_w(
            g=g, w=w, p=1, q=0.5, kfac=1.7
        ),
        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,
    ),
    EtaTherm=None,
)