#!/usr/bin/env python
#Given N-maps calculates cross-reconstruction between them

from flipper import *
#from lensTools  import *
from numpy.fft import fftshift,fftfreq,fft2,ifft2
from scipy import interpolate
from scipy import *
import os
import random
import pickle
import numpy
from fftTools import *

def powerFromKMap(k1,k2,ft,ftret):
    """
    @brief Creates a power2D object from one power 2D object and two kmaps
    """
    p2d = ftret.copy()
    mapFFT = k1.copy()
    mapFFT2 = k2.copy()

    area =ft.Nx*ft.Ny*ft.pixScaleX*ft.pixScaleY
    p2d.powerMap = numpy.real(numpy.conjugate(mapFFT)*mapFFT2)#*area/(ft.Nx*ft.Ny*1.0)**2
    
    return p2d

p = flipperDict.flipperDict()
p.read_from_file(sys.argv[1])
nullDir = p['nullDir']
lensedDir = p['lensedDir']
nMCs = p['nMCs']
rootName = p['rootName']
examplePatch = p['examplePatch']
trimAtL = p['trimAtL']
binningFile = p['binningFile']
if binningFile == None:
    binningFile = 'BIN_100_LOG'
    print "binning file set to %s"%binningFile


exampleLiteMap = liteMap.liteMapFromFits(examplePatch)
p2d = fftTools.powerFromLiteMap(exampleLiteMap)
ftex = p2d.trimAtL(trimAtL)
#indicesList = pickle.load(open(nullDir+'indicesList'+'000'+'0','r'))
numDivs = p['numDivs']

indicesList = []
indList = []
for i in xrange(numDivs):
    for j in xrange(i):
        indicesList += ["%d%d"%(j,i)]
        
        indList += [[i,j]] 
print indicesList
print indList

# gets mean kappa from mcs
for q in xrange(nMCs):
    #for i in xrange(numDivs):
    divMaps1 = []#numpy.arange(len(indicesList))
    divMaps2 = []#numpy.arange(len(indicesList))
    divMaps3 = []#numpy.arange(len(indicesList))
    divMaps4 = []#numpy.arange(len(indicesList))
    divMaps5 = []
    divMaps6 = []#numpy.arange(len(indicesList))
    divMaps7 = []#numpy.arange(len(indicesList))
    divMaps8 = []#numpy.arange(len(indicesList))
    divMaps9 = []#numpy.arange(len(indicesList))
                
    
    count = 0
    for ind in indicesList:
        divMaps1 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_004'+'_'+ind))]
        divMaps2 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_005'+'_'+ind))]
        divMaps3 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_006'+'_'+ind))]
        divMaps4 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_007'+'_'+ind))]
        divMaps5 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_008'+'_'+ind))]
        divMaps6 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_009'+'_'+ind))]
        print divMaps6
        divMaps7 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_010'+'_'+ind))]
        divMaps8 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_011'+'_'+ind))]
        divMaps9 += [pickle.load(open(nullDir+rootName+"%03d"%q+'_012'+'_'+ind))]
        count += 1
    divMaps1 = numpy.array(divMaps1)
    divMaps2 = numpy.array(divMaps2)
    divMaps3 = numpy.array(divMaps3)
    divMaps4 = numpy.array(divMaps4)
    divMaps5 = numpy.array(divMaps5)
    divMaps6 = numpy.array(divMaps6)
    divMaps7 = numpy.array(divMaps7)
    divMaps8 = numpy.array(divMaps8)
    divMaps9 = numpy.array(divMaps9)
    if q == 0:
        allDivs1 = divMaps1
        allDivs2 = divMaps2
        allDivs3 = divMaps3
        allDivs4 = divMaps4
        allDivs5 = divMaps5
        allDivs6 = divMaps6
        allDivs7 = divMaps7
        allDivs8 = divMaps8
        allDivs9 = divMaps9
    else:
        allDivs1[:] += divMaps1[:]
        allDivs2[:] += divMaps2[:]
        allDivs3[:] += divMaps3[:]
        allDivs4[:] += divMaps4[:]
        allDivs5[:] += divMaps5[:]
        allDivs6[:] += divMaps6[:]
        allDivs7[:] += divMaps7[:]
        allDivs8[:] += divMaps8[:]
        allDivs9[:] += divMaps9[:]


