#from numarray import * ## Changed to numpy Aug 2007
from numpy import *
import pyfits
#from random import *
from simpler_skymap_datons import *

# Read in each template ModelFlux file:

model1_hdulst        = pyfits.open('/AstroCode/EMC2/ModelFluxRebinndICPion1-1000GeV.060529.fits')

print 'Model 1 Data: '
print model1_hdulst[0].data
model1_hdulst.close()

model1_padded_hdulst = pyfits.open('/AstroCode/EMC2/ModelFluxPaddedICPion1-1000GeV.060529.fits')

print 'Model 1 padded Data:'
print model1_padded_hdulst[0].data
model1_padded_hdulst.close()

#
# For each element take MEAN and VARIANCE

naxis1 = model1_hdulst[0].header['NAXIS1']
naxis2 = model1_hdulst[0].header['NAXIS2']
mean = 0.
vari = 0.
for l1 in range(naxis1):
    for l2 in range(naxis2):
        mean = mean + model1_hdulst[0].data[0][l2][l1]
        vari = vari + model1_hdulst[0].data[0][l2][l1]*model1_hdulst[0].data[0][l2][l1]
mean = (mean/naxis1)/naxis2
vari = (vari/naxis1)/naxis2
sigma=sqrt(vari)

# Using same mean and variance for padded data:
naxis1_p = model1_padded_hdulst[0].header['NAXIS1']
naxis2_p = model1_padded_hdulst[0].header['NAXIS2']

print 'Mean, sigma: ', mean, sigma

model2_data        = model1_hdulst[0].data/mean
model2_padded_data = model1_padded_hdulst[0].data/mean

# For each element take sqrt( sqrt( sqrt( sqrt( data ) ) ) )  ALSO
# Chop off everything below b about 5 or 10 degrees. (For now assuming 90x120 or 128x128=padded!!!)

toobg = 1. #+ 1.(sigma/mean)
lowergalaxy = (naxis2/2)+5
for l1 in range(naxis1):
    for l2 in range(naxis2):
        mu = sqrt( sqrt( sqrt( sqrt( model2_data[0][l2][l1])  ) ) )
        if( (mu > toobg) or (l2 < lowergalaxy) ):
            model2_data[0][l2][l1] = 0.
        else:
            model2_data[0][l2][l1] = mu*mean/3.

print '\n Chopped Model 2 Data:'
print model2_data
newcomp_flux_hdu = pyfits.PrimaryHDU(data=model2_data)#,header=model1_hdulst[0].header)

lowergalaxy = (naxis2_p/2)+5
for l1 in range(naxis1_p):
    for l2 in range(naxis2_p):
        mu = sqrt( sqrt( sqrt( model2_padded_data[0][l2][l1]) ) )
        if( (mu > toobg) or (l2 < lowergalaxy) ) :
            model2_padded_data[0][l2][l1] = 0.
        else:
            model2_padded_data[0][l2][l1] = mu*mean/3.

print '\n Chopped Model Padded Data:'
print model2_padded_data
newcomp_padded_flux_hdu = pyfits.PrimaryHDU(data=model2_padded_data)#,header=model1_padded_hdulst[0].header)


#
# Store the new model component in its two formats:

newcomlst = [newcomp_flux_hdu, newcomp_padded_flux_hdu ]

NewCompFluxFileLst = ['ModelFluxRebinndNewComp.060529.fits', 'ModelFluxPaddedNewComp.060529.fits']

for ll in range( len(newcomlst) ):
    newcomlst[ll].writeto(NewCompFluxFileLst[ll])

#
# For each "old" Poisdaton file:
#
IntTimeLst = [7.0,22.0,73.0,242.0,806.0]

ExpFileLst = [
'ExposrTimPaddedIntTime7.0.060529.fits',
'ExposrTimPaddedIntTime22.0.060529.fits',
'ExposrTimPaddedIntTime73.0.060529.fits',
'ExposrTimPaddedIntTime242.0.060529.fits',
'ExposrTimPaddedIntTime806.0.060529.fits']

NewCompFluxLst = [newcomp_padded_flux_hdu,newcomp_padded_flux_hdu,newcomp_padded_flux_hdu,newcomp_padded_flux_hdu,newcomp_padded_flux_hdu]

NewCompDataFileLst = ['NewCompDataPaddedIntTime7.0.060529.fits','NewCompDataPaddedIntTime22.0.060529.fits','NewCompDataPaddedIntTime73.0.060529.fits','NewCompDataPaddedIntTime242.0.060529.fits','NewCompDataPaddedIntTime806.0.060529.fits']

OldPoisDatonsFileLst = ['PoisDatonPaddedIntTime7.0.060529.fits','PoisDatonPaddedIntTime22.0.060529.fits','PoisDatonPaddedIntTime73.0.060529.fits','PoisDatonPaddedIntTime242.0.060529.fits','PoisDatonPaddedIntTime806.0.060529.fits']

