from params import *
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid
import numpy as np
import yt
from yt.units.yt_array import YTQuantity
from functions import *

"""
The datapath,numhalos,halos,halolist,sourcelist, etc are defined in the params file.
I've also made my own function file which uses the yt functions and adds the units I want and so on but I've stripped away that part here.
There's also a ton of numpy arrays for different stuff, so just assume the proj_plot_halo_plot[l,i,j,k] actually exist.
"""

for l,simpath in enumerate(datapath):
    ts = yt.load(simpath)
    for i,index in enumerate(sourcelist):
        ds = ts[index-1] 
        for j in range(numhalos):
            curr_halo = halolist[l][j]
            cen = halos[l][curr_halo] #Center of current halo
            rad = haloradius[l][curr_halo]
            wid = 2.*rad # diameter of current halo

            if ('ramses', 'Density') in ds.field_list: #Check for gas density in sim
                #Make the cube around the halo
                left_edge  = cen - wid #A bit larger than the halo but we want that
                right_edge = cen + wid #A bit larger than the halo but we want that
                cube       = ds.region(cen, left_edge, right_edge, fields=None, ds=ds, field_parameters=None, data_source=None)

                #Make projection through cube along each axis
                for k,dim in enumerate(['x','y','z']):
                    fieldtype = 'gas'
                    field = 'density'
                    proj_plot_halo_plot[l,i,j,k] =yt.ProjectionPlot(ds, dim, (fieldtype,field), center=cen, width=wid, data_source=cube)
                    #proj_plot_halo_plot[l,i,j,k] =yt.ProjectionPlot(ds, dim, (fieldtype,field), center=cen, width=wid) #No difference between this and above.

                #radial density profile 
                sph = ds.sphere(cen, 1.*rad)
                fieldtype = 'gas'
                field = 'density'
                weight = None 
                prof_density_halo_plot[l,i,j]        = yt.create_profile(sph,'radius',(fieldtype,field),weight_field=weight)
                prof_density_halo_plot_specs[l,i,j]  = dict(linewidth=2, alpha=0.7)
                prof_density_halo_plot_labels[l,i,j] = 'Halo #%d'%curr_halo

"""After this there's some code where I combine the plots into a larger figure to compare them but that's not important."""