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

#  paths = ['../quasiNeutral_fullAblation/','../2024-09-26_11.48.04/']
#  paths = ['../2024-09-26_12.47.11/']
paths = ['../quasiNeutral_partialAblation/','../2024-09-26_13.58.24/']
#  path = '../quasiNeutral_fullAblation/'
#  path = '../quasiNeutral_partialAblatio/'


for path in paths:
    filesPhi   = sorted(glob.glob(path+'time_*_phi.csv'))
    filesMom_i = sorted(glob.glob(path+'time_*_mom_i.csv'))
    start = 0
    end   = len(filesMom_i)
    every = 100
    fig, ax = plt.subplots(4, sharex='all')
    for fileMom_i, filePhi in zip(filesMom_i[start:end+1:every], filesPhi[start:end+1:every]):
        time, r, phi, n_e            = readPhi.read(filePhi)
        time, r, n_i, u_i, T_i, Zave = readMom.read(fileMom_i)

        ax[0].plot(r, phi, label='t = {:.3f} ns'.format(time*1e9))
        ax[1].set_yscale('log')
        ax[1].set_ylim([1e14,2e25])
        ax[1].plot(r, Zave*n_i)
        ax[1].plot(r,      n_e, color='k', linestyle='dashed')
        ax[2].plot(r,      u_i)
        ax[3].plot(r,      T_i)

    ax[0].set_title(path)
    ax[0].legend()

    plt.show()