# Load data
mu_x, mu_y, mu_z = np.loadtxt('Data/displaced_dipole.H3+.dat').T

# Plot
from matplotlib.ticker import MaxNLocator

fig, ax = plt.subplots()

x = np.arange(len(mu_x))

ax.yaxis.set_major_locator(MaxNLocator(integer=True))
ax.set_xticks([x[0], x[-1]/3, 2*x[-1]/3, x[-1]])
ax.set_xticklabels(['I', 'II', 'III', 'I'])

lw = 2.5
ax.plot(x, mu_x, label = r'$x$', lw = lw)
ax.plot(x, mu_y, label = r'$y$', lw = lw)
ax.plot(x, mu_z, label = r'$z$', lw = lw)

delta = x[-1]-x[0]
ax.set_xlim(x[0] - 0.05*delta, x[-1] + 0.05*delta)

M = np.max([mu_x, mu_y, mu_z])
m = np.min([mu_x, mu_y, mu_z])
delta = M - m
ax.set_ylim(m - 0.05*delta, M + 0.05*delta)

ax.set_xlabel('Path')
ax.set_ylabel(r'$\Delta \mu /(Le)$')

ax.legend();

fig.savefig('Images/figure2c.png')

# Load data
x, y, z, mu_x, mu_y, mu_z = np.loadtxt('Data/displaced_dipole.K3Cl.dat').T

# Plot
from matplotlib.ticker import MaxNLocator

fig, ax = plt.subplots()

ax.yaxis.set_major_locator(MaxNLocator(integer=True))
ax.set_xticks([x[0], x[-1]])
ax.set_xticklabels(['I', 'I'])


# x = np.arange(len(mu_x))
lw = 2.5
ax.plot(x, mu_x, label = r'$x$', lw = lw)
ax.plot(y, mu_y, label = r'$y$', lw = lw)
ax.plot(z, mu_z, label = r'$z$', lw = lw)

delta = x[-1]-x[0]
ax.set_xlim(x[0] - 0.05*delta, x[-1] + 0.05*delta)

M = np.max([mu_x, mu_y, mu_z])
m = np.min([mu_x, mu_y, mu_z])
delta = M - m
ax.set_ylim(m - 0.05*delta, M + 0.05*delta)

ax.set_xlabel('Path')
ax.set_ylabel(r'$\Delta \mu /(Le)$')

ax.legend();

fig.savefig('Images/figure3c.png')

# Load file with conductivity
sigma_file = np.loadtxt('Data/sigma.K-KCl.core_and_wcs.dat').T

# Define timestep
dt = TAUPS*15*20*1000
# Define simulation cell linear size
ALAT = 14.070816430e-10 # meters
# Compute average temperature
TEMP_h = np.mean(np.loadtxt('Data/current.K-KCl.core_and_wcs.dat', skiprows=1).T[1])
# Conversion factor from angstroms to meters 
ang2m = 1e-10
# Conversion factor to obtain the conductivity in S/cm
conv = (ELCH*ang2m)**2 * dt*1e9 * (1/ALAT**3/TEMP_h/k_BOLTZ) / 100


sigma_core_wcs = np.zeros((2, len(sigma_file[0])-2))
sigma_core_wcs[0] = sigma_file[8, 2:]*conv
sigma_core_wcs[1] = np.sqrt(sigma_file[9, 2:])*conv

time = np.arange(1, len(sigma_core_wcs[0])+1)*dt
fitto = int(4.5/6*len(time))

sigma_core_wcs[0] *= time
sigma_core_wcs[1] *= time

(m, q), _ = curve_fit(linear, time[fitto:], sigma_core_wcs[0][fitto:],
                      sigma = (sigma_core_wcs[1][fitto:]))
sig_core_wcs = m
q_core_wcs = q

# Estimate error on sigma as the semi-dispersion of the maximum and minimum fit
(m, q), _ = curve_fit(linear, time[fitto:], sigma_core_wcs[0][fitto:] + sigma_core_wcs[1][fitto:])
dsig1 = m
(m, q), _ = curve_fit(linear, time[fitto:], sigma_core_wcs[0][fitto:] - sigma_core_wcs[1][fitto:])
dsig2 = m
dsig_core_wcs = (dsig1 - dsig2)/2

print('sigma[cores+WCs] = {:4.2f} +\- {:4.2f} S/cm'.format(sig_core_wcs, dsig_core_wcs))

