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


from flipper import *
from numpy.fft import fftshift,fftfreq,fft2,ifft2
from scipy import interpolate
from scipy import *
import scipy
import os
import random
import sys
import pickle
from lmcTools import *
import lmcTools

strList = ['allall']
indicesList = [[0,1],[0,2],[0,3],[1,2],[1,3],[2,3]]
count = 0

def weightedBinInAnnuli(p2d,weightMap,binningFile,trimAtL,powerOfL):
    binLo,binHi,binCent = fftTools.readBinningFile(binningFile)
    id = numpy.where(binHi<trimAtL)
    binHi = binHi[id]
    binLo = binLo[id]
    binCent = binCent[id]
    binnedPower = binCent.copy()*0.
    binCount = binCent.copy()*0.
    weightedBincount = binCent.copy()
    modIntLMap = numpy.array(p2d.modLMap + 0.5,dtype='int64')
    for ibin in xrange(len(binHi)):
        loc = numpy.where((modIntLMap >= binLo[ibin]) & (modIntLMap <= binHi[ibin]))
        binMap = p2d.powerMap.copy()*0.
        binMap[loc] = weightMap[loc]
        binnedPower[ibin] = numpy.sum(p2d.powerMap*binMap*p2d.modLMap**powerOfL)/numpy.sum(binMap)
        binCount[ibin] = len(loc[0])
        weightedBincount[ibin] = 1./(numpy.sum(weightMap[loc]**2)/(numpy.sum(weightMap[loc]))**2)
    #print binCount[ibin]/weightedBincount
    return binLo,binHi,binCent,binnedPower, weightedBincount/2.
                                                                                                                                    

def powerFromKMap(map1,map2,p2db,p2d2):
    power = p2db.copy()
    power.powerMap = numpy.real(numpy.conjugate(map2)*map1)

    return power

p = flipperDict.flipperDict()
p.read_from_file(sys.argv[1])
mcDir = p['mcDir']
dataDir = p['dataDir']

weightMap = pickle.load(open(mcDir+'/scatterFile'))
weightMap = 1./weightMap.powerMap*2.
weightMap[0,0] = weightMap[0,1]


nullMap4 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap4.pkl'))
nullMap5 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap5.pkl'))
nullMap6 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap6.pkl'))
nullMap7 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap7.pkl'))
nullMap8 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap8.pkl'))
nullMap9 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap9.pkl'))
nullMap10 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap10.pkl'))
nullMap11 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap11.pkl'))
nullMap12 = 1.*pickle.load(open(mcDir+'/averageNullKappaMap12.pkl'))

