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/']
#  paths = ['../2024-10-02_14.30.44/']
#  paths = ['../quasiNeutral_fullAblation/','../Poisson_fullAblation/']
paths = ['../2024-10-03_19.28.22/']
labels = [path[3:-1] for path in paths]

for path, label in zip(paths, labels):
    filesPhi   = sorted(glob.glob(path+'time_*_phi.csv'))
    filesMom_i = sorted(glob.glob(path+'time_*_mom_i.csv'))
    start = 100
    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, E, 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 = {:.1f} ns'.format(time*1e9))
        #  ax[0].plot(r, E, label='t = {:.1f} ns'.format(time*1e9))
        #  ax[0].plot(r, (Zave*n_i - n_e)/n_e, label='t = {:.1f} ns'.format(time*1e9))
        ax[1].set_yscale('log')
        ax[1].set_ylim([1e20,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(label)
    ax[0].legend()

    plt.show()