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,
    BathMemory,
)
from hops.util.abstract_truncation_scheme import (
    TruncationScheme_Simplex
)
import hops.util.bcf
import numpy as np
from stocproc import StocProc_FFT

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


wc = 2
s = 1
Ω = 1
Λ = 1
γ = 1


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


max_HO_level_set = 10
def make_config(
    max_HO_level: int,
    bcf_terms: int = 4,
    t_max: float = 10,
    k_fac: float = 1.4,
    inf_tol: float = 0.1,
        k_max: int = 7,
    sp_tol: float = 1e-3,
    bcf_scale: list[float] = [0.5, 0.5],
    truncation_scheme: str = "power",
    T: float = 0.3,
):
    global max_HO_level_set
    max_HO_level_set = max_HO_level

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

    integration = IntP(t_max=t_max, t_steps=int(t_max // 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 - 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=bcf_scale,
        T=[T, 0],
    )

    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=[k_fac] * 2
            )
            if truncation_scheme == "power"
            else BathMemory.from_system(
                system, nonlinear=True, influence_tolerance=inf_tol
            )
            if truncation_scheme == "bath_memory"
            else TruncationScheme_Simplex(k_max),
            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=sp_tol,
                intpl_tol=sp_tol,
                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=sp_tol,
                intpl_tol=sp_tol,
                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=sp_tol,
                intpl_tol=sp_tol,
                negative_frequencies=False,
            ),
            None,
        ],
    )

    return params

γ = .5
params = make_config(
   max_HO_level=9,
   bcf_terms=5,
   t_max=50,
   inf_tol=.05,
   sp_tol=1e-5,
   bcf_scale=[0.2, 0.2],
   T=.6,
   k_max=5,
   truncation_scheme="simplex",
)