import mppnp as mp
import utils
import mesa as ms


'''
m1=mp.se('/rpod2/athira/mppnp/case1/H5_out')

m2=mp.se('/rpod2/athira/mppnp/case11/H5_out')


matplotlib.rc('xtick',labelsize=15,color='k')
matplotlib.rc('ytick',labelsize=12)
matplotlib.rc('ytick.major',pad=15)
matplotlib.rc('axes',labelsize=12,labelcolor='k')


lines1=['r--','r:','r-.','r.-']
lines2=['b--','b:','b-.','b.-']

ind=800
neut_den_1=[]
neut=m1.get(ind,'Neutron-1')
mass1=m1.get(ind,'mass')
rho= m1.get(ind,'rho')
age1=m1.get(ind,'age')
for a in range(len(neut)):
    res=(neut[a]*6*10**23*rho[a])
    neut_den_1.append(res)
c13=m1.get(ind,'C-13')
ti46=m1.get(ind,'Ti-46');ti47=m1.get(ind,'Ti-47');ti48=m1.get(ind,'Ti-48');ti49=m1.get(ind,'Ti-49');ti50=m1.get(ind,'Ti-50')
ti1=ti46/46.0+ti47/47.0+ti48/48.0+ti49/49.0+ti50/50.0
y89=m1.get(ind,'Y-89')
la138=m1.get(ind,'La-138');la139=m1.get(ind,'La-139')
ba130=m1.get(ind,'Ba-130');ba132=m1.get(ind,'Ba-132');ba134=m1.get(ind,'Ba-134');ba135=m1.get(ind,'Ba-135');ba136=m1.get(ind,'Ba-136');ba137=m1.get(ind,'Ba-137');ba138=m1.get(ind,'Ba-138')

y1=y89/89.0
ba1=ba130/130.0+ba132/132.0+ba134/134.0+ba135/135.0+ba136/136.0+ba137/137.0+ba138/138.0
la1=la138/138.0+la139/139.0
c13_1=(10**-2*c13)/13.

ind=550
neut_den_2=[]
neut=m2.get(ind,'Neutron-1')
mass2=m2.get(ind,'mass')
rho= m2.get(ind,'rho')
age2=m2.get(ind,'age')
for a in range(len(neut)):
    res=(neut[a]*6*10**23*rho[a])
    neut_den_2.append(res)
c13=m2.get(ind,'C-13')
ti46=m2.get(ind,'Ti-46');ti47=m2.get(ind,'Ti-47');ti48=m2.get(ind,'Ti-48');ti49=m2.get(ind,'Ti-49');ti50=m2.get(ind,'Ti-50')
ti2=ti46/46.0+ti47/47.0+ti48/48.0+ti49/49.0+ti50/50.0
y89=m2.get(ind,'Y-89')
la138=m2.get(ind,'La-138');la139=m2.get(ind,'La-139')
ba130=m2.get(ind,'Ba-130');ba132=m2.get(ind,'Ba-132');ba134=m2.get(ind,'Ba-134');ba135=m2.get(ind,'Ba-135');ba136=m2.get(ind,'Ba-136');ba137=m2.get(ind,'Ba-137');ba138=m2.get(ind,'Ba-138')
y2=y89/89.0
la2=la138/138.0+la139/139.0
ba2=ba130/130.0+ba132/132.0+ba134/134.0+ba135/135.0+ba136/136.0+ba137/137.0+ba138/138.0
c13_2=(10**-2*c13)/13.

specs1=[y1,ba1,ti1,c13_1]
specs_labels1=['Y, case 1','Ba, case 1','Ti, case 1','$^{13}C$, case 1']
specs2=[y2,ba2,ti2,c13_2]
specs_labels2=['Y, case 4','Ba, case 4','Ti, case 4','$^{13}C$, case 1']
'''
close(40)
ifig=figure(40,figsize=(11,6))
ax1=ifig.add_subplot(111)
for a in range (len(specs2)):
        ax1.semilogy(mass1,specs1[a],lines1[a],linewidth=2,label=specs_labels1[a],markersize=6.2,alpha=0.65,markevery=15)    
        ax1.semilogy(mass2,specs2[a],lines2[a],linewidth=2,label=specs_labels2[a],markersize=6.2,alpha=0.65,markevery=15)       
ylim(1e-10,1e-5)
box1 = ax1.get_position()
ax1.set_position([box1.x0, box1.y0, box1.width * 0.90, box1.height])
ax1.legend(loc='center', bbox_to_anchor=(1.15, 0.10),prop={'size':10.0})
ax1.set_ylabel('Number fraction')
ax1.set_xlabel('$Mr/M_\mathrm{\odot}$')
ax2=ax1.twinx()
ax2.plot(mass1,log10(neut_den_1),color='violet',linestyle='-',linewidth=1,label='log Nn, case 1')
ax2.plot(mass2,log10(neut_den_2),color='c',linestyle='-',linewidth=1,label='log Nn, case 4')
ax2.legend(loc='center', bbox_to_anchor=(1.15, 0.85),prop={'size':10.0})
ax2.set_ylabel('Log Neutron Density ($N_\mathrm{n}$) ($\mathrm{cm}^{-3}$)')
ylim(0,7)
xlim(0.44,0.90)
ax2.xaxis.set_major_locator(MultipleLocator(0.05))

