import hierarchyLib
import hierarchyData
import numpy as np
import pickle
from scipy.special import gamma as gamma_func

from stocproc.stocproc import StocProc_FFT
import bcf

t_max = 25
alpha = 0.266
T = 2.09
numExpFit = 5
s = 1
wc = 0.531

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

_, g_tilde, w_tilde = good_fit_data_abs[(numExpFit, s)]
g = 1 / np.pi * gamma_func(s + 1) * wc ** (s + 1) * np.asarray(g_tilde)
w = wc * np.asarray(w_tilde)

bcf_scale = np.pi / 2 * alpha * wc ** (1 - s)


# setup hierarchy parameters, here use a simple triangular hierarchy up to level 5
# the comments show the default values
# NOTE! the normalized version of HOPS is still buggy !!!
HierachyParam = hierarchyData.HiP(
    k_max=5,
    # g_scale=None,
    # sample_method='random',
    # seed=0,
    # nonlinear=True,
    # normalized=False,
    # terminator=False,
    # result_type= hierarchyData.RESULT_TYPE_ZEROTH_ORDER_ONLY,
    # accum_only=None,
    # rand_skip=None
)

# setup the integration parameters (see scipy ode)
IntegrationParam = hierarchyData.IntP(
    t_max=t_max,
    t_steps=500,
    # integrator_name='zvode',
    # atol=1e-8,
    # rtol=1e-8,
    # order=5,
    # nsteps=5000,
    # method='bdf',
    # t_steps_skip=1
)

# setup the system parameters
# here a simple Spin Boson system, single qubit coupled to a sub-ohmic environment
# how to get the g, w is a different story
SystemParam = hierarchyData.SysP(
    H_sys=np.asarray([[0, 1], [1, 0]]),  # SIGMA X
    L=np.asarray([[-1, 0], [0, 1]]),  # SIGME Z
    psi0=np.asarray([0, 1]),  # excited qubit
    g=np.asarray(g),
    w=np.asarray(w),
    H_dynamic=[],
    bcf_scale=bcf_scale,  # some coupling strength (scaling of the fit parameters 'g_i')
    gw_hash=None,  # this is used to load g,w from some database
    len_gw=None,
)


# setup the stochastic process for the subohmic SD
# for each sample, the StocProc_FFT class is seeded differently which results in a different realization.
# for the HOPS solver, this looks like a simple callable z(t) returning complex values
Eta = StocProc_FFT(
    spectral_density=bcf.OSD(s=s, eta=1, gamma=wc),
    t_max=t_max,
    alpha=bcf.OBCF(s=s, eta=1, gamma=wc),
    intgr_tol=1e-3,
    intpl_tol=1e-3,
    seed=None,
    negative_frequencies=False,
    scale=bcf_scale,
)
# thermal noise
EtaTherm = StocProc_FFT(
    spectral_density=bcf.OFTDens(s=s, eta=1, gamma=wc, beta=1 / T),
    t_max=t_max,
    alpha=bcf.OFTCorr(s=s, eta=1, gamma=wc, beta=1 / T),
    intgr_tol=1e-3,
    intpl_tol=1e-3,
    seed=None,
    negative_frequencies=False,
    scale=bcf_scale,
)

# this guy must be digestable by the 'binfootprint' class, generating a unique binary representation of hi_key
# this allows to generate a hash value of the binary to access the produced data from some hdf5 data (see hierarchyData.py)
# NOTE! this requires that the ENTIRE sub-structure of 'hi_key' is digestable by 'binfootprint', therefore the special classes
# representing the spectral density and bath correlation function defined in bcf.py
# I developed 'binfootprit' since pickle and co. are not guaranteed to lead to unique binary data,
hi_key = hierarchyData.HIMetaKey_type(
    HiP=HierachyParam,
    IntP=IntegrationParam,
    SysP=SystemParam,
    Eta=Eta,
    EtaTherm=EtaTherm,
)

# here the hierarchy structure it set up ...
myHierarchy = hierarchyLib.HI(hi_key, number_of_samples=100, desc="run a test case")

# ... which allows to perform the integration, the data is stored in a file called 'myHierarchy.data'
# running the script again only crunches samples not done before!
myHierarchy.integrate_simple(data_name="myHierarchy.data")

# to access the data the 'hi_key' is used to find the data in the hdf5 file
with hierarchyData.HIMetaData(hid_name="myHierarchy.data", hid_path=".") as metaData:
    with metaData.get_HIData(hi_key, read_only=True) as data:
        smp = data.get_samples()
        print("{} samples found in database".format(smp))
        t = data.get_time()
        rho_t = data.get_rho_t()


# to show some dynamics
import matplotlib.pyplot as plt

plt.plot(t, rho_t[:, 1, 1].real, label="average ({})".format(smp))
plt.ylabel("rho_(0,0)")
plt.xlabel("time")
plt.legend()
plt.show()