import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
#plt.ioff()
import matplotlib.colorbar as cbar
import numpy as np
import matplotlib.colors as colors
import matplotlib.cm as cm
import matplotlib.gridspec as gridspec
import matplotlib.lines as lines
import matplotlib.patches as patch
import gc
from matplotlib.font_manager import FontProperties
from yt.mods import *
import pdb
import glob
import math
if __name__=="__main__":
    enable_parallelism()
    fdir = '/home/rosen/pfs/ORION2/RAYTRACE/'
    fstrs = ['RadDomSphereIso_64/', 'RadDomSphereIso_128/', 'RadDomSphereIso_256/']# 'RadDomSphereIso_AMR/']

    LSUN = 3.84e33
    Lum=1e6*LSUN
    CSOL = 2.99792458e10
    SIGMA = 5.670373e-5
    PCTOCM = 3.08567758e18
    SECYR = 3600.*24*365.25
    PCCM = 3.08567758e18
    n0 = 1e5
    rho0=n0*2.33*1.67e-24
    #setup the figure
    colors = ['Teal', 'DeepPink', 'MediumOrchid', 'MediumVioletRed', 'DarkSlateBlue']
    ls = ['-.','--',':','-','-.','--']
    sym = ['D','*','o','h', '^','v']
    strN = [r'$N_0=64$', r'$N_0=128$', r'$N_0=256$']
    fig, (ax1, ax2) =plt.subplots(2,1, sharex=True)
    fig.subplots_adjust(hspace=0.075,wspace=0.05, left = 0.15, right=0.98, top=0.98, bottom=0.125)
    
    j = 0
    for fstr in fstrs:
        fpath = fdir + fstr + 'data.[0-9][0-9][0-9][0-9].3d.hdf5'
        figpath = fdir +fstr + 'figs/'
        filenames = glob.glob(fpath)
        filenames.sort()
        nfiles = len(filenames)
        

        pf_last = load(filenames[nfiles-1])
        rho_min =  pf_last.find_min('density')
        rho_max =  pf_last.find_max('density')
        ts = DatasetSeries(filenames)
        Rsh = np.zeros(nfiles)
        Ra = np.zeros(nfiles)
        Resid = np.zeros(nfiles)
        tsim = np.zeros(nfiles)
        width = (4,'pc')
        i = 0
        for pf in ts.piter():
            if 0:
                slcx = SlicePlot(pf, 0, ['density', 'velocity_magnitude', 
                                     'directrad-dedt-density'], width = width)
                slcx.set_zlim('density', rho_min[0], rho_max[0])
                slcx.set_zlim('velocity_magnitude', .001, 1e7)
                time="%5.4f" % (pf.current_time.v/SECYR)
                slcx.annotate_title('t = %s yr' % (time))
                slcx.save(figpath + '%s' %(pf))
                
            #Get time and Rsh
            dd = pf.all_data()
            time =  pf.current_time.v
            tsim[i] = (time/(SECYR*1e6))
            Ra[i] = (3./(2*math.pi*rho0)*(Lum/CSOL))**0.25*(time)**0.5
            num=sum(dd["Rshock"].v*dd["density"].v*(dd["x"].v**2.+ dd["y"].v**2.+dd["z"].v**2.)**0.5)
            denom=sum(dd["Rshock"].v*dd["density"].v)
            if denom == 0:
                Rsh[i] = 0
                Resid[i] = 0
            else:
                Rsh[i] = num/denom
                Resid[i] = (Ra[i]-Rsh[i])/Ra[i]
            i = i+1
        #Make plot
        ax1.plot(tsim, Rsh/PCCM, color = colors[j], ls = ls[j], label = strN[j])
        ax2.plot(tsim, Resid, color = colors[j], ls = ls[j], label = strN[j])
        if j==2:
            #Plot R_a
            ax1.plot(tsim, Ra/PCCM, color = colors[j], ls = ls[3], label = r'$R_{\rm true}$')
        j = j+1
    fontP = FontProperties()
    fontP.set_size('medium')
    handles, labels = ax1.get_legend_handles_labels()
    ax1.legend(handles, labels, loc='lower right', ncol=1, prop=fontP)

    ax1.set_ylabel(r'$R_{\rm sh}$ [pc]')
    ax2.set_ylabel(r'$(R_{\rm true}-R_{\rm sh})/R_{\rm true}$')
    ax2.set_xlabel('Time [Myr]')
    fig.savefig('/home/rosen/pfs/ORION2/paper_figures/numPaper/RadDomSphere_Rshell.png')
    plt.clf()