savefig('post-merger_nden.png')
savefig('post-merger_nden.pdf')
savefig('post-merger_nden.eps')


'''
    ifig1=figure(40,figsize=(10,7))
    ax1=ifig1.add_subplot(111)
    for a in range (len(specs1a)):
        #ax1.plot(mass,specs1a[a],lines1a[i], label=specs1a_l[a],linewidth=2,markersize=6.2,alpha=0.65,markevery=15)
        ax1.semilogy(mass,specs1a[a],lines0[i],linewidth=2,markersize=6.2,alpha=0.65,markevery=15)
        i+=1
    #ylim(1e-5,max(c13))
    xlim(0.45,0.75)
    box1 = ax1.get_position()
    #ax1.set_position([box1.x0, box1.y0, box1.width * 0.83, box1.height])
    #ax1.legend(loc='center', bbox_to_anchor=(1.17, 0.17),prop={'size':10.0})
    ax1.set_ylabel('$^{13}C$ Mass fraction')
    ax1.set_xlabel('$Mr/M_{\odot}$')
    ax2=ax1.twinx()
    ax2.plot(mass,log10(neut_den_1),lines_n[d],linewidth=1.2)
    #ylim(max(log10(neut_den_1))/3.0,max(log10(neut_den_1)+1))
    d+=1
    ylim(0,7)
    xlim(0.45,0.75)   
    #savefig('mppnp_profile1.2a_'+'0'+str(ind)+'.png')
    #savefig('mppnp_profile1.2a_'+'0'+str(ind)+'.pdf')
    #savefig('mppnp_profile1.2a_'+'0'+str(ind)+'.eps')
    #close(10)
    #close(20)
ax2.set_ylabel('Log Neutron Density (/cc)')


    ifig2=figure(20,figsize=(10,7))
    ax2=ifig2.add_subplot(111)
    for a in range (len(specs1b)):
        ax2.semilogy(mass,specs1b[a],lines1b[j], label=specs1b_l[a],linewidth=2,markersize=6.2,alpha=0.65,markevery=5)
        j+=1
    box2 = ax2.get_position()
    ax2.set_position([box2.x0, box2.y0, box2.width * 0.83, box2.height])
    ax2.legend(loc='center', bbox_to_anchor=(1.17, 0.17),prop={'size':10.0})
    ax2.set_ylabel('Mass fraction')
    ax2.set_xlabel('$Mr/M_{\odot}$')
    ax3=ax2.twinx()
    ax3.plot(mass,log10(neut_den_1),'c-',linewidth=1.2)
    ylim(max(log10(neut_den_1))/3.0,max(log10(neut_den_1)+1))
    xlim(0.46,0.58)   
'''
'''
specs0_l=['He-4']
specs1a_l=['C-12','N-14']
specs1b_l=['H-1','C-14','N-15']
specs1c_l=['O-15','F-19']
#specs1a_l=['H-1','C-12','C-14','N-13','N-14','N-15','O-15','F-18','Ne-19','F-19']
#specs1b_l=['He-4','C-13','Ne-22','Mg-25']

#specs1b_l=['N-15','O-16','O-18','F-19','Ne-22']
#specs2_l=['Na-23','Mg-24','Al-27','Si-28','S-32','Ca-40','Ti-48']
#specs3_l=['Fe-56','Ni-58','Zn-64','Y-89','Zr-90','Ba-138','La-139']
#lines=['k-.','ro-','ks-','r--','k+:','rd-','k.-','rh-','k>-','r:']
lines0=['k:','r:','b:']
lines_n=['k-','r-','b-']
lines1a=['k-.','ko-','r-.','ro-.','b-.','bo-.']
lines1b=['ks-','k--','k+:','rs-','r--','r+:','bs-','b--','b+:']
lines1c=['k.-','kd-.','r.-','rd-.','b.-','bd-.']
index_1=[300,500,700]

i=0;j=0;k=0
d=0
close(0)
close(10)
close(20)
close(30)



for ind in index_1:
    neut_den_1=[]
    neut=m.get(ind,'Neutron-1')
    mass=m.get(ind,'mass')
    rho= m.get(ind,'rho')
    for a in range(len(neut)):
        res=(neut[a]*6*10**23*rho[a])
        neut_den_1.append(res)
    h1=m.get(ind,'H-1');he4=m.get(ind,'He-4');c12=m.get(ind,'C-12');c13=m.get(ind,'C-13');c14=m.get(ind,'C-14');n13=m.get(ind,'N-13');n14=m.get(ind,'N-14');o15=m.get(ind,'O-15');
    n15=m.get(ind,'N-15');o16=m.get(ind,'O-16');o18=m.get(ind,'O-18');f18=m.get(ind,'F-18');f19=m.get(ind,'F-19');ne19=m.get(ind,'Ne-19');ne22=m.get(ind,'Ne-22')
    na23=m.get(ind,'Na-23');mg24=m.get(ind,'Mg-24');mg25=m.get(ind,'Mg-25');al27=m.get(ind,'Al-27');si28=m.get(ind,'Si-28');s32=m.get(ind,'S-32');ca40=m.get(ind,'Ca-40');ti48=m.get(ind,'Ti-48')
    fe56=m.get(ind,'Fe-56');ni58=m.get(ind,'Ni-58');zn64=m.get(ind,'Zn-64');y89=m.get(ind,'Y-89');zr94=m.get(ind,'Zr-94');ba138=m.get(ind,'Ba-138');la139=m.get(ind,'La-139')
    specs0=[he4]
    specs1a=[c12,n14]
    specs1b=[h1,c14,n15]
    specs1c=[o15,f19]

    ifig0=figure(0,figsize=(10,7))
    ax1=ifig0.add_subplot(111)
    for a in range (len(specs0)):
        ax1.plot(mass,specs0[a],lines0[d], label=specs0_l[a],linewidth=2,markersize=6.2,alpha=0.65,markevery=15)
    ylim(0.3,1.0)
    if ind==300:
        box1 = ax1.get_position()
        ax1.set_position([box1.x0, box1.y0, box1.width * 0.83, box1.height])
        ax1.legend(loc='center', bbox_to_anchor=(1.17, 0.17),prop={'size':10.0})
        ax1.set_ylabel('Mass fraction')
        ax1.set_xlabel('$Mr/M_{\odot}$')
    #ax2=ax1.twinx()
    #ax2.plot(mass,log10(neut_den_1),lines_n[d],linewidth=1.2)
    #ylim(max(log10(neut_den_1))/3.0,max(log10(neut_den_1)+1))
    d+=1

    ifig1=figure(10,figsize=(10,7))
    ax1=ifig1.add_subplot(111)
    for a in range (len(specs1a)):
        ax1.plot(mass,specs1a[a],lines1a[i], label=specs1a_l[a],linewidth=2,markersize=6.2,alpha=0.65,markevery=15)
        i+=1
    #ylim(1e-8,1.1)
    xlim(0.46,0.58)
    if ind==300:
        box1 = ax1.get_position()
        ax1.set_position([box1.x0, box1.y0, box1.width * 0.83, box1.height])
        ax1.legend(loc='center', bbox_to_anchor=(1.17, 0.17),prop={'size':10.0})
        ax1.set_ylabel('Mass fraction')
        ax1.set_xlabel('$Mr/M_{\odot}$')
    xlim(0.46,0.58)   
    #savefig('mppnp_profile1.2a_'+'0'+str(ind)+'.png')
    #savefig('mppnp_profile1.2a_'+'0'+str(ind)+'.pdf')
    #savefig('mppnp_profile1.2a_'+'0'+str(ind)+'.eps')
    #close(10)
    #close(20)
    ifig2=figure(20,figsize=(10,7))
    ax2=ifig2.add_subplot(111)
    for a in range (len(specs1b)):
        ax2.semilogy(mass,specs1b[a],lines1b[j], label=specs1b_l[a],linewidth=2,markersize=6.2,alpha=0.65,markevery=5)
        j+=1
    if ind==300:
        box2 = ax2.get_position()
        ax2.set_position([box2.x0, box2.y0, box2.width * 0.83, box2.height])
        ax2.legend(loc='center', bbox_to_anchor=(1.17, 0.17),prop={'size':10.0})
        ax2.set_ylabel('Mass fraction')
        ax2.set_xlabel('$Mr/M_{\odot}$')

    ifig3=figure(30,figsize=(10,7))
    ax3=ifig3.add_subplot(111)
    for a in range (len(specs1c)):
        ax3.semilogy(mass,specs1c[a],lines1c[k], label=specs1c_l[a],linewidth=2,markersize=6.2,alpha=0.65,markevery=5)
        k+=1
    if ind==300:
        box3 = ax3.get_position()
        ax3.set_position([box3.x0, box3.y0, box3.width * 0.83, box3.height])
        ax3.legend(loc='center', bbox_to_anchor=(1.17, 0.17),prop={'size':10.0})
        ax3.set_ylabel('Mass fraction')
        ax3.set_xlabel('$Mr/M_{\odot}$')
    #ylim(1e-8,1.1)
    xlim(0.46,0.58)   
    #ax2=ax1.twinx()
    #ax2.plot(mass,log10(neut_den_1),linewidth=1.2,color='c')
    #ylim(max(log10(neut_den_1))/3.0,max(log10(neut_den_1)+1))
    xlim(0.46,0.58)
    #savefig('mppnp_profile1.2b_'+'0'+str(ind)+'.png')
    #savefig('mppnp_profile1.2b_'+'0'+str(ind)+'.pdf')
    #savefig('mppnp_profile1.2b_'+'0'+str(ind)+'.eps')
    #close(ind+15)

'''