#TRY TO EDIT THE EQUATIONS TO THE FORMALISM OF ZDIPY (DONATI+06,FOLSOM+18,VIDOTTO16)
#the implemented changes are m=!0 harmonics x2 since negative harmonics are not included (this is according to Vidotto 16), and leaving the second component out of alpha, beta and gamma (toroidal, or IMAGINARY?)

#calculate B from the spherical harmonics - equations from Kochukhov et. al (2014)
#NOTE THAT IN KOCHUKHOV -l <= m <= l, WHILE WAITE'S DATA IS 0 <= m <= l

lon=[4,10,17,23,29,34,40,45,50,55,59,63,66,69,72,74,75,76,77,77,76,75,74,72,69,66,63,59,55,50,45,40,34,29,23,17,10,4] #number of longitudes on each latitude

from numpy import genfromtxt,pi,sin,cos,sqrt,mean,e
from math import factorial
from scipy.special import lpmv

limit=60 #only take into account B above this latitude
lmax=10#11 #10 when taking from PyZDI, 11 from Waite

#data=genfromtxt('2008/HD35296_mag_2008.c1',skip_header=2)
#data=genfromtxt('2007/HD35296_mag.c1',skip_header=2)
data=genfromtxt('testdata/PyZDI_V1358Ori/outMagCoeff.dat',skip_header=2)

ll=data[0:65,0]
mm=data[0:65,1]

#print(ll)
#print(mm)

#radpot=data[0:77,2] #alpha_l,m
#radtor=data[0:77,3]
radpot=data[0:65,2] #when input comes from PyZDI
radtor=data[0:65,3] #tor actually the imaginary component, leave it out????

#azipot=data[77:154,2] #beta_l,m OR GAMMA???
#azitor=data[77:154,3]
azipot=data[65:130,2] #when input comes from PyZDI, this is definitely beta
azitor=data[65:130,3]

#merpot=data[154:231,2] #gamma_l,m OR BETA???
#mertor=data[154:231,3]
merpot=data[130:195,2] #when input comes from PyZDI, this is definitely gamma
mertor=data[130:195,3]

#print(mean(radpot),mean(radtor))
#print(mean(merpot),mean(mertor))
#print(mean(azipot),mean(azitor))
#print(azipot[1])

Br=[]
Bm=[]
Ba=[]

#print(factorial(7))

'''
def Legendre(n,x):
    if (n==0):
        P=1
    elif (n==1):
        P=x
    else:
        P=(1/n)*((2*n-1)*x*Legendre(n-1,x)-(n-1)*Legendre(n-1,x))
    return P

def ass_Legendre(l,m,x):
    try:
        if (m==0):
            P=Legendre(l,x)
        elif (m>0):
            P=(1/sqrt(1-x**2))*((l-m+1)*x*ass_Legendre(l,m-1,x)-(l+m-1)*ass_Legendre(l-1,m-1,x))
        elif (m<0):
            m=abs(m)
            P=((-1)**m)*(factorial(l-m)/factorial(l+m))*ass_Legendre(l,m,x)
    except ZeroDivisionError:
        P=0
    return P
'''

def C(l,m):
    return sqrt((2*l+1)*factorial(l-abs(m))/(4*pi*factorial(l+abs(m))))

def K(m,phi):
    if (m >= 0):
        return cos(abs(m)*phi)
    else:#if (m < 0):
        return sin(abs(m)*phi)

def P(l,m,x): #associated Legendre polynomial - x should be cos(theta)
    return lpmv(m,l,x) #note that m goes before l!

def dP(l,m,theta): #derivative of a-L's pol. - Eq 9 of Kohcukhov+ 2014
    return (l-abs(m)+1)*P(l+1,abs(m),cos(theta))/sin(theta)-(l+1)*cos(theta)*P(l,abs(m),cos(theta))/sin(theta)
    #return (l-abs(m)+1)*P(l+1,abs(m),cos(theta))/sin(theta)-(l+1)*cos(theta)*ass_Legendre(l,abs(m),cos(theta))/sin(theta)
    
def Y(l,m,theta,phi):
    return -C(l,m)*P(l,abs(m),cos(theta))*K(m,phi)
    #return -C(l,m)*ass_Legendre(l,abs(m),cos(theta))*K(m,phi)

def Z(l,m,theta,phi):
    return C(l,m)*dP(l,abs(m),cos(theta))*K(m,phi)/(l+1)

def X(l,m,theta,phi):
    return -C(l,m)*P(l,abs(m),cos(theta))*m*K(-m,phi)/((l+1)*sin(theta))
    #return -C(l,m)*ass_Legendre(l,abs(m),cos(theta))*m*K(-m,phi)/((l+1)*sin(theta))


latdeg=180/38.0 #the grid contains 38 latitudes
nlat=38

grid=[]
#print(radpot[0])
#print(len(radpot))
for q in range(0,nlat):
    nlon=lon[q]
    add=0.5*latdeg
    theta=(-90+add+q*latdeg)*pi/180
    if (theta*180/pi < limit):
        continue
    print('theta='+str(theta*180/pi))
    print(nlon)
    for u in range(0,nlon):
        phi=(u*360/nlon)*pi/180
        #if (theta*180/pi > 80):
         #   print('phi='+str(phi*180/pi))
        Brad=0
        Bmer=0
        Bazi=0
        ind=0
        for l in range(1,lmax+1):
            #print('l='+str(l))
            for m in range(0,int(l)+1): #from -l to l, but data is only from 0
                ind=ind+1
                #print('m='+str(m)+' ind='+str(ind-1)+' mertor[ind]='+str(mertor[ind-1]))
                alpha=radpot[ind-1]#+radtor[ind-1] #the tor parts are imaginary, and should not be included????
                beta=merpot[ind-1]#+mertor[ind-1]
                gamma=azipot[ind-1]#+azitor[ind-1]
                if (m==0):
                    Brad=Brad+alpha*Y(l,m,theta,phi) #minus comes in next step
                    Bmer=Bmer+beta*Z(l,m,theta,phi)+gamma*X(l,m,theta,phi)
                    Bazi=Bazi+beta*X(l,m,theta,phi)-gamma*Z(l,m,theta,phi)
                else:
                    Brad=Brad+2*alpha*Y(l,m,theta,phi) #minus comes in next step
                    Bmer=Bmer+2*(beta*Z(l,m,theta,phi)+gamma*X(l,m,theta,phi))
                    Bazi=Bazi+2*(beta*X(l,m,theta,phi)-gamma*Z(l,m,theta,phi))
        Br.append(round(-Brad,5)) #minus comes here
        Bm.append(round(-Bmer,5))
        Ba.append(round(-Bazi,5))
        
#print(Br)
#print(Ba)
#print(Bm)
print('number of surface elements: '+str(len(Br)))
print('mean Br: '+str(mean(Br)))
print('mean Bm: '+str(mean(Bm)))
print('mean Ba: '+str(mean(Ba)))
absBr=[]
Btot=[]
for b in range(0,len(Br)):
    absBr.append(abs(Br[b]))
    Btot.append(sqrt(Br[b]**2+Bm[b]**2+Ba[b]**2))
#print('mean |Br|: '+str(mean(absBr)))
print('mean |B|: '+str(mean(Btot)))