import mesa as ms
import utils
import matplotlib.patches as mpatches

matplotlib.rc('xtick',labelsize=15,color='k')                                   
matplotlib.rc('ytick',labelsize=15)                                             
matplotlib.rc('ytick.major',pad=5)                                              
matplotlib.rc('axes',labelsize=17,labelcolor='k')       
'''
lines=utils.symbol_list('lines2')
s1=ms.star_log('/nfs/rpod2/athira/mesa/case1/LOGS')
teff1=s1.get('log_Teff')
l1=s1.get('log_L')
age1=s1.get('star_age')
s2=ms.star_log('/nfs/rpod2/athira/mesa/case19/LOGS')
teff2=s2.get('log_Teff')
l2=s2.get('log_L')
age2=s2.get('star_age')

s3=ms.star_log('/nfs/rpod2/athira/mesa/case9/LOGS')
teff3=s3.get('log_Teff')
l3=s3.get('log_L')
age3=s3.get('star_age')
cycles3=s3.get('model_number')
s4=ms.star_log('/nfs/rpod2/athira/mesa/case11/LOGS')
teff4=s4.get('log_Teff')
l4=s4.get('log_L')
age4=s4.get('star_age')
m_dot1=s1.get('log_abs_mdot')
m_dot2=s2.get('log_abs_mdot')
m_dot3=s3.get('log_abs_mdot')
m_dot4=s4.get('log_abs_mdot')
'''

logs=[s1,s2,s3,s4]
teff_specs=[teff1,teff2,teff3,teff4]
l_specs=[l1,l2,l3,l4]
cycles1=s1.get('model_number');cycles2=s2.get('model_number');cycles4=s4.get('model_number')

#legends=['q=0.7','q=0.5']
#labels=['log Teff', 'log L']


hrd_dash1=[];hrd_dash2=[];hrd_dash3=[];hrd_dash4=[]

rcb_mod_all1=[]

teff_rcb1=[]


for a in range(len(cycles1)):
        if l1[a]>=3.5 and 3.7<=teff1[a]<=3.903:
                teff_rcb1.append(teff1[a])
                rcb_mod_all1.append(int(cycles1[a]))


cycle_min1=array(rcb_mod_all1)[list(teff_rcb1).index(min(teff_rcb1))]


mod_1_1=array(rcb_mod_all1)[np.where(array(rcb_mod_all1)<cycle_min1)]
teff_1_1=array(teff_rcb1)[np.where(array(rcb_mod_all1)<cycle_min1)]

mod_1_2=array(rcb_mod_all1)[np.where(array(rcb_mod_all1)>=cycle_min1)]
teff_1_2=array(teff_rcb1)[np.where(array(rcb_mod_all1)>=cycle_min1)]

rcb_begin1=min(rcb_mod_all1)

rcb_end1=max(rcb_mod_all1)

age_begin1=age1[rcb_begin1]
age_end1=age1[rcb_end1]
rcb_age1=[age_end1-age_begin1]

rcb_mod_all2=[]

teff_rcb2=[]

for a in range(len(cycles2)):
        if l2[a]>=3.5 and 3.7<=teff2[a]<=3.903:
                teff_rcb2.append(teff2[a])
                rcb_mod_all2.append(int(cycles2[a]))

cycle_min2=array(rcb_mod_all2)[list(teff_rcb2).index(min(teff_rcb2))]


mod_2_2=array(rcb_mod_all2)[np.where(array(rcb_mod_all2)<cycle_min2)]
teff_2_2=array(teff_rcb2)[np.where(array(rcb_mod_all2)<cycle_min2)]

mod_2_2=array(rcb_mod_all2)[np.where(array(rcb_mod_all2)>=cycle_min2)]
teff_2_2=array(teff_rcb2)[np.where(array(rcb_mod_all2)>=cycle_min2)]

rcb_begin2=min(rcb_mod_all2)

rcb_end2=max(rcb_mod_all2)


age_begin2=age2[rcb_begin2]
age_end2=age2[rcb_end2]
rcb_age2=[age_end2-age_begin2]


rcb_mod_all3=[]

teff_rcb3=[]
cycles3=s3.get('model_number')
for a in range(len(cycles3)):
        if l3[a]>=3.5 and 3.7<=teff3[a]<=3.903:
                teff_rcb3.append(teff3[a])
                rcb_mod_all3.append(int(cycles3[a]))


cycle_min3=array(rcb_mod_all3)[list(teff_rcb3).index(min(teff_rcb3))]


mod_3_3=array(rcb_mod_all3)[np.where(array(rcb_mod_all3)<cycle_min3)]
teff_3_3=array(teff_rcb3)[np.where(array(rcb_mod_all3)<cycle_min3)]

mod_3_2=array(rcb_mod_all3)[np.where(array(rcb_mod_all3)>=cycle_min3)]
teff_3_2=array(teff_rcb3)[np.where(array(rcb_mod_all3)>=cycle_min3)]

rcb_begin3=min(rcb_mod_all3)

rcb_end3=max(rcb_mod_all3)


age_begin3=age3[rcb_begin3]
age_end3=age3[rcb_end3]
rcb_age3=[age_end3-age_begin3]


rcb_mod_all4=[]

teff_rcb4=[]

for a in range(len(cycles4)):
        if l4[a]>=3.5 and 3.7<=teff4[a]<=3.903:
                teff_rcb4.append(teff4[a])
                rcb_mod_all4.append(int(cycles4[a]))


cycle_min4=array(rcb_mod_all4)[list(teff_rcb4).index(min(teff_rcb4))]


mod_4_4=array(rcb_mod_all4)[np.where(array(rcb_mod_all4)<cycle_min4)]
teff_4_4=array(teff_rcb4)[np.where(array(rcb_mod_all4)<cycle_min4)]

