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

wc = 2
s = 1

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

integration = IntP(t_max=30, t_steps=int(20 // 0.01))
system = SysP(
    H_sys=0.5 * np.array([[-1, 0], [0, 1]]),
    L=0.5 * np.array([[0, 1], [1, 0]]),
    psi0=np.array([0, 1]),
    g=g,
    w=w,
    bcf_scale=0.5,
    T=1.5,
)

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=.5 * 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.util.bcf.OhmicBCF_zeroTemp(
            s,
            1,
            wc,
        ),
        t_max=integration.t_max,
        intgr_tol=1e-5,
        intpl_tol=1e-5,
        negative_frequencies=False,
    ),
    EtaTherm=StocProc_FFT(
      spectral_density=hops.util.bcf.Ohmic_StochasticPotentialDensity(
          s, 1, wc, beta=1 / system.__non_key__["T"]
      ),
      alpha=hops.util.bcf.Ohmic_StochasticPotentialCorrelations(
          s, 1, wc, beta=1 / system.__non_key__["T"]
      ),
      t_max=integration.t_max,
      intgr_tol=1e-5,
      intpl_tol=1e-5,
      negative_frequencies=False,
  ),
)