import yt
from yt.units import dimensions
from yt.units.yt_array import YTArray
import numpy as np
#import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt

spNms  = ['H2', 'H', 'O2', 'OH', 'H2O', 'HO2', 'H2O2', 'CH3', 'CH4', 'CO', 'CO2', 'CH2O', 'C2H2', 'C2H4', 'C2H6', 'NH3', 'NO', 'HCN', 'N2']

HfO    = [0.0, 216269222.374, 0.0, 2292826.50902, -13423893.4949, 316904.97153, -4001375.87653, 9689965.02785, -4668373.7711, -3946437.03767, \
         -8941308.52674, -3859845.44206, 8707735.03834, 1870243.41835, -2788439.10336, -2695065.81111, 3009074.80936, 4937073.18373, 0.0]

omgdot = ['H2_reaction_rate mea', 'H_reaction_rate mean', 'O2_reaction_rate mea', 'OH_reaction_rate mea', 'H2O_reaction_rate me', 'HO2_reaction_rate me', \
          'H2O2_reaction_rate m', 'CH3_reaction_rate me', 'CH4_reaction_rate me', 'CO_reaction_rate mea', 'CO2_reaction_rate me', 'CH2O_reaction_rate m', \
          'C2H2_reaction_rate m', 'C2H4_reaction_rate m', 'C2H6_reaction_rate m', 'NH3_reaction_rate me', 'NO_reaction_rate mea', 'HCN_reaction_rate me', \
          'N2_reaction_rate mea']

#print 'HfO[16]',HfO[16]
#print 'omgdot[16]',omgdot[16]

#print 'HfO.shape',HfO.shape
#print 'omgdot.shape',omgdot.shape

def _HRR(field, data):
    #data._debug()
    N = data['Tmean'].shape
    HRR=np.zeros(N)
    i=0
    for sp in spNms:
        HRR=HRR+(HfO[i]*data[omgdot[i]])
        i=i+1
    HRR = -HRR
    #data._debug()
    return YTArray(HRR, 'J/m**3/s')

yt.add_field(('gas','HRR'), function=_HRR, units='J/m**3/s')#, dimensions=dimensions.'HRR')


units_override = {'length_unit':(1.0, 'm'),
                  'time_unit'  :(1.0, 's'),
                  'mass_unit'  :(1.0, 'kg')}

ds = yt.load("/data1/ahoffie/iw-dm-4/HighTempLowNOx/JICF_EXP/Scaled_GT_run_12mmPre_r2/out_jicf/post/merged_stats_00027_00034", units_override=units_override)

print 'type(ds)=',type(ds)
print 'ds.length_unit=',ds.length_unit

D        = 4e-3
xup      =-5.5  #D
xdn      = 14.5 #
xdn_chop = 12.5 #

yup      =-2.5
ydn      = 9.5

xlen_new = (abs(xup)+xdn_chop)*D
ylen     = (abs(yup)+ydn)*D

print 'xlen_new',xlen_new
print 'ylen    ',ylen

center_x = xlen_new/2.+xup*D
center_y = ylen/2.+yup*D
center_z = 0.0

print 'center_x',center_x
print 'center_y',center_y
#print 'center_z',center_z

prj = yt.ProjectionPlot( ds, 'z', 'HRR', origin='native', method='integrate', weight_field='ones',
                         center = [center_x, center_y],
                         width  = ( (xlen_new,'m'),(ylen,'m') )
                       )

#prj.set_axes_unit('m')
prj.save('HRR_prj.png')

resf=0.1992e-4 #m/pix
px_x=xlen_new/resf
px_y=ylen/resf
#
print 'px_x',px_x
print 'px_y',px_y
#
width   = ( (xlen_new,'m'),(ylen,'m') )
res     = [px_x, px_y]
center  = [center_x, center_y]
#HRR_frb = prj.to_frb(width, res, center=center)
HRR_frb = prj.frb['HRR']
#
#plt.imshow(np.array(frb['HRR']))
plt.imshow(HRR_frb)
plt.show()
plt.savefig('HRR_frb_prj.png')

                                                                                                                                                                       97,0-1        Bot