#!/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 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(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']
nMCOld = p['nMOld']
rootName = p['rootName']
examplePatch = p['examplePatch']
trimAtL = p['trimAtL']
exampleLiteMap = liteMap.liteMapFromFits(examplePatch)
p2d = fftTools.powerFromLiteMap(exampleLiteMap)
ftex = p2d.trimAtL(trimAtL)
#indicesList = pickle.load(open(nullDir+'indicesList'+'000'+'0','r'))
numDivs = p['numDivs']
mapList = p['mapList']
tag = p['tag']
indicesTag = p['indicesTag']
binFile = p['binFile']
weighting = p['weighting']
numPhases = p['numPhases']
oldNullDir = p['oldNullDir']
print 'computing null kappa maps'

allDivsList = []
for theMap in mapList:
    for q in xrange(nMCOld):
        print q
        divMaps1 = pickle.load(open(oldNullDir+rootName+"%03d"%q+'_'+theMap+'_'+indicesTag))
        divMaps1 = numpy.array(divMaps1)
        if q == 0:
            allDivs1 = divMaps1
        else:
            allDivs1 += divMaps1
    allDivs1 /= nMCs
    allDivsList += [allDivs1]
pickle.dump(allDivsList,open('allDivsList_'+tag+'.pkl','w'))

print 'finished null maps, now computing lensed'

patchPowerList = []
patchPowerSquaredList = []

# calculates mean power spectrum for each map
othercounter = 0
for theMap in mapList:
    for p in xrange(nMCs):
        print theMap, p, nMCs
        for variable in xrange(numPhases):
            if variable == 0:
                phasedMap = pickle.load(open(nullDir+rootName+"%03d"%p+'_'+theMap+'_'+indicesTag))
                phasedMapPowerTotal = powerFromKMap(phasedMap,phasedMap,p2d,ftex)
            else:
                phasedMap = pickle.load(open(nullDir+rootName+"%03d"%p+"%03d"%variable+'_'+theMap+'_'+indicesTag))
                phasedMapPower = powerFromKMap(phasedMap,phasedMap,p2d,ftex)
                phasedMapPowerTotal.powerMap += phasedMapPower.powerMap
        phasedMapPowerTotal.powerMap /= numPhases
        divMaps1 = pickle.load(open(lensedDir+rootName+"%03d"%p+'_'+theMap+'_'+indicesTag))
        divMaps1 = numpy.array(divMaps1)
        #divMaps1 -= allDivsList[othercounter]   
        crossPower1Mean = powerFromKMap(divMaps1,divMaps1,p2d,ftex)             
        if p == 0:
            patchPower = crossPower1Mean.powerMap - phasedMapPowerTotal.powerMap
            patchPowerSquared = patchPower**2.
            #justPhase = phasedMapPowerTotal.powerMap
        else:
            patchPower += crossPower1Mean.powerMap - phasedMapPowerTotal.powerMap
            patchPowerSquared += crossPower1Mean.powerMap**2.
            #justPhase += phasedMapPowerTotal.powerMap
    patchPower /= nMCs
    patchPowerSquared /= nMCs
    patchPowerList += [patchPower]
    patchPowerSquaredList += [patchPowerSquared]
    othercounter += 1

print 'now unlensed'            


print 'calculating scatter and mean, binning'
meanVar = patchPowerSquaredList
for bla in xrange(len(patchPowerList)):
    meanPowerSquared = patchPowerList[bla]**2.
    meanVar[bla] -= meanPowerSquared

clList = []
scatterList = []
for bla in xrange(len(patchPowerList)):
    scatter = numpy.sqrt(meanVar[bla])
    #pickle.dump(scatter,open('scatterFile_'+tag+'.pkl','w'))
    weightMap = 1./scatter**2.
    weightMap[0,0] = weightMap[1,0]
    loc = numpy.where(crossPower1Mean.modLMap ==0.)
    print 'loc', loc, crossPower1Mean.modLMap[loc]
    weightMap[loc] = 0.
    if weighting:
        pass
    else:
        weightMap = weightMap*0. + 1.

########### now weighted bin lensed and unlensed
    crossPowerMean = crossPower1Mean.copy()
    crossPowerMean.powerMap = patchPowerList[bla]
    crossPowerMean.powerMap[0,0] = crossPowerMean.powerMap[0,1]
    lL,lU,lBin,meanCls,weights = weightedBinInAnnuli(crossPowerMean,weightMap,binFile,trimAtL,0)
    clList += [meanCls]
    
    weightMapNew = weightMap*0. + 1.

    scatterPlot = crossPower1Mean.copy()
    scatterPlot.powerMap = 1./scatter**2.
    scatterPlot.powerMap[0,0] = scatterPlot.powerMap[0,1]
    lL,lU,lBin,newScatter,weights = weightedBinInAnnuli(scatterPlot,weightMapNew,binFile,trimAtL,0)
    #pickle.dump(numpy.sqrt(1./newScatter),open('newScatter_'+tag+'.pkl','w'))

    scatterList += [numpy.sqrt(1./newScatter)]

###### calculate optimally weighted scatter
optCls = lBin
weighting = lBin
for bla in xrange(len(patchPowerList)):
    if bla == 0:
        optCls = clList[bla]
        optCls *= 1./scatterList[bla]**2.
        sum = 1./scatterList[bla]**2.
    else:
        optCls += clList[bla]
        optCls *= 1./scatterList[bla]**2.
        sum += 1./scatterList[bla]**2.
optCls /= sum
finalScatter = numpy.sqrt(1./sum)

pickle.dump(optCls,open('cls.pkl','w'))
pickle.dump(finalScatter,open('scatter.pkl','w'))
