##########################################################################
# module: swsh (spin weighted sperical harmonic)
# author: Milinda Fernando
# email:  milinda@cs.utah.edu
# python module to generate code for far-field energy ratiation extraction (Gravitational Waves).
# (c) 2016 University of Utah, All rights reserved.
##########################################################################


import dendro as dendro
from sympy import *
from sympy.tensor.array import *
from sympy.functions.special.tensor_functions import KroneckerDelta
from sympy.utilities import numbered_symbols
from sympy.printing import print_ccode
from sympy.printing.dot import dotprint
import math as math
import numpy as np
import sys as sys


# declare the variables to evaluate spin weigneted sperical harmonics.
theta=symbols('theta') #dendro.scalar("dtheta","[pp]")
phi=symbols('phi') #dendro.scalar("dphi","[pp]")
s,l,m,t,C1,C2=symbols('s,l,m,t,C1,C2')              
#sYlm=dendro.scalar("sYml","[pp]")  

C1=Max(0,m-s)
C2=Min(l+m,l-s)

# $d^{l}_{ms}(\theta)$
#ldms=Sum((((-1)**t)*sqrt(factorial(l+m)*factorial(l-m)*factorial(l+s)*factorial(l-s))/(factorial(l+m-t)*factorial(l-s-t)*factorial(t)*factorial(t+s-m))) * ((cos(theta)/2.0)**(2*l+m-s-2*t))* ((sin(theta)/2.0)**(2*t+s-m)),(t,C1,C2))
#sYlm=((-1)**s)*sqrt((2*l+1)/(4*pi))*ldms.doit()*exp(I*m*phi)


sYlm=((-1)**(m)) *sqrt((factorial(l+m)*factorial(l-m)*(2*l+1))/(4*pi*factorial(l+s)*factorial(l-s)))*(sin(theta/2)**(2*l)) * Sum(((factorial(l-s)*factorial(l+s))/(factorial(t)*factorial(l-s-t)*factorial(t+s-m)*factorial(l-t+m))*((-1)**(l-t-s))*exp(I*m*phi)*(cot(theta/2)**(2*t+s-m))),(t,0,l-s))
#pprint(sYlm.subs({s:1,l:1,m:0,theta:0,phi:pi}).evalf())
#Y=sqrt(3/(8*pi))*sin(theta) swsh function for s=1 l=1 m=0;


dtheta=((1e-1)*(math.pi))
dphi=((1e-1)*2*math.pi)

thetaSz = math.ceil(math.pi/dtheta)
phiSz   =   math.ceil(2*math.pi/dphi)

S=-2
L=[2]
M=[0,1,2]
codeOutput=sys.stdout

## generate header files for the swsh evaluated at each point. 
print("//%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%" ,file=codeOutput )
print("//",file=codeOutput)
print("//Automatic code generated by sympy based swsh module in Dendro-GR ",file=codeOutput)
print("//python module to generate code for far-field energy ratiation extraction (Gravitational Waves).",file=codeOutput)
print("//(c) 2016 University of Utah, All rights reserved.",file=codeOutput)
print("//",file=codeOutput)
print("//%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%" ,file=codeOutput )
##########################################################################

print("#ifndef DENDRO_SWSH_H",file=codeOutput)
print("#define DENDRO_SWSH_H",file=codeOutput)


#evaluates the SWSH functions at the lebedev quadrature points. 

inFilePrefix="/home/milinda/lebedev/lebedev_" # path to the lebedev quadrature point files. 
precision=["025"] 
numPts=[230]

varNamePrefix="m2Y"
print("namespace swsh { " ,file=codeOutput)
for li in L : 
    for mj in M:
        values_real=[]
        values_imag=[]
        for p in precision:
            print("// spin weighted sperical harmonic basis evaluated at the %s lebedev quadrature points " %(p),file=codeOutput)
            inFileName=inFilePrefix+p+".txt"
            inFile=open(inFileName,'r')
            
            lebedev_phi=[]
            lebedev_theta=[]
            lebedev_weight=[]
            
            for line in inFile:
                data=line.split()
                #print(data)
                lebedev_phi.append(str("%.20f"%(((float(data[0])+180)*math.pi)/180.0))) 
                lebedev_theta.append(str("%.20f" %(float(data[1])*math.pi/180.0)) ) 
                lebedev_weight.append(str("%.20f" %(float(data[2]))))
                
            for pts in range(0,len(lebedev_theta)):
                val=complex(sYlm.subs({s:S,l:li,m:mj,theta:float(lebedev_theta[pts]),phi:float(lebedev_phi[pts])}).evalf())
                values_real.append(str("%.20f" %(val.real) ))
                values_imag.append(str("%.20f" %(val.imag) ))
            
            varName=varNamePrefix+str(li).replace('-','m')+"_"+str(mj).replace('-','m')+"_REAL"
            print("static double %s [] = { " %(varName),file=codeOutput)
            print(", ".join(values_real),file=codeOutput)
            print("};")
            

            varName=varNamePrefix+str(li).replace('-','m')+"_"+str(mj).replace('-','m')+"_IMAG"
            print("static double %s [] = { " %(varName),file=codeOutput)
            print(", ".join(values_imag),file=codeOutput)
            print("};")


print("}// end of namespace swsh" ,file=codeOutput)
print("#endif",file=codeOutput)







'''

# evaluates the SWSH functions at points define by the theta and phi step size. 

print("",file=codeOutput)
print("#define NUM_THETA_PTS %d" %(thetaSz),file=codeOutput)
print("#define NUM_PHI_PTS %d" %(phiSz),file=codeOutput)
print("",file=codeOutput);

varNamePrefix="m2Y"

print("namespace swsh { " ,file=codeOutput)

for li in L : 
    for mj in M:
        values_real=[]
        values_imag=[]
        for itheta in np.arange(0,math.pi,dtheta):
            for jphi in np.arange(0,2*math.pi,dphi):
                val=complex(sYlm.subs({s:1,l:li,m:mj,theta:itheta,phi:jphi}).evalf())
                #val1=Y.subs({s:1,l:li,m:mj,theta:itheta,phi:jphi}).evalf()
                #print("(%.6f,%.6f) l: %d m: %d sYml : %.10f Y:%.10f diff: %.10f" %(itheta,jphi,li,mj,val,val1,abs(val1-val))) #pprint(val)
                values_real.append(str("%.20f" %(val.real) ))
                values_imag.append(str("%.20f" %(val.imag) ))

        varName=varNamePrefix+str(li).replace('-','m')+"_"+str(mj).replace('-','m')+"_REAL"
        print("static double %s [] = { " %(varName),file=codeOutput)
        print(", ".join(values_real),file=codeOutput)
        print("};")

        varName=varNamePrefix+str(li).replace('-','m')+"_"+str(mj).replace('-','m')+"_IMAG"
        print("static double %s [] = { " %(varName),file=codeOutput)
        print(", ".join(values_imag),file=codeOutput)
        print("};")


print("}// end of namespace swsh" ,file=codeOutput)

print("#endif",file=codeOutput)
'''