NewCompDatonsFileLst = ['NewCompPoisDatonPaddedIntTime7.0.060529.fits','NewCompPoisDatonPaddedIntTime22.0.060529.fits','NewCompPoisDatonPaddedIntTime73.0.060529.fits','NewCompPoisDatonPaddedIntTime242.0.060529.fits','NewCompPoisDatonPaddedIntTime806.0.060529.fits']

NewPoisDatonsFileLst = ['NewSumPoisDatonPaddedIntTime7.0.060529.fits','NewSumPoisDatonPaddedIntTime22.0.060529.fits','NewSumPoisDatonPaddedIntTime73.0.060529.fits','NewSumPoisDatonPaddedIntTime242.0.060529.fits','NewSumPoisDatonPaddedIntTime806.0.060529.fits']



# Quick-o check of the other lists:
if ( ( len(IntTimeLst) != len(ExpFileLst) ) or ( len(IntTimeLst) != len(NewCompDataFileLst) ) or ( len(IntTimeLst) != len(NewCompDatonsFileLst) )  ):
    print 'Whoa! Dude! Fatal Error! Your IntTimeLst and various FileLst dont have the same number of entries!!!'

#
#
##FOR each item in lists:
for kk in range( len(IntTimeLst) ):

#   Read in old Exposure
    this_exp_hdulst = pyfits.open(ExpFileLst[kk])
#
#   Multiply fake component by associated IntTime*Exposure
#      to get FakeComp piece
    newcompdata_hdu = simpler_outer_products( NewCompFluxLst[kk], this_exp_hdulst[0] )
    newcompdata_HDU = pyfits.PrimaryHDU(data=newcompdata_hdu.data,header=newcompdata_hdu.header)
    newcompdata_HDU.writeto(NewCompDataFileLst[kk])
    this_exp_hdulst.close()
#
#   Get new Poisson Datons with the new component:
    newcompdatons_hdu = simpler_poisson_datons(newcompdata_HDU)
#
#   Write out fits and txt files for each NewComp
    print 'Sample data012 of output from outerprod:', newcompdatons_hdu.data[0][1][2]
    print 'Header of output from outerprod:', newcompdatons_hdu.header,'\n'
    newcompdatons_HDU = pyfits.PrimaryHDU(data=newcompdatons_hdu.data,header=newcompdatons_hdu.header)
    newcompdatons_HDU.writeto(NewCompDatonsFileLst[kk])


#
#
# For each array, SUM the new piece with the old one to get NewPoissDatons* .
#   Read in previous fits daton files:
    olddatons_HDU = pyfits.open(OldPoisDatonsFileLst[kk/2])
    print 'kk: ', kk, '  OldPoisDatonsFileLst[kk/2]: ', OldPoisDatonsFileLst[kk/2], '\n'
    print 'Sample data012 of 1st input to sum:', olddatons_HDU[0].data[0][1][2]
    print 'Header of 1st input to sum:', olddatons_HDU[0].header,'\n'

#   SUM
    NewPoisDatons_HDU = simpler_twosum(old_datons_HDU[0],newcompdatons_HDU)
    print 'Sample data012 of output from sum:', olddatons_HDU[0].data[0][1][2]
    print 'Header of output from sum:', olddatons_HDU[0].header,'\n'

    olddatons_HDU.close()
#
#   Now write 'em out, fits ant txt
    NewPoisDatons_HDU.writeto(NewPoisDatonsFileLst[kk])


# All done?? For now??

# ModelFluxPaddedICPion1-1000GeV.060529.fits
# ModelFluxRebinndICPion1-1000GeV.060529.fits
# PoisDatonPaddedIntTime7.0.060529.fits
# ModelDataPaddedIntTime7.0.060529.fits
# ExposrTimPaddedIntTime7.0.060529.fits
# ExposrTimRebinndIntTime7.0.060529.fits
# PoisDatonPaddedIntTime22.0.060529.fits
# ModelDataPaddedIntTime22.0.060529.fits
# ExposrTimPaddedIntTime22.0.060529.fits
# ExposrTimRebinndIntTime22.0.060529.fits
# PoisDatonPaddedIntTime73.0.060529.fits
# ModelDataPaddedIntTime73.0.060529.fits
# ExposrTimPaddedIntTime73.0.060529.fits
# ExposrTimRebinndIntTime73.0.060529.fits
# ModelDataPaddedIntTime242.0.060529.fits
# PoisDatonPaddedIntTime242.0.060529.fits
# ExposrTimPaddedIntTime242.0.060529.fits
# ExposrTimRebinndIntTime242.0.060529.fits
# ModelDataPaddedIntTime806.0.060529.fits
# PoisDatonPaddedIntTime806.0.060529.fits
# ExposrTimPaddedIntTime806.0.060529.fits
# ExposrTimRebinndIntTime806.0.060529.fits
# ModelFluxPaddedBremICPion1-1000GeV.060529.fits
# ModelFluxRebinndBremICPion1-1000GeV.060529.fits