import hierarchyLib
import hierarchyData
import numpy as np

from stocproc.stocproc import StocProc_FFT
import bcf
t_max = 25
coupling_strength = 0.3


# 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([-0.00783725-0.03065741j,       # fit for sub-ohmic SD / BCF
                                                 -0.00959114-0.12684515j,
                                                 0.12327218-0.34116637j,
                                                 0.71838859-0.3019772j,
                                                 0.52298016+0.66270003j,
                                                 -0.20912288+0.13755251j]),
                                 w = np.asarray([0.08064116+7.08122348e-03j,    # fit for sub-ohmic SD / BCF
                                                 0.40868336+4.73041920e-02j,
                                                 1.18403114+2.25925466e-01j,
                                                 2.81218754+9.72139059e-01j,
                                                 5.63558761+3.41859384e+00j,
                                                 9.49687769+9.75760546e+00j]),
                                 H_dynamic = [],
                                 bcf_scale = coupling_strength,                 # 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=0.25, eta=1, gamma=3),
                   t_max = t_max,
                   alpha = bcf.OBCF(s=0.25, eta=1, gamma=3),
                   intgr_tol=1e-2,
                   intpl_tol=1e-2,
                   seed=None,
                   negative_frequencies=False,
                   scale=coupling_strength)
# no thermal noise -> zero temperature, however, this works ion the same fashion as above
EtaTherm = None

# 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()
        psi_zt_0 = data.get_stoc_traj(idx = 0)


# to show some dynamics
import matplotlib.pyplot as plt
plt.plot(t, rho_t[:,0,0].real, label="average ({})".format(smp))
plt.plot(t, np.abs(psi_zt_0[:,0]), label="single psi_z_t")
plt.plot(t, np.sqrt(np.sum(np.abs(psi_zt_0)**2, axis=(1,))), label="|psi_z_t|", color='0.4')
plt.ylabel("rho_(0,0)")
plt.xlabel("time")
plt.legend()
plt.show()