from astropy.table import Table
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import griddata

cmap = plt.get_cmap('hsv')
masses = ['0.6','0.8','1.0','1.2','1.5','2.0','2.5','3.0','4.0',
        '5.0','6.0','7.0','8.0','10','12','15','20','25','30','40','50']
z_arr = ['0.02']
markers = ['k-', 'g-']
for i,zval in enumerate(z_arr):
        for j,mval in enumerate(masses):
                path = '/data1/kirg/stellar_models/M%sZ%s/LOGS/history.data' % (mval, zval)
                tbl = Table.read(path, format='ascii', header_start=4, data_start=5)
                init_ind = np.arange(len(tbl))
                init_ind_sftd = np.where(((tbl['log_L']-tbl['log_LH']) > 0.0))
                tbl_prehburn = tbl[(((tbl['log_L']-tbl['log_LH']) < 0) & (tbl['center_h1']>1e-3*tbl['center_h1'][0]))]
                c_coeff = (j + 0.)/(len(masses) + 0.)
                plt.plot(tbl_prehburn['star_age'],tbl_prehburn['log_L'], c=cmap(c_coeff), label=mval)
                #plt.text(10**tbl_prehburn['log_Teff'][-1],tbl_prehburn['log_L'][-1], mval, horizontalalignment='right', verticalalignment='top')
plt.xscale('log')
plt.xlabel(r'Age, y')
plt.ylabel(r'$lg(L/L_{\odot})$')
plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=5, mode="expand", borderaxespad=0.)
plt.tight_layout()
plt.savefig('task2_L.pdf', dpi=150)
plt.clf()
for i,zval in enumerate(z_arr):
        for j,mval in enumerate(masses):
                path = '/data1/kirg/stellar_models/M%sZ%s/LOGS/history.data' % (mval, zval)
                tbl = Table.read(path, format='ascii', header_start=4, data_start=5)
                init_ind = np.arange(len(tbl))
                init_ind_sftd = np.where(((tbl['log_L']-tbl['log_LH']) > 0.0))
                tbl_prehburn = tbl[(((tbl['log_L']-tbl['log_LH']) < 0) & (tbl['center_h1']>1e-3*tbl['center_h1'][0]))]
                c_coeff = (j + 0.)/(len(masses) + 0.)
                plt.plot(tbl_prehburn['star_age'],tbl_prehburn['log_R'], c=cmap(c_coeff), label=mval)
plt.xscale('log')
plt.xlabel(r'Age, y')
plt.ylabel(r'$lg(R/R_{\odot})$')
plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=5, mode="expand", borderaxespad=0.)
#plt.tight_layout()
plt.savefig('task2_R.pdf', inches='tight',dpi=150)
plt.clf()
for i,zval in enumerate(z_arr):
        for j,mval in enumerate(masses):
                path = '/data1/kirg/stellar_models/M%sZ%s/LOGS/history.data' % (mval, zval)
                tbl = Table.read(path, format='ascii', header_start=4, data_start=5)
                init_ind = np.arange(len(tbl))
                init_ind_sftd = np.where(((tbl['log_L']-tbl['log_LH']) > 0.0))
                tbl_prehburn = tbl[(((tbl['log_L']-tbl['log_LH']) < 0) & (tbl['center_h1']>1e-3*tbl['center_h1'][0]))]
                c_coeff = (j + 0.)/(len(masses) + 0.)
                plt.plot(tbl_prehburn['star_age'],tbl_prehburn['log_center_T'], c=cmap(c_coeff), label=mval)
plt.xscale('log')
plt.xlabel(r'Age, y')
plt.ylabel(r'$lg T$')
plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=5, mode="expand", borderaxespad=0.)
#plt.tight_layout()
plt.savefig('task2_T.pdf', inches='tight',dpi=150)

plt.clf()
for i,zval in enumerate(z_arr):
        for j,mval in enumerate(masses):
                path = '/data1/kirg/stellar_models/M%sZ%s/LOGS/history.data' % (mval, zval)
                tbl = Table.read(path, format='ascii', header_start=4, data_start=5)
                init_ind = np.arange(len(tbl))
                init_ind_sftd = np.where(((tbl['log_L']-tbl['log_LH']) > 0.0))
                tbl_prehburn = tbl[(((tbl['log_L']-tbl['log_LH']) < 0) & (tbl['center_h1']>1e-3*tbl['center_h1'][0]))]
                c_coeff = (j + 0.)/(len(masses) + 0.)
                plt.plot(tbl_prehburn['star_age'],tbl_prehburn['log_center_Rho'], c=cmap(c_coeff), label=mval)
plt.xscale('log')
plt.xlabel(r'Age, y')
plt.ylabel(r'$lg \rho$')
plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=5, mode="expand", borderaxespad=0.)
#plt.tight_layout()
plt.savefig('task2_rho.pdf', inches='tight',dpi=150)