# Load file with conductivity
sigma_file = np.loadtxt('Data/sigma.K-KCl.nominal_charges_and_lone_pair.dat').T

# Define timestep
dt = TAUPS*15*20*1000
# Define simulation cell linear size
ALAT = 14.070816430e-10 # meters
# Compute average temperature
TEMP_h = np.mean(np.loadtxt('Data/current.K-KCl.nominal_charges_and_lone_pair.dat', skiprows=1).T[1])
# Conversion factor from angstroms to meters 
ang2m = 1e-10
# Conversion factor to obtain the conductivity in S/cm
conv = (ELCH*ang2m)**2 * dt * 1e9 * (1/ALAT**3/TEMP_h/k_BOLTZ) / 100

# Load data
sigma_ions_lp = np.zeros((2, len(sigma_file[0])-2))
sigma_ions_lp[0] = sigma_file[8, 2:]*conv
sigma_ions_lp[1] = np.sqrt(sigma_file[9, 2:])*conv

# 
time = np.arange(1, len(sigma_ions_lp[0])+1)*dt
fitto = int(4.5/6*len(time))

sigma_ions_lp[0] *= time
sigma_ions_lp[1] *= time

(m, q), _ = curve_fit(linear, time[fitto:], sigma_ions_lp[0][fitto:],
                      sigma = (sigma_ions_lp[1][fitto:]))
sig_ions_lp = m
q_ions_lp = q

# Estimate error on sigma as the semi-dispersion of the maximum and minimum fit
(m, q), _ = curve_fit(linear, time[fitto:], sigma_ions_lp[0][fitto:] + sigma_ions_lp[1][fitto:])
dsig1 = m
(m, q), _ = curve_fit(linear, time[fitto:], sigma_ions_lp[0][fitto:] - sigma_ions_lp[1][fitto:])
dsig2 = m
dsig_ions_lp = (dsig1 - dsig2)/2

print('sigma[ionic charges + lone pair] = {:4.2f} +\- {:4.2f} S/cm'.format(sig_ions_lp, dsig_ions_lp))

# Load file with conductivity
sigma_file = np.loadtxt('Data/sigma.K-KCl.cross.dat').T

# Define timestep
dt = TAUPS*15*20*1000
# Define simulation cell linear size
ALAT = 14.070816430e-10 # meters
# Compute average temperature
TEMP_h = np.mean(np.loadtxt('Data/current.K-KCl.nominal_charges_and_lone_pair.dat', skiprows=1).T[1])
# Conversion factor from angstroms to meters 
ang2m = 1e-10
# Conversion factor to obtain the conductivity in S/cm
conv = (ELCH*ang2m)**2 * dt*1e9 * (1/ALAT**3/TEMP_h/k_BOLTZ) / 100

sigma_cross = np.zeros((2, len(sigma_file[0])-2))
sigma_cross[0] = sigma_file[12, 2:]*conv
sigma_cross[1] = np.sqrt(sigma_file[13, 2:])*conv

time = np.arange(1, len(sigma_cross[0])+1)*dt
fitto = int(4.5/6*len(time))

sigma_cross[0] *= time
sigma_cross[1] *= time

(m, q), _ = curve_fit(linear, time[fitto:], sigma_cross[0][fitto:],
                      sigma = (sigma_cross[1][fitto:]))
sig_cross = m
q_cross = q

# Estimate error on sigma as the semi-dispersion of the maximum and minimum fit
(m, q), _ = curve_fit(linear, time[fitto:], sigma_cross[0][fitto:] + sigma_cross[1][fitto:])
dsig1 = m
(m, q), _ = curve_fit(linear, time[fitto:], sigma_cross[0][fitto:] - sigma_cross[1][fitto:])
dsig2 = m
dsig_cross = (dsig1 - dsig2)/2

print('sigma[cross-term] = {:4.2f} +\- {:4.2f} S/cm'.format(sig_cross, dsig_cross))

# Load file with conductivity
sigma_file = np.loadtxt('Data/sigma.K-KCl.ionic.dat').T

