from nugridpy import nugridse as mp
import matplotlib.pyplot as plt
from collections import OrderedDict

# load data
umbe3m=mp.se('m2z1m3/H5_restart','restart.h5') 
#umbe3m=mp.se('m2z1m3/ne22pg/H5_restart','restart.h5')
#umbe3m=mp.se('m2z1m3/ne22ng/H5_restart','restart.h5')
#umbe3m=mp.se('m2z6m3/H5_restart','restart.h5')
#umbe3m=mp.se('m2z6m3/ne22pg/H5_restart','restart.h5')

# datasets
prot=[]
he4=[]
c12=[]
o16=[]
n14=[]
ne22=[]
na23=[]
mass=[]
i=[]

targets=[33800,35400]
h=0
for cycle in targets:
    prot1=umbe3m.se.get(cycle,'iso_massf','H-1')
    he41=umbe3m.se.get(cycle,'iso_massf','He-4')
    c121=umbe3m.se.get(cycle,'iso_massf','C-12')
    o161=umbe3m.se.get(cycle,'iso_massf','O-16')
    n141=umbe3m.se.get(cycle,'iso_massf','N-14')
    ne221=umbe3m.se.get(cycle,'iso_massf','Ne-22')
    na231=umbe3m.se.get(cycle,'iso_massf','Na-23')
    mass1=umbe3m.se.get(cycle,'mass')
    
    prot.append(prot1)
    he4.append(he41)
    c12.append(c121)
    o16.append(o161)
    n14.append(n141)
    ne22.append(ne221)
    na23.append(na231)
    mass.append(mass1)

    print(cycle)
    i.append(h)
    h=h+1

## uncomment to add extra tests and add 1 to umbe3m = umbe3m1
## also uncomment the plotting command
#for cycle in targets1:
#
#    na232=umbe3m1.se.get(cycle,'iso_massf','Na-23')
#    mass2=umbe3m1.se.get(cycle,'mass')
#    na_23.append(na232)
#    mass_1.append(mass2)
#
#    print(cycle)
#    i.append(t)
#    t=t+1

# plot datasets
for k in i:
    plt.figure(k)
    plt.semilogy(mass[k],prot[k],'k.',label='$H^1$',lw=4.5)
    plt.semilogy(mass[k],he4[k],'r-',label='$He^4$',lw=3)
    plt.semilogy(mass[k],c12[k],'g--',label='$C^{12}$',lw=3)
    plt.semilogy(mass[k],o16[k],'b-.',label='$O^{16}$',lw=3)
    plt.semilogy(mass[k],n14[k],'m-.',label='$N^{14}$',lw=3)
    plt.semilogy(mass[k],ne22[k],'r.',label='$Ne^{22}$',lw=3)
    plt.semilogy(mass[k],na23[k],'y-',label='$Na^{23}$',lw=3)
    #plt.semilogy(mass_1[k],na_23[k],'r-.',label='$Na^{23}$ $Ne22pg$',lw=3)
    
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = OrderedDict(zip(labels, handles))
    plt.legend(by_label.values(), by_label.keys(),prop={'size':25})
    plt.ylabel('Mass Fraction (X)', fontsize=40)
    plt.xlabel('Mass Coordinate (M$_\odot$)', fontsize=40)
    ax = plt.gca()
    ax.xaxis.get_majorticklocs()
    plt.rcParams['text.usetex'] = True
    plt.rcParams['axes.linewidth'] = 2
    plt.rcParams['xtick.major.size'] = 20
    plt.rcParams['xtick.major.width'] = 2
    plt.rcParams['ytick.major.size'] = 20
    plt.rcParams['ytick.major.width'] = 2
    plt.rcParams['ytick.minor.size'] = 10
    plt.rcParams['ytick.minor.width'] = 2
    plt.rcParams['xtick.labelsize'] = 30
    plt.rcParams['ytick.labelsize'] = 30
    plt.tick_params(axis='both',which='both',pad=10,direction='in')
    plt.axis([0.59, 0.63, 0.00001, 1])
plt.show()