import os,sys,math
import numpy as np

factorial = math.factorial
comb = math.comb
perm = math.perm

from .amspolyfuns import poly_pow
from .amspolyfuns import poly_derivative
from .amspolyfuns import polyval


def legendre_poly(L):
    """
    generates polynomial coefficients for Legendre 
    polynomial of order L
    """

    L = int(L)
    if(L<0):
        return np.float64([0.0])
    elif(L==0):
        return np.float64([1.0])

    p = np.float64([-1,0,1])
    p = poly_pow(p,L)
    
    for I in range(1,L+1):
        p = poly_derivative(p)
        p = p/2/I

    return p

def legendre_eval(x,L):
    p = legendre_poly(L)
    y = polyval(x,p)
    return y


def alegendre_polypart(L,M):
    """
    Polynomial part of associated legendre function
     The full associated legendre function, including condon-shortley phase
     is (1-x^2)*(M/2) * alegendre_polypart(L,M)
    """

    M = int(M)
    L = int(L)
    if(L<0):
        return np.float64([0])
    if(abs(M)>L):
        return np.float64([0])
    
    Ma = abs(M)
    CSP = (-1)**Ma
    p = legendre_poly(L)

    if(Ma>0):
        p = poly_derivative(p,Ma)
    
    p = CSP*p
    if(M<0):
        Nrm = factorial(L-Ma)/factorial(L+Ma)*(-1)**Ma*p
        p = Nrm*p

    return p


def alegendre_eval(x,L,M):
    """
    Evaluates associated legendre function 
        y = P^m_l(x)
    """

    L = int(L)
    M = int(M)
    Ma = abs(M)

    if(L<0):
        y = np.zeros(x.shape)
        return y
    if(Ma>L):
        y = np.zeros(x.shape)
        return y

    if(Ma>0):
        #sqpart = np.sign(1-x**2)*np.sqrt(1-x**2)**(Ma)
        #sign doesn't work here - need to convert sqrt(1-cos2) to sin
        sqpart = np.sqrt(1-x**2)**(Ma)
    else:
        sqpart = np.ones(x.shape)
    
    p = alegendre_polypart(L,M)

    y = polyval(x,p)
    y = y*sqpart

    return y

def gso_cosinealegendre(theta, m, n):
    """
    Returns the Graham-Schmidt semi-normalized associated legendre function
    as described in Winch et. al. 2005.
    y = P^m_n(cos(theta))

    P^m_n(cos(theta)) = sqrt((2-del^0_m)(n-m)!/(n+m)!) * sin(theta)^m *(d/dcos(theta))^m [P_n(cos(theta))]

    normalization of associated cosine and sine harmonics is 4*pi/(2*n+1)

    for use in the IGRF-13 model
    """ 

    n = int(n)
    m = int(m)
    theta = np.array(theta)

    ma = abs(m)

    if(n<0):
        y = np.zeros(theta.shape)
        return y
    if(ma>n):
        y = np.zeros(theta.shape)
        return y
    
    if(ma>0):
        #sqpart = np.sign(1-x**2)*np.sqrt(1-x**2)**(Ma)
        #sign doesn't work here - need to convert sqrt(1-cos2) to sin
        sqpart = np.sin(theta)**(ma)
        nrm = np.sqrt(2*factorial(n-m)/factorial(n+m))
    else:
        sqpart = np.ones(theta.shape)
        nrm = np.sqrt(factorial(n-m)/factorial(n+m))
    
    p = legendre_poly(n)
    if(ma>0):
        p = poly_derivative(p,ma)

    y = polyval(np.cos(theta),p)

    y = y*nrm*sqpart

    return y