# this is a python interactive script. It is meant to be run with the
# ed command from an interactive ipython -pylab session. All loaded
# variables will be in your local name space.

# need to read the files again?
read_new=False
read_new=True

noffset=750  # get rid of some pre-ms                                                                   

# Select cases from run.list file or list
list='specified'    # use list below, include the 'LOGS' dir name
dirs=['M2.00Z0.010_1.1_i11/LOGS_1','M2.00Z0.010_1.1_i11/LOGS']
list='run.list'   # use run.list

xlim_mod_min=67000
xlim_mod_max=79500
xlim_mod_min=-1       # automatic xlim

# things plot vs star_age for each case:
things=['log_L','log_LH','log_LHe']
things_ylims=[1,7]

# user defined parameters above this line
import mesa as ms
import utils
import os
symbs=utils.symbol_list('lines1')

if list is 'specified':
    print "User specified list of cases:"
    print dirs
elif list is 'run.list':
    print "Taking cases from run.list file: "
    rl=file('run.list')
    lines=rl.read()
    all_lines=lines.splitlines()
    dirs=[]
    for line in all_lines:
        if 'end' in line:
            break
        else:
            dirs.append(line.split()[0])
else:
    print "EROOR: list type not known"

if read_new:
    m=[]
    for this_run in dirs:
        if list is 'specified':
            this_m=ms.star_log(this_run)
        else:
            this_m=ms.star_log(this_run+'/LOGS')
        m.append(this_m)

# Now I have all data and I can start doing some plots below:

figure(1)
i=0
for case in m:
    logTeff=case.get('log_Teff')
    logL=case.get('log_L')
    plot(logTeff[noffset:],logL[noffset:],symbs[i],label=dirs[i])
    i += 1

case.xlimrev()
legend(loc=2)
xlabel('log Teff')
ylabel('log L')

figure(2)
i=0
for case in m:
    star_age=case.get('star_age')
    h1_boundary_mass=case.get('h1_boundary_mass')
    star_mass=case.get('star_mass')
    plot(star_age[noffset:],h1_boundary_mass[noffset:],symbs[i],label=dirs[i])
    plot(star_age[noffset:],star_mass[noffset:],symbs[i],label=dirs[i])
    i += 1
legend(loc=2)
xlabel('time / yr')
ylabel('M_H - core mass')


figure(3)
i=0
for case in m:
    model_number=case.get('model_number')
    h1_boundary_mass=case.get('h1_boundary_mass')
    star_mass=case.get('star_mass')
    plot(model_number[noffset:],h1_boundary_mass[noffset:],symbs[i],label=dirs[i])
    plot(model_number[noffset:],star_mass[noffset:],symbs[i],label=dirs[i])
    i += 1
legend(loc=2)
xlabel('model number')
ylabel('M_H - core mass')

figure(4)
i=0
for case in m:
    star_age=case.get('star_age')
    model_number=case.get('model_number')
    C=case.get('surface_c12')
    O=case.get('surface_o16')
    CO=C*4./(O*3.)
    plot(model_number[noffset:],CO[noffset:],symbs[i],label=dirs[i])
    i += 1
legend(loc=2)
xlabel('model number')
ylabel('C/O ratio')

figure(55)
i=0
for case in m:
    case.plot('model_number','v_div_csound_surf',legend=dirs[i],shape=symbs[i])
    i += 1
legend(loc=2)
xlabel('model number')
ylabel('v_div_csound_surf')
if xlim_mod_min >= 0:
    xlim(xlim_mod_min,xlim_mod_max)
ylim(-10.,10.)

figure(6)
i=0
for case in m:
    case.plot('model_number','log_abs_mdot',legend=dirs[i],shape=symbs[i])
    i += 1
legend(loc=2)
xlabel('model number')
ylabel('log_abs_mdot')
if xlim_mod_min >= 0:
    xlim(xlim_mod_min,xlim_mod_max)
ylim(-7,-3.5)

figure(7)
i=0
for case in m:
    case.plot('star_age','log_R',legend=dirs[i],shape=symbs[i])
    i += 1
legend(loc=2)
xlabel('model number')
ylabel('log_R')
if xlim_mod_min >= 0:
    xlim(xlim_mod_min,xlim_mod_max)


i=10
for case in m:
    case.kippenhahn(i,'model',t0_model=0)
#    title('Mini= '+str(case.header_attr['initial_mass']))
    title(dirs[i-10])
    i += 1

i=20
for case in m:
    figure(i)
    j = 0
    for thing in things:
        case.plot('model_number',thing,legend=thing,shape=symbs[j])
        j += 1
    title(dirs[i-20])
    ylim(things_ylims[0],things_ylims[1])
    i += 1