allDivs1[:] /= nMCs
allDivs2[:] /= nMCs
allDivs3[:] /= nMCs
allDivs4[:] /= nMCs
allDivs5[:] /= nMCs
allDivs6[:] /= nMCs
allDivs7[:] /= nMCs
allDivs8[:] /= nMCs
allDivs9[:] /= nMCs


pickle.dump(allDivs1,open('averageNullKappaMap004.pkl','w'))
pickle.dump(allDivs2,open('averageNullKappaMap005.pkl','w'))
pickle.dump(allDivs3,open('averageNullKappaMap006.pkl','w'))
pickle.dump(allDivs4,open('averageNullKappaMap007.pkl','w'))
pickle.dump(allDivs5,open('averageNullKappaMap008.pkl','w'))
pickle.dump(allDivs6,open('averageNullKappaMap009.pkl','w'))
pickle.dump(allDivs7,open('averageNullKappaMap010.pkl','w'))
pickle.dump(allDivs8,open('averageNullKappaMap011.pkl','w'))
pickle.dump(allDivs9,open('averageNullKappaMap012.pkl','w'))

patch1Power = []
patch2Power = []
patch3Power = []
patch4Power = []
patch5Power = []
patch6Power = []
patch7Power = []
patch8Power = []
patch9Power = []

# calculates mean power spectrum for each map
for p in xrange(nMCs):
    divMaps1 = []
    divMaps2 = []
    divMaps3 = []
    divMaps4 = []
    divMaps5 = []
    divMaps6 = []
    divMaps7 = []
    divMaps8 = []
    divMaps9 = []
    count = 0
    for ind in indicesList:
        divMaps1 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_004'+'_'+ind))]
        divMaps2 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_005'+'_'+ind))]
        divMaps3 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_006'+'_'+ind))]
        divMaps4 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_007'+'_'+ind))]
        divMaps5 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_008'+'_'+ind))]
        divMaps6 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_009'+'_'+ind))]
        divMaps7 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_010'+'_'+ind))]
        divMaps8 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_011'+'_'+ind))]
        divMaps9 += [pickle.load(open(lensedDir+rootName+"%03d"%p+'_012'+'_'+ind))]
        count += 1

    divMaps1 = numpy.array(divMaps1)
    divMaps2 = numpy.array(divMaps2)
    divMaps3 = numpy.array(divMaps3)
    divMaps4 = numpy.array(divMaps4)
    divMaps5 = numpy.array(divMaps5)
    divMaps6 = numpy.array(divMaps6)
    divMaps7 = numpy.array(divMaps7)
    divMaps8 = numpy.array(divMaps8)
    divMaps9 = numpy.array(divMaps9)

  #  divMaps1[:] -= allDivs1[:]
  #  divMaps2[:] -= allDivs2[:]
  #  divMaps3[:] -= allDivs3[:]
  #  divMaps4[:] -= allDivs4[:]
  #  divMaps5[:] -= allDivs5[:]
  #  divMaps6[:] -= allDivs6[:]
  #  divMaps7[:] -= allDivs7[:]
  #  divMaps8[:] -= allDivs8[:]
  #  divMaps9[:] -= allDivs9[:]
    print 'banannaaas'
    nDivs = len(divMaps1)
    count = 0
    for i in xrange(nDivs):
        for j in xrange(i):
            difference = 1
            for k in xrange(2):
                for lll in xrange(2):
                    if (indList[i][k] == indList[j][lll]):
                        difference = 0
            if difference:
                print i,j,indList[i],indList[j], difference
                crossPower1Old = powerFromKMap(divMaps1[i],divMaps1[j],p2d,ftex)
                crossPower2Old = powerFromKMap(divMaps2[i],divMaps2[j],p2d,ftex)                
                crossPower3Old = powerFromKMap(divMaps3[i],divMaps3[j],p2d,ftex)
                crossPower4Old = powerFromKMap(divMaps4[i],divMaps4[j],p2d,ftex)
                crossPower5Old = powerFromKMap(divMaps5[i],divMaps5[j],p2d,ftex)
                crossPower6Old = powerFromKMap(divMaps6[i],divMaps6[j],p2d,ftex)                
                crossPower7Old = powerFromKMap(divMaps7[i],divMaps7[j],p2d,ftex)
                crossPower8Old = powerFromKMap(divMaps8[i],divMaps8[j],p2d,ftex)
                crossPower9Old = powerFromKMap(divMaps9[i],divMaps9[j],p2d,ftex)
                if count == 0:
                    crossPower1 = crossPower1Old.copy()
                    crossPower2 = crossPower2Old.copy()
                    crossPower3 = crossPower3Old.copy()
                    crossPower4 = crossPower4Old.copy()
		    crossPower5 = crossPower5Old.copy()
                    crossPower6 = crossPower6Old.copy()
                    crossPower7 = crossPower7Old.copy()
                    crossPower8 = crossPower8Old.copy()
                    crossPower9 = crossPower9Old.copy()
                else:
                    crossPower1.powerMap[:] += crossPower1Old.powerMap[:]
                    crossPower2.powerMap[:] += crossPower2Old.powerMap[:]
                    crossPower3.powerMap[:] += crossPower3Old.powerMap[:]
                    crossPower4.powerMap[:] += crossPower4Old.powerMap[:]
                    crossPower5.powerMap[:] += crossPower5Old.powerMap[:]
                    crossPower6.powerMap[:] += crossPower6Old.powerMap[:]
                    crossPower7.powerMap[:] += crossPower7Old.powerMap[:]
                    crossPower8.powerMap[:] += crossPower8Old.powerMap[:]
                    crossPower9.powerMap[:] += crossPower9Old.powerMap[:]
                    print 'adding power, n=', count
            
                count +=1
    crossPower1Mean = crossPower1.copy()
    crossPower2Mean = crossPower2.copy()
    crossPower3Mean = crossPower3.copy()
    crossPower4Mean = crossPower4.copy()
    crossPower5Mean = crossPower5.copy()
    crossPower6Mean = crossPower6.copy()
    crossPower7Mean = crossPower7.copy()
    crossPower8Mean = crossPower8.copy()
    crossPower9Mean = crossPower9.copy()

    crossPower1Mean.powerMap[:] /= count
    crossPower2Mean.powerMap[:] /= count
    crossPower3Mean.powerMap[:] /= count
    crossPower4Mean.powerMap[:] /= count
    crossPower5Mean.powerMap[:] /= count
    crossPower6Mean.powerMap[:] /= count
    crossPower7Mean.powerMap[:] /= count
    crossPower8Mean.powerMap[:] /= count
    crossPower9Mean.powerMap[:] /= count
    lL,lU,lBin,clBin1, plErr_12,weights = crossPower1Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin2, plErr_12,weights = crossPower2Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin3, plErr_12,weights = crossPower3Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin4, plErr_12,weights = crossPower4Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin5, plErr_12,weights = crossPower5Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin6, plErr_12,weights = crossPower6Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin7, plErr_12,weights = crossPower7Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin8, plErr_12,weights = crossPower8Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin9, plErr_12,weights = crossPower9Mean.binInAnnuli(binningFile)
    patch1Power += [clBin1]
    patch2Power += [clBin2]
    patch3Power += [clBin3]
    patch4Power += [clBin4]
    patch5Power += [clBin5]
    patch6Power += [clBin6]
    patch7Power += [clBin7]
    patch8Power += [clBin8]
    patch9Power += [clBin9]


