from yt.mods import *
from matplotlib import use; use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from yt.analysis_modules.halo_finding.api import *
import sys
import glob
import re

#try: 
#    from mpi4py import MPI
#    comm = MPI.COMM_WORLD
#    nprocs = comm.size
#    rank = comm.rank
#except:
#    nprocs=1
#    rank=0


fn, = glob.glob('RedshiftOutput????') 
output = int(re.sub(r'\D',"",fn))
pf = load(fn)

vl=pf["RefineRegionLeftEdge"]
vr=pf["RefineRegionRightEdge"]
#sv = pf.h.region((vl+vr)/2, vl, vr)
width=(vr-vl).max()
c=(vl+vr)/2
z=pf["CosmologyCurrentRedshift"]

wkpc=width*40000/(1.0+z)
#width=wkpc*(z+1.0)/40.0/1000
#c=np.array([0.4341635423, 0.4348677931, 0.5726667823])
#c=[0.5,0.5,0.5]
RL=c-width/2 #RL=[0.425, 0.425, 0.425]
RR=c+width/2 #RR=[0.575, 0.575, 0.575]
sv = pf.h.region(c, RL, RR)

#deltaz= 0.02
#c[2]=c[2]+deltaz
#haloes = LoadHaloes(pf, "RD%04i"%output)

#pc = PlotCollection(pf, center=c)
#p=pc.add_projection('Density',2,data_source=sv)
#p = ProjectionPlot(pf, 2, "Density", c, width=(2*wkpc, 'kpc'),data_source=sv)
#pc.add_projection('Density',1,data_source=sv)
#pc.add_projection('Density',0,data_source=sv)

if 2 == 2:
 for dz in range(-5,6,1):
   deltaz=0.01*dz
   p = SlicePlot(pf, 1, ["H_Ionization_Fraction","Density"], [c[0],c[1]+deltaz,c[2]],width=(wkpc, 'kpc'))#,data_source=sv)
   #p.annotate_particles(wkpc/pf['kpc'],p_size=0.5,stars_only=True)
   #p.annotate_particles(wkpc/pf['kpc'],p_size=0.5,col='w',ptype=7)
   #p.annotate_hop_circles(haloes,min_size=200,width=wkpc/pf['kpc'])
   p.set_zlim("H_Ionization_Fraction",1e-4,1.0)
   p.set_zlim("Density",1e-30,1e-25)
   #p.plots['H_Ionization_Fraction'].hide_colorbar()
   #p.plots['H_Ionization_Fraction'].hide_axes()
   p.save('%3.1fMpc_%4.2f_%4.2f'%(wkpc*(z+1.0)/1000,z,deltaz))

if 1 == 0:
   p = SlicePlot(pf, 1, ["H_Ionization_Fraction","Density"], c, width=(wkpc, 'kpc'))#,data_source=sv)
   #p.annotate_particles(wkpc/pf['kpc'],p_size=0.5,stars_only=True)
   #p.annotate_particles(wkpc/pf['kpc'],p_size=0.5,col='w',ptype=7)
   #p.annotate_hop_circles(haloes,min_size=200,width=wkpc/pf['kpc'])
   p.set_zlim("H_Ionization_Fraction",1e-4,1.0)
   p.set_zlim("Density",1e-30,1e-25)
   p.save('%3.1fMpc_%4.2f'%(wkpc*(z+1.0)/1000,z))

if 1 == 0:
   p = SlicePlot(pf, 0, ["H_Ionization_Fraction","Density"], c, width=(wkpc, 'kpc'))#,data_source=sv)
   #p.annotate_particles(wkpc/pf['kpc'],p_size=0.5,stars_only=True)
   #p.annotate_particles(wkpc/pf['kpc'],p_size=0.5,col='w',ptype=7)
   #p.annotate_hop_circles(haloes,min_size=200,width=wkpc/pf['kpc'])
   p.set_zlim("H_Ionization_Fraction",1e-4,1.0)
   p.set_zlim("Density",1e-30,1e-25)
   p.save('%3.1fMpc_%4.2f'%(wkpc*(z+1.0)/1000,z))


#p = ProjectionPlot(pf, 1, ["Density","TotalMetallicity"], c, width=(2*wkpc, 'kpc'),data_source=sv)
#p.annotate_particles(2*wkpc/pf['kpc'],ptype=5)#,stars_only=True)
#p.annotate_hop_circles(haloes,min_size=200,width=wkpc/pf['kpc'])
#p.save('%3.1fkpc_%3.1f'%(2*wkpc,z))
#p = ProjectionPlot(pf, 0, ["Density","TotalMetallicity"], c, width=(2*wkpc, 'kpc'),data_source=sv)
#p.annotate_particles(2*wkpc/pf['kpc'],ptype=5)#,stars_only=True)
#p.annotate_hop_circles(haloes,min_size=200,width=wkpc/pf['kpc'])
#p.save('%3.1fkpc_%3.1f'%(2*wkpc,z))

#pc.add_projection('TotalMetallicity',2,data_source=sv)
#pc.add_projection('Temperature',2,weight_field='Density',data_source=sv)
#pc.add_projection('TotalMetallicity',2,weight_field='Density',data_source=sv)
#pc.set_width(2*wkpc,'kpc')
#pc.set_zlim(1e-1,1e-3)
#pc.save('%3.1fkpc_%3.1f'%(2*wkpc,z))
#del pc

#pc = PlotCollection(pf, center=c)
#pc.add_projection('Density',2,data_source=sv)
#pc.add_projection('TotalMetallicity',2,data_source=sv)
#pc.add_projection('Temperature',2,weight_field='Density',data_source=sv)
#pc.add_projection('Temperature',1,weight_field='Density',data_source=sv)
#pc.add_projection('Temperature',0,weight_field='Density',data_source=sv)
#pc.add_projection('TotalMetallicity',2,weight_field='Density',data_source=sv)
#pc.set_width(2*width,'1')
#pc.set_zlim(1e-1,1e-3)
#pc.save("Halo")
#del pc

#pc = PlotCollection(pf, center=c)
#pc.add_projection('Density',2,data_source=sv)
#pc.add_projection('TotalMetallicity',2,data_source=sv)
#pc.add_projection('Temperature',2,weight_field='Density',data_source=sv)
#pc.add_projection('TotalMetallicity',2,weight_field='Density',data_source=sv)
#pc.add_projection('TotalMetallicity',1,weight_field='Density',data_source=sv)
#pc.add_projection('TotalMetallicity',0,weight_field='Density',data_source=sv)
#pc.set_width(2*width,'1')
#pc.set_zlim(1e-1,1e-6)
#pc.save("Halo")
#del pc

#c=na.array([0.4882009076, 0.4725569205, 0.4654072348])
#pc = PlotCollection(pf, center=c)
#sv = pf.h.region(c,c-5.0/pf['kpc'],c+5.0/pf['kpc'])
#pc.add_projection('Density',2,data_source=sv)
#pc.add_projection('Temperature',2,weight_field='Density',data_source=sv)
#pc.add_projection('TotalMetallicity',2,weight_field='Density',data_source=sv)
#pc.set_width(10,'kpc')
#pc.set_zlim(2e-1,1e-4)
#pc.save("most_massive_halo")
#del pc