import utils
import mesa as ms
import utils

#mods=[1,700,1150,2500]
mods=[2500]
logdir='LOGS'
p=ms.mesa_profile(logdir,100,num_type='model')
cycles=p.log_ind.keys()
cycles.sort()
#mods=(1,cycles[-1],50)
lines=utils.symbol_list('lines2')
s=ms.star_log(logdir)

specs=['n14','log_D_mix','log_mlt_D_mix','logT','eps_nuc','gradr','grada','logS']
labels=['$X_{^{14}N}$','log $D_\mathrm{total}$','log $D_\mathrm{MLT}$', 'log $T$ ','${\epsilon}_\mathrm{nuc}$','{\\nabla}_{r}','{\\nabla}_{a}','log $S$']

specs2=['log_mlt_D_mix','cp','log_opacity','logRho']

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


for mod in mods:
    mass=[]
    radius=[]
    neut=[]
    h1=[]
    he4=[]
    f19=[]
    c12=[]
    n14=[]
    c13=[]
    o16=[]
    o18=[]
    ne22=[]
    log_mlt_D_mix=[]
    eps_nuc=[]
    grada=[]
    gradr=[]
    luminosity=[]
    log_opacity=[]        
    logT=[]
    log_D_mix=[]
    log_D_conv=[]
    log_D_th=[]
    den_new=[]
    logRho=[]
    logS=[]
    entropy=[]
    p=ms.mesa_profile(logdir,mod,num_type='model')
    mass=p.get('mass')
    for i in range(len(specs)):
        eval(specs[i]).append(p.get(specs[i]))
    for i in range(len(specs2)):
        eval(specs2[i]).append(p.get(specs2[i]))
    j=0
    fig=mod
    close(fig)
    ifig=figure(fig,figsize=(7,9.5))
    ax=ifig.add_subplot(111)
    subplot(311)
    ax1=ifig.add_subplot(311)
    ax1.semilogy(mass,eval(specs[j])[0],lines[j],label=labels[j],linewidth=3)
    box = ax1.get_position()
    ax1.set_position([box.x0, box.y0, box.width * 0.83, box.height])
    #ax1.legend(loc='center', bbox_to_anchor=(1.22, 0.16),prop={'size':10.0})
    #ax1.legend(loc=4,prop={'size':10.0})
    ax1.set_ylabel(labels[j])
    setp(ax1.get_xticklabels(), visible=False)
    ylim(1e-4,1)
    j+=1
    ax2 = ax1.twinx()
    ax2.plot(mass,eval(specs[j])[0],lines[j+11],label=labels[j],linewidth=3)
    j+=1
    ax2.plot(mass,eval(specs[j])[0],lines[j-1],label=labels[j],linewidth=3)
    #ax2.legend(loc='center', bbox_to_anchor=(1.22, 1.15),prop={'size':10.0})
    ax1.legend(loc=2,prop={'size':10.0})
    ax2.legend(loc=3,prop={'size':10.0})
    ax2.set_ylabel(labels[j])
    #show()
    ylim(0,20)
    j+=1
    xlim(0.40,0.92)
    subplot(312)
    ax1=ifig.add_subplot(312)
    ax1.plot(mass,eval(specs[j])[0],lines[j],label=labels[j],linewidth=3)
    ylim(6,9)
    ax1.yaxis.set_major_locator(MultipleLocator(0.5))
    ax1.set_ylabel(labels[j])
    box = ax1.get_position()
    ax1.set_position([box.x0, box.y0, box.width * 0.83, box.height])
    #ax1.legend(loc='center', bbox_to_anchor=(1.22, 0.15),prop={'size':10.0})
    ax1.legend(loc=3,prop={'size':10.0})
    j+=1
    ax2 = ax1.twinx()
    ax2.set_ylabel(labels[j])
    ax2.semilogy(mass,eval(specs[j])[0],lines[j+10],label=labels[j],linewidth=3)
    #ax2.legend(loc='center', bbox_to_anchor=(1.22, 1.0),prop={'size':10.0})
    ax2.legend(loc=4,prop={'size':10.0})
    ylim(1,1e8)
    j+=1
    xlim(0.40,0.92)
    setp(ax1.get_xticklabels(), visible=False)
    subplot(313)
    ax1=ifig.add_subplot(313)
    ax1.plot(mass,eval(specs[j])[0],color='orange',linestyle='-',label=labels[j],linewidth=3)
    j+=1
    ax1.plot(mass,eval(specs[j])[0],lines[j],label=labels[j],linewidth=3)
    ax1.set_ylabel('{\\nabla}')
    box = ax1.get_position()
    ax1.set_position([box.x0, box.y0, box.width * 0.83, box.height])
    #ax1.legend(loc='center', bbox_to_anchor=(1.22, 0.15),prop={'size':10.0})
    j+=1
    ylim(0.0,0.40)
    ax1.set_xlabel('${M_{r}/M_\mathrm{\odot}}$')
    ax1.legend(loc=4,prop={'size':10.0})
    ax2 = ax1.twinx()
    ax2.plot(mass,eval(specs[j])[0],lines[j+3],label=labels[j],linewidth=3)
    #ax2.legend(loc='center', bbox_to_anchor=(1.22, 1.0),prop={'size':10.0})
    ax2.set_ylabel(labels[j])
    #ifig.text(0.5, 0.03,'$M/M_{\odot}$', ha='center', va='center', rotation='horizontal',color='b',fontsize=18)
    ax2.legend(loc=3,prop={'size':10.0})
    show()
    xlim(0.40,0.92)
    #ylim(6.0,9.5)
    savefig('D_profile_'+'0'+str(mod)+'.png')
    savefig('D_profile_'+'0'+str(mod)+'.pdf')
    savefig('D1_profile_'+'0'+str(mod)+'.eps')
    #close(ifig)

def vt(T,k,d,c):
    return log10((4.0/3)*(7.56577*10**-15)*3*10**10)+3*T-(k+2*d+log10(c))


for mod in mods:
    fig=mod+10
    close(fig)
    ifig=figure(fig,figsize=(5,5))
    ax3=ifig.add_subplot(111)
    log_vt=[]
    p=ms.mesa_profile(logdir,mod,num_type='model')
    mass=p.get('mass')
    for a in range(len(mass)):
        log_vt.append(vt(log_T[0][a],log_opacity[0][a],logRho[0][a],cp[0][a]))
    ax3.plot(mass,log_vt,linestyle='--',color='brown',linewidth=2,label='log $\nu_\mathrm{T}$')
    ax3.plot(mass,log_D_mix,linestyle='-',color='cyan',linewidth=2,label='log $D_\mathrm{total}$')
    ax3.set_xlabel('${M_{r}/M_\mathrm{\odot}}$')
    legend(loc=2)
    savefig('D_vt_'+'0'+str(mod)+'.png')
    savefig('D_vt_'+'0'+str(mod)+'.pdf')    
    savefig('D_vt_'+'0'+str(mod)+'.eps')