#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
vari = vari - (mean*mean)
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 rotate and rescale.

toobg = 1. #+ 1.(sigma/mean)
tooltl = 0.2
print 'tooltl, toobg:', tooltl, toobg
center1 = naxis1/2					# center of image
center2 = naxis2/2					#
cosrot = 0.						# rotate by about 90 deg
sinrot = 1.						# 
model3_data = zeros((1,naxis2,naxis1),dtype=float)	# zero out nu2 comp

for l1 in range(naxis1):
    mid_l1 = ((l1 - center1)/2)				 # First it is shrunk
    for l2 in range(naxis2):
        mid_l2 = l2 - center2
        new_l1 = int(mid_l1*cosrot - mid_l2*sinrot )+ center1 # Then both axes rotated
        new_l2 = int(mid_l1*sinrot + mid_l2*cosrot) + center2 #
        if( (0 <= new_l1) and (new_l1 < naxis1) ):
            if( (0 <= new_l2) and (new_l2 < naxis2) ):
                mu = model2_data[0][l2][l1]*model2_data[0][l2][l1]
                if( mu< tooltl):
                    model3_data[0][new_l2][new_l1] = 0.
                elif(mu > toobg):
                    model3_data[0][new_l2][new_l1] = 0.
                else:
                    model3_data[0][new_l2][new_l1] = model3_data[0][new_l2][new_l1] + mu*mean*4.


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


center1 = naxis1_p/2					# center of image
center2 = naxis2_p/2					#
cosrot = 0.						# rotate by about 30 deg
sinrot = 1.						# 
model3_padded_data = zeros((1,naxis2_p,naxis1_p),dtype=float)	# zero out nu2 comp

for l1 in range(naxis1_p):
    mid_l1 = ((l1 - center1)/2)				 # First it is shrunk
    for l2 in range(naxis2_p):
        mid_l2 = l2 - center2
        new_l1 = int(mid_l1*cosrot - mid_l2*sinrot )+ center1 # Then both axes rotated
        new_l2 = int(mid_l1*sinrot + mid_l2*cosrot) + center2 #
        if( (0 <= new_l1) and (new_l1 < naxis1_p) ):
            if( (0 <= new_l2) and (new_l2 < naxis2_p) ):
                mu = model2_padded_data[0][l2][l1]*model2_padded_data[0][l2][l1]
                if( (mu< tooltl) or (mu > toobg) ):
                    model3_padded_data[0][new_l2][new_l1] = 0.
                else:
                    model3_padded_data[0][new_l2][new_l1] = model3_padded_data[0][new_l2][new_l1] + mu*mean*4.

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


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

nu2comlst = [nu2comp_flux_hdu, nu2comp_padded_flux_hdu ]

Nu2CompFluxFileLst = ['ModelFluxRebinndNu2Comp.060529.fits', 'ModelFluxPaddedNu2Comp.060529.fits']

for ll in range( len(nu2comlst) ):
    nu2comlst[ll].writeto(Nu2CompFluxFileLst[ll])

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

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

Nu2CompFluxLst = [nu2comp_padded_flux_hdu,nu2comp_padded_flux_hdu,nu2comp_padded_flux_hdu,nu2comp_padded_flux_hdu,nu2comp_padded_flux_hdu,nu2comp_padded_flux_hdu]

Nu2CompDataFileLst = ['Nu2CompDataPaddedIntTime7.0.060529.fits','Nu2CompDataPaddedIntTime22.0.060529.fits','Nu2CompDataPaddedIntTime73.0.060529.fits','Nu2CompDataPaddedIntTime242.0.060529.fits','Nu2CompDataPaddedIntTime539.0.060529.fits','Nu2CompDataPaddedIntTime806.0.060529.fits']

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

Nu2CompDatonsFileLst = ['Nu2CompPoisDatonPaddedIntTime7.0.060529.fits','Nu2CompPoisDatonPaddedIntTime22.0.060529.fits','Nu2CompPoisDatonPaddedIntTime73.0.060529.fits','Nu2CompPoisDatonPaddedIntTime242.0.060529.fits','Nu2CompPoisDatonPaddedIntTime539.0.060529.fits','Nu2CompPoisDatonPaddedIntTime806.0.060529.fits']

Nu2PoisDatonsFileLst = ['Nu2SumPoisDatonPaddedIntTime7.0.060529.fits','Nu2SumPoisDatonPaddedIntTime22.0.060529.fits','Nu2SumPoisDatonPaddedIntTime73.0.060529.fits','Nu2SumPoisDatonPaddedIntTime242.0.060529.fits','Nu2SumPoisDatonPaddedIntTime539.0.060529.fits','Nu2SumPoisDatonPaddedIntTime806.0.060529.fits']



# Quick-o check of the other lists:
if ( ( len(IntTimeLst) != len(ExpFileLst) ) or ( len(IntTimeLst) != len(Nu2CompDataFileLst) ) or ( len(IntTimeLst) != len(Nu2CompDatonsFileLst) )  ):
    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
    nu2compdata_hdu = simpler_outer_products( Nu2CompFluxLst[kk], this_exp_hdulst[0] )
    nu2compdata_HDU = pyfits.PrimaryHDU(data=nu2compdata_hdu.data,header=nu2compdata_hdu.header)
    nu2compdata_HDU.writeto(Nu2CompDataFileLst[kk])
    this_exp_hdulst.close()
#
#   Get nu2 Poisson Datons with the nu2 component:
    nu2compdatons_hdu = simpler_poisson_datons(nu2compdata_HDU)
#
#   Write out fits and txt files for each Nu2Comp
    print 'Sample data012 of output from outerprod:', nu2compdatons_hdu.data[0][1][2]
    print 'Header of output from outerprod:', nu2compdatons_hdu.header,'\n'
    nu2compdatons_HDU = pyfits.PrimaryHDU(data=nu2compdatons_hdu.data,header=nu2compdatons_hdu.header)
    nu2compdatons_HDU.writeto(Nu2CompDatonsFileLst[kk])


#
#
# For each array, SUM the nu2 piece with the old one to get Nu2PoissDatons* .
#   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
    Nu2PoisDatons_hdu = simpler_twosum(olddatons_HDU[0],nu2compdatons_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
    Nu2PoisDatons_HDU = pyfits.PrimaryHDU(data=Nu2PoisDatons_hdu.data)
    Nu2PoisDatons_HDU.writeto(Nu2PoisDatonsFileLst[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