import pyfits
import healpy as hp
import numpy as np
import time
import pymp
from math import sqrt

start_time = time.time()
print('Starting')

threads=16
lmax=2002
lmax=int(lmax+1)

########## IMPORTING DATA ##########
alm=pyfits.open('alm.fits')
almimag=np.array(alm[1].data['IMAG'])
almreal=np.array(alm[1].data['REAL'])
almreal=list(almreal)
almimag=list(almimag)
cl=np.array((pyfits.open('cl.fits'))[1].data['TEMPERATURE'],dtype='float64')
#mastermatrix=1;
#gaussianbeam
#noise
#<f>ab,<f>dopp e <f>abdopp==<f>ab+<f>dopp (import or generate)

########## DEFINING FUNCTIONS ##########

def almf( l, m ):
    ellement=almlist[(l*(l+1))+m]
    return ellement;
def almfconjugate( l, m ):
    ellement=almlistconjugate[(l*(l+1))+m]
    return ellement;
def cps0factor( l, m):
    return sqrt(((l+1+m)*(l+1-m))/((4*(l+1)*(l+1)-1)))
def cms0factor( l, m):
    return -1*sqrt(((l+m)*(l-m))/(4*l*l-1))
def cps1factor( l, m):
    return -1*sqrt(((l+2+m)*(l+m+1))/((4*(l+1)*(l+1)-1)*(2)))
def cms1factor( l, m):
    return -1*sqrt(((l+m-1)*(l+m))/((4*l*l-1)*(2)))

########## CALCULATING c+ AND c- ##########

cps0abfactorlist=pymp.shared.array((hp.Alm.getsize(lmax=lmaxalm)*2-lmaxalm,), dtype='float64')
cms0abfactorlist=pymp.shared.array((hp.Alm.getsize(lmax=lmaxalm)*2-lmaxalm,), dtype='float64')
cps0doppfactorlist=pymp.shared.array((hp.Alm.getsize(lmax=lmaxalm)*2-lmaxalm,), dtype='float64')
cms0doppfactorlist=pymp.shared.array((hp.Alm.getsize(lmax=lmaxalm)*2-lmaxalm,), dtype='float64')
with pymp.Parallel(threads) as p:
    for l in p.range(lmaxalm+1):
        idxl=l*(l+1)
        for m in range(-l,l+1):
            idxlm=idxl+m
            cps0=cps0factor( l, m)
            cms0=cms0factor( l, m)
            cps0abfactorlist[idxlm]=(l+2)*cps0
            cms0abfactorlist[idxlm]=(l-1)*cms0
            cps0doppfactorlist[idxlm]=-cps0
            cms0doppfactorlist[idxlm]=cms0

elapsed_time = time.time() - start_time
print('cp and cm generated')
print(elapsed_time,'secs')

########## CALCULATING <f> AND <g> (h) ##########

symarraysize=hp.Alm.getsize(lmax=lmax)*2-lmax-2
notsymarraysize=hp.Alm.getsize(lmax=lmax)

hs0ablist=pymp.shared.array((hp.Alm.getsize(lmax=lmax),), dtype='float64')
hs0dopplist=pymp.shared.array((hp.Alm.getsize(lmax=lmax),), dtype='float64')
hs02ablist=pymp.shared.array((hp.Alm.getsize(lmax=lmax),), dtype='float64')
hs02dopplist=pymp.shared.array((hp.Alm.getsize(lmax=lmax),), dtype='float64')

with pymp.Parallel(threads) as p:
    for l in p.range(lmax+1):
        idxl=l*(l+1);idxlp1=(l+1)*(l+2);
        for m in range(1,l):
            idxlm=idxl+m;idxlp1m=idxlp1+m;
            abcoef=(cms0abfactorlist[idxlp1m]*cl[l]+cps0abfactorlist[idxlm]*cl[l+1])
            doppcoef=(cms0doppfactorlist[idxlp1m]*cl[l]+cps0doppfactorlist[idxlm]*cl[l+1])
            hs0ablist[int((idxl/2)+m)]=abcoef
            hs02ablist[int((idxl/2)+m)]=2*abcoef
            hs0dopplist[int((idxl/2)+m)]=doppcoef
            hs02dopplist[int((idxl/2)+m)]=2*doppcoef
        for m in [l]:
            idxlm=idxl+m;idxlp1m=idxlp1+m;
            abcoef=(cms0abfactorlist[idxlp1m]*cl[l]+cps0abfactorlist[idxlm]*cl[l+1])
            doppcoef=(cms0doppfactorlist[idxlp1m]*cl[l]+cps0doppfactorlist[idxlm]*cl[l+1])
            hs0ablist[int((idxl/2)+m)]=abcoef
            hs02ablist[int((idxl/2)+m)]=2*abcoef
            hs0dopplist[int((idxl/2)+m)]=doppcoef
            hs02dopplist[int((idxl/2)+m)]=2*doppcoef
        for m in [0]:
            idxlm=idxl+m;idxlp1m=idxlp1+m;
            hs0ablist[int((idxl/2)+m)]=(cms0abfactorlist[idxlp1m]*cl[l]+cps0abfactorlist[idxlm]*cl[l+1])
            hs02ablist[int((idxl/2)+m)]=(cms0abfactorlist[idxlp1m]*cl[l]+cps0abfactorlist[idxlm]*cl[l+1])
            hs0dopplist[int((idxl/2)+m)]=(cms0doppfactorlist[idxlp1m]*cl[l]+cps0doppfactorlist[idxlm]*cl[l+1])
            hs02dopplist[int((idxl/2)+m)]=(cms0doppfactorlist[idxlp1m]*cl[l]+cps0doppfactorlist[idxlm]*cl[l+1])