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)

#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 rank == 2:
   p = ProjectionPlot(pf, 2, ["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('all',1e-4,1e-2)
   p.save('%3.1fMpc_%3.1f'%(wkpc*(z+1.0)/1000,z))

if rank == 1:
   p = ProjectionPlot(pf, 1, ["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('all',1e-4,1e-2)
   p.save('%3.1fMpc_%3.1f'%(wkpc*(z+1.0)/1000,z))

if rank == 0:
   p = ProjectionPlot(pf, 0, ["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('all',1e-4,1e-2)
   p.save('%3.1fMpc_%3.1f'%(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