#import asciidata as asc
import mppnp as mp  
import numpy as np
import matplotlib.pyplot as plt
import time as tm
import mesa as ms
import scipy 
from scipy import interpolate
from scipy import integrate

#Least squares fit to obtain envelope masses

mhe_data=np.array([0.179, 0.195, 0.234, 0.259, 0.300, 0.331, 0.414])

mc_data=np.array([ 0.1552,  0.1782,  0.2220,  0.2524,  0.2960  ,  0.3281,  0.4116, 0.546, 0.565, 0.598, 0.644])

menv_data=np.array([ 0.048289,  0.030768,  0.013098,  0.007232,  0.004746,  0.003744, 0.002175, 5.01e-4, 2.51e-4, 1.905e-4, 1.23e-4])

A = np.vstack([mc_data, np.ones(len(mc_data))]).T

slope,coeff = np.linalg.lstsq(A,log10(menv_data))[0]

mc_new=np.arange(0.1552,0.6441,0.0001)

m_env=10**(slope*mc_new+coeff)

#Abundance profile plot of AGB model

m=ms.star_log('/nfs/rpod2/fherwig/SEE/data/set1.2/see_wind/M3.00Z2.0e-02/LOGS')
h1=m.get('h1_boundary_mass')
model=m.get('model_number')
rad=m.get('log_R')

agb=2364
import mppnp as mp
p=mp.se('/rpod2/fherwig/SEE/data/set1.2/ppn_wind/M3.00Z2.0e-02.standard/H5_out')
Y=p.get(agb,'iso_massf')#This gives all the mass fractions of  all the isotopes for all th models
isos=p.se.isotopes#And this gives all the isotopes
mass=p.get(agb,'mass')
H_agb=p.get(agb,'H-1')
He_agb=p.get(agb,'He-4')
C_agb=p.get(agb,'C-12')
N_agb=p.get(agb,'N-14')
O16_agb=p.get(agb,'O-16')
O18_agb=p.get(agb,'O-18')
Ne_agb=p.get(agb,'Ne-22')

diff1=[]
for a in range(len(H_agb)):
	res1=abs(H_agb[a]-0.37)
        diff1.append(res1)

H_boun_ind=diff1.index(min(diff1))
H_boun=H_agb[H_boun_ind]
m_boun=mass[H_boun_ind]

mco_env=10**(slope*m_boun+coeff)
mco=m_boun+mco_env

diff2=[]
for a in range(len(mass)):
	res1=abs(mass[a]-mco)
        diff2.append(res1)

mass_co_ind=diff2.index(min(diff2))
mass_co=mass[mass_co_ind]
H_co=H_agb[mass_co_ind]

mass_H=integrate.cumtrapz(H_agb,x=mass)
mass_H_env=mass_H[mass_co_ind]-mass_H[m_boun]

print mass_co, mass_H_env, H_co 


ind=np.where(mass<mass_co)
Y_wd=Y[ind]
mass_wd=mass[ind]
ind1=ind

mass_plot=1-(mass_wd/mass_co)

fig=figure()
ax1 = fig.add_subplot(111)
'''
#Plot: 1-(mass_wd/mass_co)
ax1.loglog(mass_plot,H_agb[:len(mass_plot)],label='$^1H$',color='orange')
ax1.loglog(mass_plot,He_agb[:len(mass_plot)],label='$^4He$')
ax1.loglog(mass_plot,C_agb[:len(mass_plot)],label='$^{12}C$')
ax1.loglog(mass_plot,N_agb[:len(mass_plot)],label='$^{14}N$')
ax1.loglog(mass_plot,O16_agb[:len(mass_plot)],label='$^{16}O$')
ax1.loglog(mass_plot,O18_agb[:len(mass_plot)],label='$^{18}O$')
ax1.loglog(mass_plot,Ne_agb[:len(mass_plot)],label='$^{22}Ne$',color='blue')

xlim(3e-5,1)
ylim(1e-5,1.1)
ax1.set_xlabel('$1-M_{r}/M_{CO}$')
ax1.set_ylabel('Mass fraction')

box = ax1.get_position()
ax1.set_position([box.x0, box.y0, box.width * 0.88, box.height])
ax1.legend(loc='center left', bbox_to_anchor=(1.05, 0.20),prop={'size':12})
show()
'''

#Plot :mass_wd/mass_co
#figure(10)
ax1.semilogy(mass_wd,H_agb[:len(mass_wd)],label='$^1H$',color='orange')
ax1.semilogy(mass_wd,He_agb[:len(mass_wd)],label='$^4He$')
ax1.semilogy(mass_wd,C_agb[:len(mass_wd)],label='$^{12}C$')
ax1.semilogy(mass_wd,N_agb[:len(mass_wd)],label='$^{14}N$')
ax1.semilogy(mass_wd,O16_agb[:len(mass_wd)],label='$^{16}O$')
ax1.semilogy(mass_wd,O18_agb[:len(mass_wd)],label='$^{18}O$')
ax1.semilogy(mass_wd,Ne_agb[:len(mass_wd)],label='$^{22}Ne$',color='blue')

xlim(0.0,0.90)
ylim(1e-5,1.1)
ax1.set_xlabel('$M_{r}/M_{CO}$')
ax1.set_ylabel('Mass fraction')

box = ax1.get_position()
ax1.set_position([box.x0, box.y0, box.width * 0.88, box.height])
ax1.legend(loc='center left', bbox_to_anchor=(1.05, 0.20),prop={'size':12})
show()


#Legends and labels
'''
axvline(x=0.2,color='k',linestyle='--')
axvline(x=0.1341,color='k',linestyle='--')
axvline(x=0.0661,color='k',linestyle='--')
axvline(x=0.00042,color='k',linestyle='--')

text(0.1615,0.032,"Core (Hydro)",rotation=90,color='k',fontsize=13)
text(0.0943,0.0017,"He-shell",rotation=90,color='k',fontsize=13)
text(0.0075,0.070,"H-shell ashes",rotation=90,color='k',fontsize=13)
text(0.0001,0.0017,"Envelope",rotation=90,color='k',fontsize=13)


'''