# Define timestep
dt = TAUPS*15*20*1000
# Define simulation cell linear size
ALAT = 14.070816430e-10 # meters
# Compute average temperature
TEMP_h = np.mean(np.loadtxt('Data/current.K-KCl.nominal_charges_and_lone_pair.dat', skiprows=1).T[1])
# Conversion factor from angstroms to meters 
ang2m = 1e-10
# Conversion factor to obtain the conductivity in S/cm
conv = (ELCH*ang2m)**2 * dt*1e9 * (1/ALAT**3/TEMP_h/k_BOLTZ) / 100


sigma_ions = np.zeros((2, len(sigma_file[0])-2))
sigma_ions[0] = sigma_file[8, 2:]*conv
sigma_ions[1] = np.sqrt(sigma_file[9, 2:])*conv

time = np.arange(1, len(sigma_ions[0])+1)*dt
fitto = int(4.5/6*len(time))

sigma_ions[0] *= time
sigma_ions[1] *= time

(m, q), _ = curve_fit(linear, time[fitto:], sigma_ions[0][fitto:],
                      sigma = (sigma_ions[1][fitto:]))
sig_ions = m
q_ions = q

# Estimate error on sigma as the semi-dispersion of the maximum and minimum fit
(m, q), _ = curve_fit(linear, time[fitto:], sigma_ions[0][fitto:] + sigma_ions[1][fitto:])
dsig1 = m
(m, q), _ = curve_fit(linear, time[fitto:], sigma_ions[0][fitto:] - sigma_ions[1][fitto:])
dsig2 = m
dsig_ions = (dsig1 - dsig2)/2

print('sigma[ionic contribution] = {:4.2f} +\- {:4.2f} S/cm'.format(sig_ions, dsig_ions))

fig, ax = plt.subplots()

#######################################################
s = sigma_core_wcs
m = sig_core_wcs
q = q_core_wcs

time = np.arange(len(s[0]))*dt
pf = ax.plot(time/1000, s[0]/1000, lw=0)
pffill = ax.fill_between(time/1000, (s[0] - s[1])/1000,
                (s[0] + s[1])/1000, color = pf[0].get_color(), alpha = 0.5)
pffit = ax.plot(time/1000, linear(time, m, q)/1000,
        ls = '--',
        lw = 3,
        color = pf[0].get_color())
#######################################################
s = sigma_ions_lp
m = sig_ions_lp
q = q_ions_lp

time = np.arange(len(s[0]))*dt
p = ax.plot(time/1000, s[0]/1000, lw=0)
pfill = ax.fill_between(time/1000, (s[0] - s[1])/1000,
                (s[0] + s[1])/1000, color = p[0].get_color(), alpha = 0.5)
pfit = ax.plot(time/1000, linear(time, m, q)/1000,
        ls = '--',
        lw = 3,
        color = p[0].get_color())
#######################################################
s = sigma_ions
m = sig_ions
q = q_ions

time = np.arange(len(s[0]))*dt
pi = ax.plot(time/1000, s[0]/1000, lw=0)
pifill = ax.fill_between(time/1000, (s[0] - s[1])/1000,
                (s[0] + s[1])/1000, color = pi[0].get_color(), alpha = 0.5)
pifit = ax.plot(time/1000, linear(time, m, q)/1000,
        ls = '--',
        lw = 3,
        color = pi[0].get_color())
#######################################################
s = sigma_cross
m = sig_cross
q = q_cross

time = np.arange(len(s[0]))*dt
pc = ax.plot(time/1000, s[0]/1000, lw=0)
pcfill = ax.fill_between(time/1000, (s[0] - s[1])/1000,
                (s[0] + s[1])/1000, color = pc[0].get_color(), alpha = 0.5)
pcfit = ax.plot(time/1000, linear(time, m, q)/1000,
        ls = '--',
        lw = 3,
        color = pc[0].get_color())
#######################################################
ax.legend([(pffill, pffit[0]), (pfill, pfit[0]), (pifill, pifit[0]), (pcfill, pcfit[0])],
              [r'From Eq. (2)', r'From Eq. (4)', r'Ionic contribution', r'Cross-term'],
              loc = 'upper left')

ax.set_xlabel(r'Time (ps)')
ax.set_ylabel(r'$\langle | \Delta \mu | ^2 \rangle / (6 k_B T L^3)$ (S cm$^{-1}$ ps)')

# ax.tick_params(axis='both', which='major', labelsize = fs)
ax.set_xticks(range(4))
ax.set_yticks(range(0, 80, 20))
ax.set_xlim(0,3.4)
ax.set_ylim(-5,75)

