#!/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]))
        print loc
        print binLo[ibin], binHi[ibin]
        binMap = numpy.nan_to_num(p2d.powerMap.copy())*0.
        print binMap.shape, weightMap.shape, modIntLMap.shape
        binMap[loc] = weightMap[loc]
        binnedPower[ibin] = numpy.sum(numpy.nan_to_num(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['nullDataDir']+'/'
lensedDir = p['lensedDataDir']+'/'

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'))
mapList = p['mapList']

tag = p['tag']
indicesTag = p['indicesTag']
binFile = p['binFile']
weighting = p['weighting']
numPhases = p['numPhases']


print "phase randomization per SI %d"%numPhases


# calculates mean power spectrum for each map


for mapNo in xrange(len(mapList)):

    theMap = mapList[mapNo]

    print mapList[mapNo]

    for variable in xrange(numPhases):
        if variable == 0:
            phasedMap = pickle.load(open(nullDir+rootName+'_'+theMap+'_'+indicesTag))
            phasedMapPowerTotal = powerFromKMap(phasedMap,phasedMap,p2d,ftex)
        else:
            phasedMap = pickle.load(open(nullDir+rootName+"%03d"%variable+'_'+theMap+'_'+indicesTag))
            phasedMapPower = powerFromKMap(phasedMap,phasedMap,p2d,ftex)
            phasedMapPowerTotal.powerMap += phasedMapPower.powerMap
    phasedMapPowerTotal.powerMap /= numPhases
    divMaps1 = pickle.load(open(lensedDir+rootName+'_'+theMap+'_'+indicesTag))
    divMaps1 = numpy.array(divMaps1)

    crossPower1Mean = powerFromKMap(divMaps1,divMaps1,p2d,ftex)        
    
    patchPower = crossPower1Mean.powerMap - phasedMapPowerTotal.powerMap
    patchPowerSquared = patchPower**2.
    
    weightMap = patchPower*0. + 1.0
    crossPower1Mean.powerMap -= phasedMapPowerTotal.powerMap

    print weightMap.shape, crossPower1Mean.powerMap.shape
    lL,lU,lBin,clBin,junk = weightedBinInAnnuli(crossPower1Mean,weightMap,binFile,trimAtL,0)


    print mapList, mapNo
    
    pickle.dump(patchPower,open('pPowerData'+mapList[mapNo],'w'))
    
    pickle.dump([lBin,clBin],open('binnedPowerData'+mapList[mapNo],'w'))
    
