import matplotlib
matplotlib.use('Agg')
from yt.mods import *
from yt.funcs import *
from matplotlib import pyplot as plt
import numpy as np
from matplotlib import colors


pf = EnzoStaticOutput('RedshiftOutput0028')

#z=pf.h["CosmologyCurrentRedshift"]
#halo_list=HaloFinder(pf)

#pc.add_phase_sphere(200.0, 'kpc', ["Density", "MagneticField", "CellVolume"],weight = None)
#pc.add_phase_sphere(200.0, 'kpc', ["Density", "MagneticField", "MagneticEnergy"],weight = None)

vl=pf["RefineRegionLeftEdge"]
vr=pf["RefineRegionRightEdge"]
region = pf.h.region((vl+vr)/2, vl, vr)
#region = pf.h.region([0.5,0.5,0.5],[0.49,0.49,0.49],[0.51,0.51,0.51])

#sphere = pf.h.sphere(c,h_a)

r_min = pf.h.get_smallest_dx()*pf['cm']

d_min = 0.005 #region['Density'].max()
d_max = 10000 #region['Density'].min()

#T_max = 1e4 #region['MagneticField'].max()
#T_min = 1e8 #region['MagneticField'].min()
IF_min = 0.0
IF_max = 1.0

d_bins = 300    #60
I_bins = 100  #200

prof2d = BinnedProfile2D(region, d_bins, "Baryon_Overdensity", d_min, d_max, True,I_bins, "H_Ionization_Fraction", IF_min, IF_max, False, True)
prof2d.add_fields('CellVolume',weight=None)
prof2d.add_fields('CellMass',weight=None)
#prof2d.add_fields('MagneticEnergy',weight=None)
#prof2d.add_fields('beta',weight='CellVolume')
#prof2d.add_fields('beta_k',weight='CellVolume')

intmin = None
intmax = None

figurename = 'phase_OverDensity_HIF'
cbarlabel=r'$CellVolume [cm^{3}]$'
xlabel=r'$Over Density$'
ylabel='$x$'


#fig = plt.figure()
#ax = fig.add_subplot(111)

        # determine the axes values
        # for attr in self.odmach2d.attrs.listnames():
        #    if attr.__contains__('x-axis'):
        #        xvals = self.odmach2d.attrs[attr]
        #    if attr.__contains__('y-axis'):
        #        yvals = self.odmach2d.attrs[attr]

xvals = prof2d['Baryon_Overdensity']
yvals = prof2d['H_Ionization_Fraction']

dy = yvals[1]-yvals[0]
dlogx = np.log10(xvals[1])-np.log10(xvals[0])
        
for field,cbarlabel in zip(['CellVolume','CellMass'],[r'$CellVolume [cm^{3}]$',r'$CellMass [g]$']):
#for field,cbarlabel in zip(['CellVolume','CellMass'],[r'$CellVolume [cm^{3}]$',r'$CellMass [g]$']):
    fig = plt.figure()
    ax = fig.add_subplot(111)

    data=prof2d[field]
    xvals = prof2d['Baryon_Overdensity']
    yvals = prof2d['H_Ionization_Fraction']

    data /= dlogx*dy

    if field == 'CellVolume':
       temp_min = r_min**3/dlogx/dy
    else :
       temp_min = data[data > 0].min()

    intmin = int(np.log10(temp_min)) + 1
    intmax = int(np.log10(data[1:,1:].max())) + 1


    mynorm = colors.LogNorm([intmin, intmax])
    im = ax.contourf(xvals, yvals, data.T, np.logspace(intmin,intmax,3), norm=mynorm)

    if xvals.size == data.T.shape[0]:
        xvals = np.append(xvals, 2.*10.**np.log10(xvals[-1]) - 10.**np.log10(xvals[-2]))
    if yvals.size == data.T.shape[0]:
        yvals = np.append(yvals, 2.*(yvals[-1] -yvals[-2]))

    X, Y = np.meshgrid(xvals,yvals)
    im = ax.pcolor(X,Y,data.T, norm=mynorm)

    ax.set_xscale('log')
    #ax.set_yscale('log')
    cbar = fig.colorbar(im, fraction=0.05, format = matplotlib.ticker.LogFormatterMathtext(), ticks=np.logspace(intmin,intmax, (intmax-intmin)+1), pad=0.01, shrink=0.9)
    cbar.set_label('%s' % cbarlabel)

#ax.axis([])

    plt.xlim(d_min, d_max)
    plt.ylim(0, 1)

    plt.xlabel('%s' % xlabel)
    plt.ylabel('%s'% ylabel)
    plt.savefig('%s_BaryonOverDensity_HIF_%s'%(pf,field))
    fig.clf()