patch1PowerUn = []
patch2PowerUn = []
patch3PowerUn = []
patch4PowerUn = []
patch5PowerUn = []
patch6PowerUn = []
patch7PowerUn = []
patch8PowerUn = []
patch9PowerUn = []

for p in xrange(nMCs):
    divMaps1 = []
    divMaps2 = []
    divMaps3 = []
    divMaps4 = []
    divMaps5 = []
    divMaps6 = []
    divMaps7 = []
    divMaps8 = []
    divMaps9 = []
    count = 0
    for ind in indicesList:
        divMaps1 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_004'+'_'+ind))]
        divMaps2 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_005'+'_'+ind))]
        divMaps3 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_006'+'_'+ind))]
        divMaps4 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_007'+'_'+ind))]
        divMaps5 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_008'+'_'+ind))]
        divMaps6 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_009'+'_'+ind))]
        divMaps7 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_010'+'_'+ind))]
        divMaps8 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_011'+'_'+ind))]
        divMaps9 += [pickle.load(open(nullDir+rootName+"%03d"%p+'_012'+'_'+ind))]
        count += 1

    divMaps1 = numpy.array(divMaps1)
    divMaps2 = numpy.array(divMaps2)
    divMaps3 = numpy.array(divMaps3)
    divMaps4 = numpy.array(divMaps4)
    divMaps5 = numpy.array(divMaps5)
    divMaps6 = numpy.array(divMaps6)
    divMaps7 = numpy.array(divMaps7)
    divMaps8 = numpy.array(divMaps8)
    divMaps9 = numpy.array(divMaps9)

    divMaps1[:] -= allDivs1[:]
    divMaps2[:] -= allDivs2[:]
    divMaps3[:] -= allDivs3[:]
    divMaps4[:] -= allDivs4[:]
    divMaps5[:] -= allDivs5[:]
    divMaps6[:] -= allDivs6[:]
    divMaps7[:] -= allDivs7[:]
    divMaps8[:] -= allDivs8[:]
    divMaps9[:] -= allDivs9[:]
    nDivs = len(divMaps1)
    count = 0
    for i in xrange(nDivs):
        for j in xrange(i):
            difference = 1
            for k in xrange(2):
                for lll in xrange(2):
                    if (indicesList[i][k] == indicesList[j][lll]):
                        difference = 0
            if difference:
                print i,j,indicesList[i],indicesList[j], difference
                crossPower1Old = powerFromKMap(divMaps1[i],divMaps1[j],p2d,ftex)
                crossPower2Old = powerFromKMap(divMaps2[i],divMaps2[j],p2d,ftex)                
                crossPower3Old = powerFromKMap(divMaps3[i],divMaps3[j],p2d,ftex)
                crossPower4Old = powerFromKMap(divMaps4[i],divMaps4[j],p2d,ftex)
                crossPower5Old = powerFromKMap(divMaps5[i],divMaps5[j],p2d,ftex)
                crossPower6Old = powerFromKMap(divMaps6[i],divMaps6[j],p2d,ftex)                
                crossPower7Old = powerFromKMap(divMaps7[i],divMaps7[j],p2d,ftex)
                crossPower8Old = powerFromKMap(divMaps8[i],divMaps8[j],p2d,ftex)
                crossPower9Old = powerFromKMap(divMaps9[i],divMaps9[j],p2d,ftex)
                if count == 0:
                    crossPower1 = crossPower1Old.copy()
                    crossPower2 = crossPower2Old.copy()
                    crossPower3 = crossPower3Old.copy()
                    crossPower4 = crossPower4Old.copy()
                    crossPower5 = crossPower5Old.copy()
                    crossPower6 = crossPower6Old.copy()
                    crossPower7 = crossPower7Old.copy()
                    crossPower8 = crossPower8Old.copy()
                    crossPower9 = crossPower9Old.copy()
                else:
                    crossPower1.powerMap[:] += crossPower1Old.powerMap[:]
                    crossPower2.powerMap[:] += crossPower2Old.powerMap[:]
                    crossPower3.powerMap[:] += crossPower3Old.powerMap[:]
                    crossPower4.powerMap[:] += crossPower4Old.powerMap[:]
                    crossPower5.powerMap[:] += crossPower5Old.powerMap[:]
                    crossPower6.powerMap[:] += crossPower6Old.powerMap[:]
                    crossPower7.powerMap[:] += crossPower7Old.powerMap[:]
                    crossPower8.powerMap[:] += crossPower8Old.powerMap[:]
                    crossPower9.powerMap[:] += crossPower9Old.powerMap[:]
                    print 'adding power, n=', count
            
                count +=1
    crossPower1Mean = crossPower1.copy()
    crossPower2Mean = crossPower2.copy()
    crossPower3Mean = crossPower3.copy()
    crossPower4Mean = crossPower4.copy()
    crossPower5Mean = crossPower5.copy()
    crossPower6Mean = crossPower6.copy()
    crossPower7Mean = crossPower7.copy()
    crossPower8Mean = crossPower8.copy()
    crossPower9Mean = crossPower9.copy()

    crossPower1Mean.powerMap[:] /= count
    crossPower2Mean.powerMap[:] /= count
    crossPower3Mean.powerMap[:] /= count
    crossPower4Mean.powerMap[:] /= count
    crossPower5Mean.powerMap[:] /= count
    crossPower6Mean.powerMap[:] /= count
    crossPower7Mean.powerMap[:] /= count
    crossPower8Mean.powerMap[:] /= count
    crossPower9Mean.powerMap[:] /= count
    lL,lU,lBin,clBin1, plErr_12,weights = crossPower1Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin2, plErr_12,weights = crossPower2Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin3, plErr_12,weights = crossPower3Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin4, plErr_12,weights = crossPower4Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin5, plErr_12,weights = crossPower5Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin6, plErr_12,weights = crossPower6Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin7, plErr_12,weights = crossPower7Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin8, plErr_12,weights = crossPower8Mean.binInAnnuli(binningFile)
    lL,lU,lBin,clBin9, plErr_12,weights = crossPower9Mean.binInAnnuli(binningFile)
    patch1PowerUn += [clBin1]
    patch2PowerUn += [clBin2]
    patch3PowerUn += [clBin3]
    patch4PowerUn += [clBin4]
    patch5PowerUn += [clBin5]
    patch6PowerUn += [clBin6]
    patch7PowerUn += [clBin7]
    patch8PowerUn += [clBin8]
    patch9PowerUn += [clBin9]

