'''contour_multiplot.py

purpose: to run a grid of reasonable cloud densities and surface
densities and look at the CO, CII and OI abundnace contours; 

the SFR will be explicitly mandated though, not generated by any
subgrid model

'''

# Import the despotic library and various standard python libraries
from despotic import cloud,zonedcloud
from despotic.chemistry import NL99,NL99_GO
import numpy as np
import matplotlib.pyplot as plt
from astropy import units as u
from astropy import constants as constants
import pdb
from matplotlib.ticker import MaxNLocator
import matplotlib.cm as cm


########################################################################
# User-settable options
########################################################################
ngrids = 5

# Set up a range of volume densities
lognHgrid = np.linspace(1.,5,ngrids)
logsigmahgrid = np.linspace(1.7,4,ngrids) #Sigma_H in Msun/pc^2
sfrgrid = [1,10,100,1000]


nH = 10.**lognHgrid * u.cm**-3
column_density = 10.**logsigmahgrid * u.Msun/u.pc**2.
column_density_cgs = (column_density/constants.m_p).cgs

NZONES=8

########################################################################
# Program code
########################################################################
 
#set up initial plots


fig = plt.figure()
fig.subplots_adjust(hspace=0.25,wspace=0.25)
plotnum = 0 

for sf in sfrgrid:

    #make 2D lists to hold results
    xCO = np.zeros([ngrids,ngrids])
    xCp = np.zeros([ngrids,ngrids])
    xC = np.zeros([ngrids,ngrids])
    xHep =  np.zeros([ngrids,ngrids])
    xHe = np.zeros([ngrids,ngrids])
    xe =  np.zeros([ngrids,ngrids])
    xO =  np.zeros([ngrids,ngrids])
    xCHx =  np.zeros([ngrids,ngrids])
    xOHx = np.zeros([ngrids,ngrids])
    
    Tg = np.zeros([ngrids,ngrids])
    Td = np.zeros([ngrids,ngrids])
    sfr_grid =  np.zeros([ngrids,ngrids])
    colden_grid =  np.zeros([ngrids,ngrids])
    num_clouds_grid =  np.zeros([ngrids,ngrids])
    m_cloud_grid =  np.zeros([ngrids,ngrids])

    #these will serve as memory temperatures of the last iteration to
    #try to help the solver find temperatures more quickly
    oldTd = 30
    oldTg = 30
    
    for i in range(ngrids): #column_density
        for j in range(ngrids): #nH
            

            gmc=zonedcloud(colDen = np.linspace(column_density_cgs[i].value/NZONES,column_density_cgs[i].value,NZONES))
            gmc.addEmitter('c+',1.e-100)
            gmc.addEmitter('c',2.e-4)
            gmc.addEmitter('o', 4.e-4)
            gmc.addEmitter('co',1.e-100)
            

            
            # Turn on extrapolation
            for nz in range (NZONES):
                gmc.comp[nz].xH2 = 0.5
                gmc.comp[nz].xHe = 0.1
                gmc.emitters[nz]['co'].extrap=True
                gmc.emitters[nz]['c+'].extrap=True
                gmc.emitters[nz]['o'].extrap = True
                gmc.emitters[nz]['c'].extrap = True

    
            gmc.nH = nH[j].value
            gmc.Td = oldTd
            gmc.Tg = oldTg
            gmc.rad.TradDust = 30
            gmc.ionRate = 1.e-17 * sf
            gmc.rad.ionRate = 1.e-17 * sf
            gmc.chi = sf
            gmc.rad.chi = sf


            
            # Set virial ratio to unity
            gmc.setVirial()
            
            # Calculate chemical abundances
            gmc.setTempEq()
            oldTg = gmc.Tg
            oldTd = gmc.Td
            gmc.setChemEq(network=NL99_GO, evolveTemp = 'iterate', verbose=True)

            xCO[i,j] = gmc.chemabundances['CO']
            xCp[i,j] = gmc.chemabundances['C+']
            xC[i,j] = gmc.chemabundances['C']
           
        
       
            
   
    plotnum += 1
    ax = plt.subplot(2,2,plotnum)
    y = np.log10(nH.value)
    x = np.log10(column_density.value)
    X,Y = np.meshgrid(x,y)
    Z = np.log10(xCO)
    CS = plt.contourf(X,Y,Z,levels=np.linspace(-13,np.max(Z),20),cmap=cm.coolwarm)
    plt.colorbar(CS,format="%0.2f")
    
    #putting the fixings on the plots
    if plotnum == 2 or plotnum == 4:
        plt.setp(ax.get_yticklabels(),visible=False)
    if plotnum == 1 or plotnum == 2:
        plt.setp(ax.get_xticklabels(),visible=False)
    if plotnum == 3 or plotnum == 4:
        ax.set_xlabel('log$_{10}(\Sigma_H)$ (M$_\odot$ pc$^{-2}$)',fontsize=8)
   
       
    ax.set_ylabel('log$_{10}(nH)$ (cgs)',fontsize=8)
    

    ax.set_title('SFR = %s M$_\odot$ yr$^{-1}$' % sf)

   

fig.savefig('co_contours.png',dpi=300)
