import readPhi
import readMom
import readF
import readZlist
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/','../Poisson_partialAblation/']
#  paths = ['../2024-10-02_14.30.44/']
#  paths = ['../quasiNeutral_fullAblation/','../Poisson_fullAblation/']
#  paths = ['../2025-04-10_11.59.02/']
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')
    filesPhi   = sorted(glob.glob(path+'time_*_phi.csv'))
    start = 80
    end   = 85#len(filesPhi)
    every = 1
    fig, ax = plt.subplots(3, sharex='all')
    ax[1].set_yscale('log')
    ax[1].set_ylim(bottom=1e10, top=1e24)
    _, r, _, _, _         = readPhi.read(filesPhi[0])
    for t in range(start,end+1,every):
        sum_Zni = np.zeros(len(r))
        ave_ui  = np.zeros(len(r))
        ave_Ti  = np.zeros(len(r))
        for Z in Zlist:
            filename='time_*_Z_{:.1f}_mom_i.csv'.format(Z)
            filesMom_i = sorted(glob.glob(path+filename))
            fileMom_i = filesMom_i[t]
            time, r, n_i, u_i, T_i, Zave = readMom.read(fileMom_i)
            sum_Zni += Zave*n_i
            ave_ui  += Zave*n_i*u_i
            ave_Ti  += Zave*n_i*T_i

        ave_ui = np.divide(ave_ui, sum_Zni, out=np.zeros_like(ave_ui), where=sum_Zni!=0.0)
        ave_Ti = np.divide(ave_Ti, sum_Zni, out=np.zeros_like(ave_Ti), where=sum_Zni!=0.0)

        filePhi = filesPhi[t]
        time, r, phi, E, n_e  = readPhi.read(filePhi)

        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), label='t = {:.1f} ns'.format(time*1e9))
        ax[1].plot(r, sum_Zni)
        ax[1].plot(r,      n_e, color='k', linestyle='dashed')
        ax[2].plot(r, ave_ui)

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

    plt.show()