fig.savefig('Images/figure8.png')

wind = 0.1

dt = TAUPS*15*20*1000
ALAT = 14.070816430e-10 # metres
VOL = 14.070816430**3
fstar = 35

current_file = np.loadtxt('Data/current.K-KCl.core_and_wcs.dat', skiprows=1).T[1:]
TEMP = np.mean(current_file[0])

conv = np.sqrt(TEMP)/10

j_core_wcs = current_file[1:].T * conv

hc_core_wcs = tc.HeatCurrent(j_core_wcs, 'metal', dt, TEMP, VOL, wind)
hcr_core_wcs = hc_core_wcs.resample_current(fstar_THz = fstar, plot = True, PSD_FILTER_W = wind)[0]
hcr_core_wcs.cepstral_analysis(Kmin_corrfactor = 0.8)

wind=0.1

dt = TAUPS*15*20*1000
ALAT = 14.070816430e-10 # metres
VOL = 14.070816430**3
fstar = 35

current_file = np.loadtxt('Data/current.K-KCl.nominal_charges_and_lone_pair.dat', skiprows=1).T[1:]
TEMP = np.mean(current_file[0])

conv = np.sqrt(TEMP)/10

j_ions_lp = current_file[1:].T * conv

hc_ions_lp = tc.HeatCurrent(j_ions_lp, 'metal', dt, TEMP, VOL, wind)
hcr_ions_lp = hc_ions_lp.resample_current(fstar_THz = fstar, plot = True, PSD_FILTER_W = wind)[0]
hcr_ions_lp.cepstral_analysis(Kmin_corrfactor = 0.8)

fig, ax = plt.subplots()

alpha = 0.3
fmax = fstar
ymax=25

f = hcr_core_wcs
conv = f.kappa_scale/2
psd = ax.plot(f.freqs_THz, conv * f.psd,
              alpha = alpha, zorder = 0, lw=0)
fpsd = ax.plot(f.freqs_THz, conv * f.fpsd,
               color = psd[0].get_color(), zorder = 1)
ceps = ax.plot(f.freqs_THz, conv * f.dct.psd, lw = 2,
               label="From Eq. (2)",
               color = psd[0].get_color(), zorder = 2)
###########################################################################################
f = hcr_ions_lp
conv = f.kappa_scale/2
psd = ax.plot(f.freqs_THz, conv * f.psd,
#               label = "Bare PSD",
              alpha = alpha, zorder = 0, lw=0)
fpsd = ax.plot(f.freqs_THz, conv * f.fpsd,
#                label = "Window filter",
               color = psd[0].get_color(), zorder = 1)
ceps = ax.plot(f.freqs_THz, conv * f.dct.psd, lw = 2,
               label= "From Eq. (4)", #"Cepstral estimate",
               color = psd[0].get_color(), zorder = 2)


ax.set_xlim(0, 8)
ax.set_ylim(0, ymax)

ax.set_xlabel(r'Frequency (THz)')
ax.set_ylabel(r'Conductivity  (S/cm)')

ax.legend()

fig.savefig('Images/figure9.png')

fig, ax = plt.subplots()

TEMP = np.loadtxt('Data/current.K-KCl.core_and_wcs.dat', skiprows = 1).T[1].mean()

gap1 = np.loadtxt('Data/gap_HOMOLUMO.K33Cl31.dat')
gap2 = np.loadtxt('Data/gap_HOMO-1LUMO.K33Cl31.dat')

dt = TAUPS*15*20 
time = np.arange(len(gap1))*dt

lw = 0.5

gap1p = ax.plot(time, gap1, label = 'HOMO/LUMO',
                zorder = 0,
                lw = lw);
gap2p = ax.plot(time, gap1+gap2, label = 'HOMO-1/LUMO',
                zorder = 1,
                color='orange',
                lw = lw);

ax.set_xlabel('Time (ps)');
ax.set_ylabel(r'Energy (eV)');

kT = 8.617333262145e-5 * TEMP
ax.axhline(kT,
           lw = 3, color = 'red')

stoic_gap = 2.59 # eV
q1 = 2.01 #2.0619
q2 = 3.00 #2.9374

gapSp = ax.axhline(stoic_gap,
                   lw = 3, color = 'darkgreen', zorder=2)


ax.set_ylim(0,4)
ax.tick_params(axis='both', which='major')

fig.savefig('Images/figure4.png')