import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
try:
    import sportran as st
except ImportError:
    from sys import path
    path.append('..')
    import sportran as st

c = plt.rcParams['axes.prop_cycle'].by_key()['color']

%matplotlib notebook

"""
Explanation: Example 2: Cepstral Analysis of liquid NaCl
This example shows the basic usage of sportran to compute the thermal conductivity of a classical MD simulation of molten NaCl, which requires the multi-component theory.
End of explanation
"""


jfile = st.i_o.TableFile('./data/NaCl.dat', group_vectors=True)

"""
Explanation: 1. Load trajectory
Read the heat current from a simple column-formatted file. The desired columns are selected based on their header (e.g. with LAMMPS format).
For other input formats see the corresponding example.
End of explanation
"""


jfile.read_datalines(start_step=0, NSTEPS=0, select_ckeys=['Temp', 'flux', 'vcm[1]'])

"""
Explanation: We need to load the energy flux (flux) and one mass flux (vcm[1], that is the velocity of the center of mass of one species).
End of explanation
"""


DT_FS = 5.0           # time step [fs]
TEMPERATURE = np.mean(jfile.data['Temp'])   # temperature [K]
VOLUME = 40.21**3     # volume [A^3]
print('T = {:f} K'.format(TEMPERATURE))
print('V = {:f} A^3'.format(VOLUME))

j = st.HeatCurrent([jfile.data['flux'], jfile.data['vcm[1]']], UNITS='metal', DT_FS=DT_FS,
                   TEMPERATURE=TEMPERATURE, VOLUME=VOLUME)

# trajectory
f, ax = plt.subplots(2, sharex=True)
ax[0].plot(j.timeseries()/1000., j.traj);
ax[1].plot(j.timeseries()/1000., j.otherMD[0].traj);
plt.xlim([0, 1.0])
plt.xlabel(r'$t$ [ps]')
ax[0].set_ylabel(r'$J^0$ [eV A/ps]');
ax[1].set_ylabel(r'$J^1$ [A/ps]');

"""
Explanation: 2. Heat Current
Define a HeatCurrent from the trajectory, with the correct parameters. The difference with respect to the single-component case is only here: we load a list of currents instead of a single one.
End of explanation
"""


# Periodogram with given filtering window width
ax = j.plot_periodogram(PSD_FILTER_W=0.4, kappa_units=True, label=r'$\bar{\mathcal{S}}^0_k$')
print(j.Nyquist_f_THz)

# compare with the spectrum of the energy flux
jen = st.HeatCurrent(jfile.data['flux'], UNITS='metal', DT_FS=DT_FS,
                     TEMPERATURE=TEMPERATURE, VOLUME=VOLUME)
ax = jen.plot_periodogram(axes=ax, PSD_FILTER_W=0.4, kappa_units=True, label=r'$\mathcal{S}^0_k$')
plt.xlim([0, 20])
ax[0].set_ylim([0, 0.8]);
ax[1].set_ylim([7, 18]);
ax[0].legend(); ax[1].legend();

"""
Explanation: Compute the Reduced periodogram $\bar{\mathcal{S}}^0_k$ and filter it for visualization. You can notice the difference with respect to the energ-flux periodogram $\mathcal{S}^0_k$.
End of explanation
"""


FSTAR_THZ = 14.0
jf, ax = j.resample(fstar_THz=FSTAR_THZ, plot=True, freq_units='thz')
plt.xlim([0, 20])
ax[1].set_ylim([7, 18]);

ax = jf.plot_periodogram(PSD_FILTER_W=0.1)
ax[1].set_ylim([7, 18]);

"""
Explanation: 3. Resampling
If the Nyquist frequency is very high (i.e. the sampling time is small), such that the log-spectrum goes to low values, you may want resample your time series to obtain a maximum frequency $f^$.
Before performing that operation, the time series is automatically filtered to reduce the amount of aliasing introduced. Ideally you do not want to go too low in $f^$. In an intermediate region the results should not change. 
To perform resampling you can choose the resampling frequency $f^$ or the resampling step (TSKIP). If you choose $f^$, the code will try to choose the closest value allowed.
The resulting PSD is visualized to ensure that the low-frequency region is not affected.
End of explanation
"""


jf.cepstral_analysis()

# Cepstral Coefficients
print('c_k = ', jf.dct.logpsdK)

ax = jf.plot_ck()
ax.set_xlim([0, 25])
ax.set_ylim([-0.2, 1.0])
ax.grid();

# AIC function
f = plt.figure()
plt.plot(jf.dct.aic, '.-', c=c[0])
plt.xlim([0, 50])
plt.ylim([1400, 1500]);

print('K of AIC_min = {:d}'.format(jf.dct.aic_Kmin))
print('AIC_min = {:f}'.format(jf.dct.aic_min))

"""
Explanation: 4. Cepstral Analysis
Perform Cepstral Analysis. The code will:
 1. the parameters describing the theoretical distribution of the PSD are computed
 2. the Cepstral coefficients are computed by Fourier transforming the log(PSD)
 3. the Akaike Information Criterion is applied
 4. the resulting $\kappa$ is returned
End of explanation
"""


# L_0 as a function of cutoff K
ax = jf.plot_L0_Pstar()
ax.set_xlim([0, 50])
ax.set_ylim([14.75, 16.5]);

print('K of AIC_min = {:d}'.format(jf.dct.aic_Kmin))
print('AIC_min = {:f}'.format(jf.dct.aic_min))

# kappa as a function of cutoff K
ax = jf.plot_kappa_Pstar()
ax.set_xlim([0, 50])
ax.set_ylim([0, 1.0]);

print('K of AIC_min = {:d}'.format(jf.dct.aic_Kmin))
print('AIC_min = {:f}'.format(jf.dct.aic_min))

"""
Explanation: Plot the thermal conductivity $\kappa$ as a function of the cutoff $P^*$
End of explanation
"""


results = jf.cepstral_log
print(results)

"""
Explanation: Print the results :)
End of explanation
"""


# filtered log-PSD
ax = j.plot_periodogram(0.5, kappa_units=True)
ax = jf.plot_periodogram(0.5, axes=ax, kappa_units=True)
ax = jf.plot_cepstral_spectrum(axes=ax, kappa_units=True)
ax[0].axvline(x = jf.Nyquist_f_THz, ls='--', c='r')
ax[1].axvline(x = jf.Nyquist_f_THz, ls='--', c='r')
plt.xlim([0, 20])
ax[0].set_ylim([0, 0.8]);
ax[1].set_ylim([7, 18]);
ax[0].legend(['original', 'resampled', 'cepstrum-filtered'])
ax[1].legend(['original', 'resampled', 'cepstrum-filtered']);

"""
Explanation: You can now visualize the filtered PSD...
End of explanation
"""


