# native
import pathlib
import traceback
import pickle

# third party
import numpy as np
from scipy.special import gamma as gamma_func
import qutip

# hops
import hops.core.hierarchyData as hid
import hops.util.bcf as bcf
import hops.util.matrixLib as ml
import hops.util.truncation_schemes as truncation_schemes
import hops.stg_helper as stg_helper


from hops.data.path import get_path as get_hops_data_path

################################################
#
# OVERALL PARAMETERS
#
# may be overwritten in later definitions
#
################################################

max_HO_level = 20
OMEGA = 1

wc = 2
s = 1
__BCF_FIT_n = 5
__T = 0.5

__HI_eta_type = "fft"
__HI_eta_temperature_type = "fft"

eta = 0.2

p_trunc = 1
q_trunc = 0.5

n_eq = None
kfac = 1.4
kmax_list_pre_sqrt = True

__HI_k_max = None

t_max = 25

sp_tol = 1e-6
__HI_seed = 0


__HI_accum_only = False
__HI_number_of_samples = 10_000


__BCF_FIT_good_data_path = get_hops_data_path() + "/good_fit_data_abs_brute_force"
with open(__BCF_FIT_good_data_path, "rb") as f:
    good_fit_data_abs = pickle.load(f)

try:
    _, g_tilde, w_tilde = good_fit_data_abs[(__BCF_FIT_n, s)]
    __g = 1 / np.pi * gamma_func(s + 1) * wc ** (s + 1) * np.asarray(g_tilde)
    __w = wc * np.asarray(w_tilde)
    __HI_truncation_scheme = stg_helper.get_truncation_scheme(
        kmax_list_pre_sqrt=kmax_list_pre_sqrt,
        g=__g,
        w=__w,
        q=q_trunc,
        p=p_trunc,
        kfac=kfac,
    )
except KeyError:
    raise KeyError(
        "no fit data in {} for n={} and s={}".format(
            __BCF_FIT_good_data_path, __BCF_FIT_n, s
        )
    )

__BCF_FIT_gw_hash = None
__temp_method = "stoc_pot"

__bcf_scale = eta
# print("eta", __bcf_scale)
__bcf_zero_temp = bcf.OhmicBCF_zeroTemp(s=s, eta=1, w_c=wc)
__spd_zero_temp = bcf.OhmicSD_zeroTemp(s=s, eta=1, w_c=wc)

if __T == 0:
    __bcf = __bcf_zero_temp
    __spec_dens = __spd_zero_temp
    __stoc_temp_corr = None
    __stoc_temp_dens = None
    beta = np.inf
    __HI_save_therm_rng_seed = False
else:
    __bcf = bcf.OhmicBCF_nonZeroTemp(s=s, eta=1, w_c=wc, beta=1 / __T)
    __spec_dens = bcf.PseudoSD(sd_at_t_zero=__spd_zero_temp, T=__T)
    __stoc_temp_corr = bcf.Ohmic_StochasticPotentialCorrelations(
        s=s, eta=1, w_c=wc, beta=1 / __T
    )
    __stoc_temp_dens = bcf.Ohmic_StochasticPotentialDensity(
        s=s, eta=1, w_c=wc, beta=1 / __T
    )
    beta = 1 / __T
    __HI_save_therm_rng_seed = True

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
__L = 0.5 * q
__H_sys = 0.25 * OMEGA * (p ** 2 + q ** 2)

__L = ml.Operator(data=__L, fmt="coo")
__H_sys = ml.Operator(data=__H_sys, fmt="coo")


#################################################
#
# settings for stoc proc
#
#################################################
__SP_t_max = t_max

__SP_db_name = "SPCache"
__SP_db_path = "."

__SP_KLE_tol = sp_tol
__SP_KLE_ng_fac = 4
__SP_KLE_meth = "fourpoint"
# __SP_KLE_meth = 'trapz'
# __SP_KLE_diff_method = 'full'
# __SP_KLE_dm_random_samples = None

__SP_KLE_diff_method = "random"
__SP_KLE_dm_random_samples = 10000


if __temp_method == "stoc_pot":
    __SP_FFT_neg_frequ = False
else:
    if __T == 0:
        __SP_FFT_neg_frequ = False
    else:
        __SP_FFT_neg_frequ = True

# __SP_FFT_intgr_tol = 5e-3
# __SP_FFT_intpl_tol = 5e-3

__SP_FFT_intgr_tol = sp_tol
__SP_FFT_intpl_tol = sp_tol

################################################
#
# System and Hierarchy parameters
#
################################################

# __g = None
# __w = None
__H_dyn = []

__psi0 = np.array(qutip.states.fock(max_HO_level, n=0).data.todense()).flatten()

# for image: stoch_trj.png
# __INTGR_t_max = 6
# __INTGR_t_steps = 750
__INTGR_t_max = t_max
__INTGR_t_steps = int(t_max / 0.01)

__INTGR_name = "zvode"
__INTGR_atol = "1e-8"
__INTGR_rtol = "1e-8"
__INTGR_order = 5
__INTGR_nsteps = 10_000
__INTGR_method = "bdf"

__HI_g_scale = None
__HI_sample_method = "random"

__HI_nonlinear = True
__HI_normalized = False
__HI_normalized_by_hand = True
__HI_with_terminator = False
# __HI_result_data = hid.RESULT_TYPE_ZEROTH_ORDER_ONLY
__HI_result_data = hid.ResultType.ZEROTH_AND_FIRST_ORDER


#################################################
#
# DB and Server/Client parameters
#
#################################################
__db_name = "SBM"
__db_path = "."
__host = "localhost"
__verbose = 1
__authkey = __db_name
__desc = __db_name
__port = 35254
__nproc = 0
__nice = 19