import readPhi
import readF
import readZlist
import matplotlib.pyplot as plt
import glob
import numpy as np
from scipy.constants import e, k

m_i = 1.9712e-25

#  paths = ['../quasiNeutral_fullAblation/','../Poisson_fullAblation/']
#  paths = ['../quasiNeutral_partialAblation/','../Poisson_partialAblation/','../PoissonTi_partialAblation/','../2024-11-05_13.56.23/']
#  paths = ['../2024-11-11_13.58.54/']
#  paths = ['../quasiNeutral_fullAblation/','../Poisson_fullAblation/', '../quasiNeutral_partialAblation/', '../Poisson_partialAblation/']
#  paths = ['../2024-12-02_21.07.52/', '../Poisson_50ns_T30Z11/']
#  paths = ['../2024-12-10_18.45.17/', '../Poisson_50ns_T30Z11/']
#  paths = ['../2024-12-11_12.38.27/', '../Poisson_polytropic_fa_T30Z11/', '../Poisson_fa_T30Z11/']
#  paths = ['../Poisson_partialAblation/','../Poisson_partialAblation_lowerT/','../Poisson_partialAblation_lowT/','../Poisson_partialAblation_highT/']
paths = ['../polytropic_80ns_T30/']
labels = [path[3:-1] for path in paths]

for path, label in zip(paths, labels):
    Zlist      = readZlist.read(path+'ZList.csv')
    filesCum_i = sorted(glob.glob(path+'time_*_fCum_i.csv'))
    _, _, v, _ = readF.read(filesCum_i[-1])
    sumF = np.zeros(len(v))

    for Z in Zlist:
        filename='time_*_Z_{:.1f}_fCum_i.csv'.format(Z)
        filesCum_i = sorted(glob.glob(path+filename))
        time, x, v, f_i = readF.read(filesCum_i[-1])
        sumF += f_i[0]
        plt.plot(v**2*m_i*0.5/e, f_i[0]*e/m_i/v, label=Z)
    print(time)
    plt.plot(v**2*m_i*0.5/e, sumF*e/m_i/v, label='sum', color='k', linestyle='dashed')

plt.yscale('log')
plt.ylim([1e16,5e27])
plt.ylabel('Sum f(e) / sqrt(e) (m^-3 eV^-1)')
plt.xscale('log')
plt.xlim([1e0,1e4])
plt.xlabel('e (eV)')

plt.legend()

#  plt.title('r = {:.1f} mm, time_max={:.1f} ns, '.format(x[0]*1e3, time*1e9) + label)

plt.show()