patchPowerTotal = patch1Power
for var in xrange(len(patchPowerTotal)):
    patchPowerTotal[var] += patch2Power[var] + patch3Power[var] + patch4Power[var]
    patchPowerTotal[var] += patch5Power[var] + patch6Power[var] + patch7Power[var] + patch8Power[var]
    patchPowerTotal[var] /= 9.

patchPowerTotalUn = patch1PowerUn
for var in xrange(len(patchPowerTotalUn)):
    patchPowerTotalUn[var] += patch2PowerUn[var] + patch3PowerUn[var] + patch4PowerUn[var]
    patchPowerTotalUn[var] += patch5Power[var] + patch6Power[var] + patch7Power[var] + patch8Power[var]
    patchPowerTotalUn[var] /= 9.

# get mean power spectrum overall by averaging map power spectra
for r in xrange(nMCs):
    if r == 0:
        meanCls = patchPowerTotal[r].copy()
    else:
        meanCls += patchPowerTotal[r]
meanCls[:] /= nMCs

# scatter for each mean spec point
for r in xrange(nMCs):
    aCl = patchPowerTotal[r].copy()
    if r == 0:
        scatter = (aCl - meanCls)**2.
    else:
        scatter += (aCl - meanCls)**2.
scatter[:] = numpy.sqrt(scatter[:]/nMCs)

