#!/usr/bin/python3
import os,sys,math
import numpy as np

from .amspyigrf_loader import *

from .amssphharmonics import gso_cosharmonic
from .amssphharmonics import gso_sinharmonic

pi = np.pi
sin = np.sin
cos = np.cos
sqrt = np.sqrt

def ams_xyz_rthetaphi(xyz):
    """
    Returns astronomical latitude and longitude [radians]
    (elevation and azimuth)
    Input:
    xyz [3] or [3xN] array
    Output:
    [radius lat lon] [radians] [3] or [3xN] array
    """
    xyz = np.array(xyz)
    singles = 0

    if(len(xyz.shape)==1):
        singles = 1
        xyz = np.reshape(xyz,[3,1])

    r = np.sqrt(np.sum(xyz*xyz,axis=0))
    #rho = np.sqrt(xyz[0,:]*xyz[0,:]+xyz[1,:]*xyz[1,:])
    #theta = np.where(rho>0.0,np.arctan(xyz[2,:]/rho))
    theta = np.arcsin(xyz[2,:]/r)
    phi = np.arctan2(xyz[1,:],xyz[0,:])

    rtp = np.zeros([3,xyz.shape[1]])
    rtp[0,:] = r
    rtp[1,:] = theta
    rtp[2,:] = phi

    if(singles==1):
        rtp = np.reshape(rtp,[3])

    return rtp

def igrf_potential(xyz, coefstruct):
    """
    Returns the IGRF potential function [nT - km] as a function of ECEF
    coordinates [x,y,z] in [km] in an earth-centered earth-fixed frame.

    Ref for potential function is Aiken et. al. 2021

    Inputs:
        xyz [3] or [3xN] set of positions [km]
        coefstruct - contains gh, n, m, and a arrays
    Outputs:
        V [scalar] or [N]
    """

    xyz = np.array(xyz)
    single = 0
    if(len(xyz.shape)==1):
        single = 1
        xyz = np.reshape(xyz,[3,1])

    
    gh = coefstruct["gh"] #whether using a sin or cosine harmonic
    n = coefstruct["n"] #
    m = coefstruct["m"]
    a = coefstruct["a"]

    rtp = ams_xyz_rthetaphi(xyz)
    theta = pi/2.0-rtp[1,:]
    phi = rtp[2,:]
    rref = 6371.2 #[km] - reference diameter


    V = np.zeros(xyz.shape[1])
    for I in range(0,a.shape[0]):
        V = V + a[I]*np.float64(gh[I]==0)*rref*(rref/rtp[0,:])**(n[I]+1)*gso_cosharmonic(theta,phi,n[I],m[I])
        V = V + a[I]*np.float64(gh[I]==1)*rref*(rref/rtp[0,:])**(n[I]+1)*gso_sinharmonic(theta,phi,n[I],m[I])

    if(single==1):
        V = V[0]

    return V


def igrf_bfield(xyz,coefstruct):
    """
    Takes a quick and dirty numeric derivative of the potential function
    to return the B-field [nT] [x,y,z] from the IGRF-13 potential function

    Inputs:
        xyz [3] or [3xN] set of positions [km] ECEF
        coefstruct - contains gh, n, m, and a arrays
    Outputs:
        B [nT] size [3] or [3xN]
    """
    xyz = np.array(xyz)

    xyz = np.array(xyz)
    single = 0
    if(len(xyz.shape)==1):
        single = 1
        xyz = np.reshape(xyz,[3,1])

    N = xyz.shape[1]
    dx = 0.1 #km
    xp = np.array([dx,0,0]).reshape(3,1)
    xp = np.repeat(xp,N,axis=1)
    yp = np.array([0,dx,0]).reshape(3,1)
    yp = np.repeat(yp,N,axis=1)
    zp = np.array([0,0,dx]).reshape(3,1)
    zp = np.repeat(zp,N,axis=1)


    phixp = igrf_potential(xyz+xp,coefstruct)
    phixm = igrf_potential(xyz-xp,coefstruct)
    phiyp = igrf_potential(xyz+yp,coefstruct)
    phiym = igrf_potential(xyz-yp,coefstruct)
    phizp = igrf_potential(xyz+zp,coefstruct)
    phizm = igrf_potential(xyz-zp,coefstruct)

    B = np.zeros([3,N])
    B[0,:] = -(phixp-phixm)/2.0/dx
    B[1,:] = -(phiyp-phiym)/2.0/dx
    B[2,:] = -(phizp-phizm)/2.0/dx
    
    if(single==1):
        B = np.squeeze(B[:,0])
    
    return B

    