for i in xrange(1):
    holder = liteMap.liteMapFromFits('/scratch/sudeep/lensMC/randomMapsRealNoiseBoth/lensedMaps_CMBplusNoise_071/patch_148_000_11')
    p2d = fftTools.powerFromLiteMap(holder)
    p2d = p2d.trimAtL(3300)
    ftex = p2d.copy()
    for j in xrange(1):
        difference = 1
        if difference:
            print i,j,indicesList[i],indicesList[j], difference
            crossPower4Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_004_'+strList[i]))-nullMap1[i],pickle.load(open(dataDir+'/saveKappas_004_'+strList[j]))-nullMap1[j],p2d,ftex)
        
            crossPower5Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_005_'+strList[i]))-nullMap2[i],pickle.load(open(dataDir+'/saveKappas_005_'+strList[j]))-nullMap2[j],p2d,ftex)
            crossPower6Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_006_'+strList[i]))-nullMap3[i],pickle.load(open(dataDir+'/saveKappas_006_'+strList[j]))-nullMap3[j],p2d,ftex)
            crossPower7Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_007_'+strList[i]))-nullMap4[i],pickle.load(open(dataDir+'/saveKappas_007_'+strList[j]))-nullMap4[j],p2d,ftex)                                                                        
            crossPower8Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_008_'+strList[i]))-nullMap1[i],pickle.load(open(dataDir+'/saveKappas_008_'+strList[j]))-nullMap1[j],p2d,ftex)
        
            crossPower9Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_009_'+strList[i]))-nullMap2[i],pickle.load(open(dataDir+'/saveKappas_009_'+strList[j]))-nullMap2[j],p2d,ftex)
            crossPower10Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_010_'+strList[i]))-nullMap3[i],pickle.load(open(dataDir+'/saveKappas_010_'+strList[j]))-nullMap3[j],p2d,ftex)
            crossPower11Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_011_'+strList[i]))-nullMap4[i],pickle.load(open(dataDir+'/saveKappas_011_'+strList[j]))-nullMap4[j],p2d,ftex)                                                                        
            crossPower12Old = powerFromKMap(pickle.load(open(dataDir+'/saveKappas_012_'+strList[i]))-nullMap4[i],pickle.load(open(dataDir+'/saveKappas_012_'+strList[j]))-nullMap4[j],p2d,ftex)                                                                        

            pickle.dump(pickle.load(open(dataDir+'/saveKappas_004_'+strList[i]))-nullMap4[i],open('kappa4.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_005_'+strList[i]))-nullMap5[i],open('kappa5.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_006_'+strList[i]))-nullMap6[i],open('kappa6.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_007_'+strList[i]))-nullMap7[i],open('kappa7.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_008_'+strList[i]))-nullMap8[i],open('kappa8.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_009_'+strList[i]))-nullMap9[i],open('kappa9.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_010_'+strList[i]))-nullMap10[i],open('kappa10.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_011_'+strList[i]))-nullMap11[i],open('kappa11.pkl','w'))
            pickle.dump(pickle.load(open(dataDir+'/saveKappas_012_'+strList[i]))-nullMap12[i],open('kappa12.pkl','w'))
             
            if count == 0:
                crossPower4 = crossPower4Old.copy()
                crossPower5 = crossPower5Old.copy()
                crossPower6 = crossPower6Old.copy()
                crossPower7 = crossPower7Old.copy()

                crossPower8 = crossPower8Old.copy()
                crossPower9 = crossPower9Old.copy()
                crossPower10 = crossPower10Old.copy()
                crossPower11 = crossPower11Old.copy()

                crossPower12 = crossPower12Old.copy()

            count +=1

averagePower = powerFromKMap(nullMap2[0],nullMap2[0],p2d,ftex)
averagePower.powerMap = (crossPower4.powerMap+crossPower5.powerMap+crossPower6.powerMap+crossPower7.powerMap)

averagePower.powerMap += (crossPower8.powerMap+crossPower9.powerMap+crossPower10.powerMap+crossPower11.powerMap+crossPower12.powerMap)
averagePower.powerMap /= 9.
averagePower.powerMap[0,0] = averagePower.powerMap[0,1]

lL,lU,lBin,clBin1,weights = weightedBinInAnnuli(averagePower,weightMap,'BIN_LENS',3300,0)
print lBin, clBin1

pickle.dump(clBin1,open('clData_south.pkl','w'))
pickle.dump(lBin,open('lBinData.pkl','w'))
averagePower.powerMap = weightMap
bla,bla,l,ww,bla,bla = averagePower.binInAnnuli('BIN_LENS')
bl
X = numpy.loadtxt('/project/rbond/sudeep/theoryFiles/bode_almost_wmap5_lmax_1e4\
_scalCls.dat')
tcmb = 2.726e6
clkk = X[:,4]/(4.*tcmb**2)
l = X[:,0]

tag = 'New'
clBinUn = pickle.load(open(mcDir+"/meanClsUn.pkl"))
clBinSigma = pickle.load(open(mcDir+"/scatter.pkl"))
fac = 1.
pylab.plot(lBin,clBin1,"o",label='data')
pylab.plot(lBin,clBinUn*fac,label='unlensed')
pylab.errorbar(lBin,(clBin1 -clBinUn*fac),clBinSigma,fmt="o",label='data-unlensed')
pylab.plot(l,clkk)
pylab.legend(loc='best')
g = pylab.gca()
g.set_yscale("log")
g.set_xscale("log")
pylab.xlim(10,3300)
pylab.ylim(1e-10,1e-4)
pylab.savefig("cls%s.png"%tag)