# calculate bias from unlensed mean spec
for r in xrange(nMCs):
    if r == 0:
        meanClsUn = patchPowerTotalUn[r].copy()
    else:
        meanClsUn += patchPowerTotalUn[r]
meanClsUn[:] /= nMCs

pickle.dump(lBin,open('lBin.pkl','w'))
pickle.dump(meanCls,open('meanClsEq.pkl','w'))
pickle.dump(meanClsUn,open('meanClsUnEq.pkl','w'))
pickle.dump(scatter,open('scatterEq.pkl','w'))


ymin = 1e-10

g = pylab.gca()
g.set_yscale('log',nonposy='clip')
g.set_xscale('log')
#pylab.plot(ells,kappaTrueAvg,label='true kappa')
#pylab.plot(ells,kappaRecMinusAvg+nullbias,label='recwithoutnullbias')
pylab.plot(lBin,meanCls,'o',label='lensed')
pylab.plot(lBin,meanClsUn,label='unlensed')
kappaRecMinusAvg = meanCls-meanClsUn
pylab.plot(lBin,kappaRecMinusAvg,label='reconswithoutbias')
#pylab.plot(ells,kappaBiasAvg,label='bias')
#pylab.plot(ells,kappaRecMinusAvg,'o',label='rec. kappa - bias')
kk = kappaRecMinusAvg.copy()


err = scatter.copy()#/numpy.sqrt(4.)
id = numpy.where((kappaRecMinusAvg-err) <0)
print id
errorlow = err.copy()
errorlow[id] = kappaRecMinusAvg[id]*0.99

pylab.errorbar(lBin,kappaRecMinusAvg,[errorlow,err],fmt="o",ecolor='k')

id = numpy.where(kk<0)
kk[id] = ymin
#pylab.errorbar(ells[id],kk[id],[err[id]*0,err[id]+kappaRecMinusAvg[id]-kk[id]],fmt=None,ecolor='k')

g.set_xlabel(r'$\ell$',fontsize=20)
g.set_ylabel(r'$C_\ell$',fontsize=20)
#pylab.ylim(1e-10,1e-4)
#pylab.xlim(30,3000)
g.set_xlim(49,3000)
#g.set_ylim(ymin,1e-3)
pylab.legend(loc='best')
pylab.savefig('masterPlot.png')
pylab.clf()
