#!/usr/bin/python3

import os,sys,math
import numpy as np
import matplotlib.pyplot as plt

from amspyigrf13 import *

pi = np.pi

def test_loading_igrfdata():
    q = load_igrfdata()
    print(q)

def test_getcomponents():

    igrfdata = load_igrfdata()
    yrs = np.linspace(1895,2035,1000)
    q = []
    for I in range(0,len(yrs)):
        a = igrf_getcomponents(yrs[I],igrfdata)
        q.append(a)
    
    q = np.array(q)

    for I in range(0,q.shape[1]):
        plt.plot(yrs,q[:,I])
        
    plt.show()

    return

def test_gsosph_normalization(n1,m1,n2,m2):

    Ntheta = 50
    Nphi = 100

    sum1 = 0.0
    sum2 = 0.0
    sum3 = 0.0

    for I in range(0,Ntheta):
        for J in range(0,Nphi):
            theta = I/Ntheta*pi
            phi = J/Nphi*2*pi
            dtheta = pi/Ntheta
            dphi = 2*pi/Nphi
            dA = np.sin(theta)*dtheta*dphi

            sum1 = sum1 + gso_cosharmonic(theta,phi,n1,m1)*gso_cosharmonic(theta,phi,n2,m2)*dA
            sum2 = sum2 + gso_sinharmonic(theta,phi,n1,m1)*gso_sinharmonic(theta,phi,n2,m2)*dA
            sum3 = sum3 + gso_cosharmonic(theta,phi,n1,m1)*gso_sinharmonic(theta,phi,n2,m2)*dA
            
    return [sum1,sum2,sum3]

def test_norm2():

    for I in range(0,5):
        for J in range(0,5):
            [a,b,c] = test_gsosph_normalization(I,J,I,J)
            a2 = a*(2*I+1)/(4*pi)
            b2 = b*(2*I+1)/(4*pi)
            c2 = c*(2*I+1)/(4*pi)
            print("{} {}: {} {} {}".format(I,J,a2,b2,c2))
            #print("{} {}".format(I,J))

    return

def test_igrf_potential():

    coefstruct = load_igrfdata()
    coefstruct = igrf_getcompstruct(2024,coefstruct)

    xyz = [6371,0,0]
    V1 = igrf_potential(xyz,coefstruct)
    B1 = igrf_bfield(xyz,coefstruct)
    print("Potential at {} = {} [nT-km], B = {} [nT]".format(xyz,V1,B1))

    xyz = [-6371,0,0]
    V1 = igrf_potential(xyz,coefstruct)
    B1 = igrf_bfield(xyz,coefstruct)
    print("Potential at {} = {} [nT-km], B = {} [nT]".format(xyz,V1,B1))

    xyz = [6371+100,0,0]
    V1 = igrf_potential(xyz,coefstruct)
    B1 = igrf_bfield(xyz,coefstruct)
    print("Potential at {} = {} [nT-km], B = {} [nT]".format(xyz,V1,B1))

    xyz = [0,0,-6371]
    V1 = igrf_potential(xyz,coefstruct)
    B1 = igrf_bfield(xyz,coefstruct)
    print("Potential at {} = {} [nT-km], B = {} [nT]".format(xyz,V1,B1))

    xyz = [0,0,6371]
    V1 = igrf_potential(xyz,coefstruct)
    B1 = igrf_bfield(xyz,coefstruct)
    print("Potential at {} = {} [nT-km], B = {} [nT]".format(xyz,V1,B1))

    return 

def mycontourf(x,y,s,N=100):

    s = np.copy(s)
    s[np.isinf(s)] = 0
    s[np.isnan(s)] = 0

    # s[s<smin] = smin
    # s[s>smax] = smax
    smin = np.min(s)
    smax = np.max(s)
    if(np.abs(smax-smin)<1E-15):
        smax = smin + 1
        s[1,0] = smax
    lvls = np.linspace(smin,smax,N)

    plt.contourf(x,y,s,cmap="jet",levels=lvls)
    plt.colorbar()

    return

def getuvw(xyz):



    return [u,v,w]


def test_igrf_plotsurfaceB():

    coefstruct = load_igrfdata()
    coefstruct = igrf_getcompstruct(2020,coefstruct)

    Nlats = 25
    Nlons = 50
    lats = np.linspace(-pi/2,pi/2,Nlats)
    lons = np.linspace(-pi,pi,Nlons)
    [lons2,lats2] = np.meshgrid(lons,lats)
    lons2 = lons2.transpose(); lats2=lats2.transpose()
    lons2 = np.array(lons2,order="F")
    lats2 = np.array(lats2,order="F")
    lons3 = np.reshape(lons2,[Nlons*Nlats],order="F")
    lats3 = np.reshape(lats2,[Nlons*Nlats],order="F")
    

    Re = 6378
    
    x = np.cos(lats3)*np.cos(lons3)*Re
    y = np.cos(lats3)*np.sin(lons3)*Re
    z = np.sin(lats3)*Re
    # xyz = np.zeros([3,Nlons*Nlats],order="F")
    # xyz[0,:] = x
    # xyz[1,:] = y
    # xyz[2,:] = z
    xyz = np.stack([x,y,z],axis=0)

    B = igrf_bfield(xyz,coefstruct)
    [lonlatalt,Buvw] = grs80_xyzvectors_to_uvw(xyz,B)

    print(Buvw.shape)
    Bmag = np.squeeze(np.sqrt(np.sum(B**2,axis=0)))

    lats2 = lats2*180/pi
    lons2 = lons2*180/pi
    Bmag2 = np.reshape(Bmag,[Nlons,Nlats],order="F")

    fig = plt.figure(1)
    mycontourf(lons2,lats2,Bmag2)
    ax = plt.gca()
    ax.set_aspect("equal")
    plt.title("IGRF-13 Surface B magnitude [nT]")
    plt.xlabel("Longitude [deg]")
    plt.ylabel("Latitude [deg]")
    

    # fig = plt.figure(2).add_subplot(projection="3d")
    # sc = 1/60
    # plt.quiver(xyz[0,:],xyz[1,:],xyz[2,:],B[0,:]*sc,B[1,:]*sc,B[2,:]*sc,color=[0,0,0])
    # ax = plt.gca()
    # ax.set_aspect("equal")

    print(np.max(lonlatalt[0,:]),np.min(lonlatalt[0,:]))
    print(np.max(lonlatalt[1,:]),np.min(lonlatalt[1,:]))
    print(np.max(lonlatalt[2,:]),np.min(lonlatalt[2,:]))
    


    fig = plt.figure(2)
    sc = 1/60
    plt.quiver(lonlatalt[0,:],lonlatalt[1,:],Buvw[0,:]*sc,Buvw[1,:]*sc,color=[0,0,0])
    ax = plt.gca()
    ax.set_aspect("equal")
    plt.title("IGRF-13 Surface B Heading")
    plt.xlabel("Longitude [deg]")
    plt.ylabel("Latitude [deg]")

    plt.show()

    return


if(__name__=="__main__"):
    
    #test_loading_igrfdata()
    #test_getcomponents()
    #test_norm2()
    #test_igrf_potential()
    test_igrf_plotsurfaceB()

    
    pass