mod_4_2=array(rcb_mod_all4)[np.where(array(rcb_mod_all4)>=cycle_min4)]
teff_4_2=array(teff_rcb4)[np.where(array(rcb_mod_all4)>=cycle_min4)]

rcb_begin4=min(rcb_mod_all4)

rcb_end4=max(rcb_mod_all4)


age_begin4=age4[rcb_begin4]
age_end4=age4[rcb_end4]
rcb_age4=[age_end4-age_begin4]


#s1.xlimrev()    

close(100)
figure(100)

xlim(5.0,3.5)
ylim(2.2,4.1)


ifig=figure(100,figsize=(8,6))

subplot(1,1,1)

ax1=ifig.add_subplot(1,1,1)
box = ax1.get_position()
ax1.set_position([box.x0, box.y0, box.width * 0.90, box.height]) 


xy=3.9,3.6
width=-0.3
height=0.4
p1 = mpatches.Rectangle(xy, width, height, facecolor="black", edgecolor="green", alpha=0.15)
draw()
gca().add_patch(p1)
draw()

text(x=3.8,y=3.5,s='RCB',fontsize=17,color='k')

for ind,item in enumerate(cycles1):
    if item <= 1300:
        hrd_dash1.append(ind)

ax1.plot(teff1[hrd_dash1[-1]:],l1[hrd_dash1[-1]:],'orange',linestyle='-',alpha=1.0,linewidth=3,label='case 1')
ax1.plot(teff1[hrd_dash1],l1[hrd_dash1],'orange',linestyle='-',alpha=1.0,linewidth=1)
#ax1.plot(teff1[hrd_dash1],l1[hrd_dash1],'k-',alpha=0.3,linewidth=2)
ax1.plot(4.1,2.62,'k',marker='o',markersize=7,mfc='None')
ax1.plot(teff1[650],l1[650],'cyan',marker='o',markersize=7,alpha=0.8)
ax1.plot(teff1[1150],l1[1150],'violet',marker='o',markersize=7)
ax1.plot(teff1[1750],l1[1750],'brown',marker='o',markersize=7,alpha=0.8)
ax1.plot(teff1[2500],l1[2500],'k',marker='o',markersize=7,alpha=0.8)


for ind,item in enumerate(cycles2):
    if item <= 1300:
        hrd_dash2.append(ind)

ax1.plot(teff2[hrd_dash2[-1]:],l2[hrd_dash2[-1]:],color='r',linestyle=':',alpha=1.0,linewidth=3,label='case 2')
ax1.plot(teff2[hrd_dash2],l2[hrd_dash2],color='r',linestyle=':',linewidth=1)

for ind,item in enumerate(cycles3):
    if item <= 1300:
        hrd_dash3.append(ind)

ax1.plot(teff3[hrd_dash3[-1]:],l3[hrd_dash3[-1]:],color='b',linestyle='-.',alpha=1.0,markevery=20,linewidth=2,label='case 3')
ax1.plot(teff3[hrd_dash3],l3[hrd_dash3],color='b',linestyle='-.',alpha=1.0,linewidth=0.7,markevery=20)

for ind,item in enumerate(cycles4):
    if item <= 1300:
        hrd_dash4.append(ind)


ax1.plot(teff4[hrd_dash4[-1]:],l4[hrd_dash4[-1]:],'g--',alpha=1.0,linewidth=2,label='case 4')
ax1.plot(teff4[hrd_dash4],l4[hrd_dash4],'g--',alpha=1,linewidth=0.7)


ax1.legend(loc='center', bbox_to_anchor=(1.14, 0.08),prop={'size':11.0})
ax1.set_ylabel('log $ L/L_\mathrm{\odot}$')
ax1.set_xlabel('log $ T_\mathrm{eff} (K)$')
xlim(5.0,3.5)
ylim(2.6,4.01)
'''
for num,spec in enumerate(specs):
    i=0
    close(num);figure(num)
    for ind,log in enumerate(logs):
        age=[]
        age = log.get('star_age')-log.get('star_age')[0]
        plot(age,log.get(spec),lines[i],label=legends[ind])
        xlabel('age(years)')
        ylabel(labels[num])
        i+=1

legend(loc=2)

rcb_teff=[]
rcb_model=[]
rcb_age=[]
rcb_teffu=[]
rcb_modelu=[]
rcb_ageu=[]


i=0
close(100)
for ind,log in enumerate(logs):
    age=[]
    teff=[]
    model=[]
    figure(100)
    print i
    log.plot(specs[0],specs[1],shape=lines[i],legend=legends[ind],labelx=labels[0],labely=labels[1],title='')
    i+=1
    age=log.get('star_age')
    teff=log.get(specs[0])
    model=log.get('model_number')
    diff1=[]
    for a,b in enumerate(teff):
        res1=abs(b-3.9030)
        diff1.append(res1)
    rcb_teff.append(teff[diff1.index(min(diff1))])
    rcb_age.append(age[diff1.index(min(diff1))])
    rcb_model.append(model[diff1.index(min(diff1))])

log.xlimrev()    
legend(loc=4)
xlim(5.0,3.5)
ylim(2.2,4.1)
'''
'''
xy=3.9,3.6
width=-0.3
height=0.4
p1 = mpatches.Rectangle(xy, width, height, facecolor="k", edgecolor="k",alpha=0.3)
plt.draw()
plt.gca().add_patch(p1)
plt.draw()
'''
text(x=3.8,y=3.5,s='RCB',fontsize=17,color='k')
#savefig('HRD_runs.png');savefig('HRD_runs.eps');savefig('HRD_runs.pdf')