# -*- coding: utf-8 -*-
"""
Created on Tue May 12 15:45:20 2020

@author: Daniel
"""

import matplotlib.pyplot as plt
import numpy as np
import datetime
import tsyg
from scipy.spatial import ConvexHull
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.cm as cm
from matplotlib.colors import Normalize


def KF94_field(Rbs, Rmp, Bimfx, Bimfy, Bimfz, x, y, z):
    """Magnetosheath magnetic field (Kobela and Fluckiger 1994)."""
    A = (2 * Rbs - Rmp) / (2 * (Rbs - Rmp))
    l = 3 * Rmp / 2 - x
    Bmsx = -A * (-Bimfx * (1 - Rmp / (2 * l))
                 + Bimfy * (y / l)
                 + Bimfz * (z / l))
    Bmsy = A * (-Bimfx * (y / (2 * l))
                + Bimfy * (2 - (y**2) / (l * Rmp))
                - Bimfz * (y * z / (l * Rmp)))
    Bmsz = A * (-Bimfx * (z / (2 * l))
                - Bimfy * (y * z / (l * Rmp))
                + Bimfz * (2 - (z**2) / (l * Rmp)))

    return Bmsx, Bmsy, Bmsz


def Sibeck_Rmp(pdyn, A0, B0, C0, p0):
    """
    Magnetopause stand-off distance, Sibeck, 1991.

     R^2 + A*x^2+B*(p0/p)^(1/6)*x + C * (p0/p)^(1/3) = 0
     Rmp (stand-off dist) =  x for R = 0
     pdyn : solar wind pressure [nPa ]
     p0 : reference solar wind pressure [nPa]
     A0, B0, C0 : magnetopause coefficients
    """
    p = pdyn
    p1 = (p0/p)**(1/6)
    p2 = (p0/p)**(1/3)
    Rmp = (-B0*p1+np.sqrt(B0**2*p1**2 - 4*A0*C*p2))/(2*A0)
    return Rmp


def centroid(arr):
    """
    Compute face centroid (intersection of medians).

     This is the position where the value of the shear angle for the face will
     be computed. It will give the face color.
    """
    length = arr.shape[0]
    midpoints_ = np.zeros((length, 3))
    for i in range(0, length):
        sum_x = np.sum(arr[i, :,  0])
        sum_y = np.sum(arr[i, :,  1])
        sum_z = np.sum(arr[i, :,  2])
        midpoints_[i] = [sum_x/3, sum_y/3, sum_z/3]
    return midpoints_


def Sibeck_to_KF94(xs, ys, zs, Rmp, A, B, C):
    """
    Map  Sibeck's magnetopause points to KF94 magnetopause.

     xs, ys, zs : coordinates of a point on the Sibeck magnetopause [R_E]
     Rmp : magnetopause stand-off distance  [R_E]
     A, B, C : Sibeck's magnetopause parameters

     Compute the intersection of the normal to Sibeck magnetopause
     at  xs,ys,zs with KF94 magnetopause
    """
    t0 = (
          2*B*Rmp + 4*A*Rmp*xs + 4*ys**2 + 4*zs**2 -
          np.sqrt((-2*B*Rmp - 4*A*Rmp*xs - 4*ys**2 - 4*zs**2)**2 -
                  4*(-4*ys**2 - 4*zs**2) *
                  (2*Rmp**2 - 2*Rmp*xs - ys**2 - zs**2))
        )/(2*(-4*ys**2 - 4*zs**2))
    xin0 = xs + (B+2*A*xs)*t0
    yin0 = ys + 2*ys*t0
    zin0 = zs + 2*zs*t0
    return xin0, yin0, zin0  # point on  KF94 mpause


def find_color_for_point(pt, cmap, R_bs, R_mp,
                         A, B, C,
                         pdyn, dst, ps, Bimfx, Bimfy, Bimfz):
    """
    Compute the shear angle and the corresponding face color for point pt.

     pt : numpy array of size 3
     cmap : color map
     R_bs : bow shock standoff distance [R_E]
     R_mp : magnetopause standoff distance [R_E]
     A, B, C : Sibeck's magnetopause parameters'
     ps : dipole tilt
     Bimfx, Bimfy, Bimfz: IMF components
    """
    iopt = 0  # dummy variable for T96
    parmod1 = [pdyn, dst, Bimfy, Bimfz, 0, 0, 0, 0, 0, 0]
    x_, y_, z_ = pt
    bxgsm_, bygsm_, bzgsm_ = [0., 0., 0.]
    Bmsx_, Bmsy_, Bmsz_ = [0., 0., 0.]
    B_, b_, alpha_, dp_ = [0., 0., 0., 0.]
    hxgsw_, hygsw_, hzgsw_ = [0.0, 0.0, 0.0]
    bx1, by1, bz1 = [0., 0., 0.]
    bxgsm0_, bygsm0_, bzgsm0_ = tsyg.igrf_gsw_08(x_, y_, z_,
                                                 hxgsw_, hygsw_, hzgsw_)
    bxgsm1_, bygsm1_, bzgsm1_ = tsyg.t96_01(iopt, parmod1,
                                            ps, x_, y_, z_, bx1, by1, bz1)
    bxgsm_ = bxgsm0_ + bxgsm1_
    bygsm_ = bygsm0_ + bygsm1_
    bzgsm_ = bzgsm0_ + bzgsm1_
    x_k, y_k, z_k = Sibeck_to_KF94(x_, y_, z_, R_mp, A, B, C)
    Bmsx_, Bmsy_, Bmsz_ = KF94_field(R_bs, R_mp, Bimfx, Bimfy, Bimfz,
                                     x_k, y_k, z_k)
#    print('bxgsm_,bygsm_,bzgsm_: ', bxgsm_,bygsm_,bzgsm_)
#    print('Bmsx_, Bmsy_, Bmsz_: ', Bmsx_, Bmsy_, Bmsz_)
    dp_ = Bmsx_ * bxgsm_ + Bmsy_ * bygsm_ + Bmsz_ * bzgsm_  # dot product
    B_ = np.sqrt(Bmsx_**2 + Bmsy_**2 + Bmsz_**2)
    b_ = np.sqrt(bxgsm_**2 + bygsm_**2 + bzgsm_**2)
    # angle between B msheath and b msphere
    alpha_ = np.degrees(np.arccos(dp_ / b_ / B_))
    # print('x,y,z, alpha_: ',x_, y_, z_,alpha_)
    norm = Normalize(vmin=0, vmax=180)
    col = cmap(norm(alpha_))
    return col


def main(): """Plot the  shear angle on a 3D magnetopause."""


# input------------------------------------------------------
year, month, day, hour, min, sec = [1997, 5, 9, 12, 10, 0]
Bimfx = 1.0  # nT
Bimfy = -2.6
Bimfz = -1.8
dst = -9   # DST index
vsw = 300.  # solar wind speed
n = 11.0  # solar wind number density cm^-3
# Sibeck s magnetopause parameters
#   Use the coefficients for the pressure bin [1.47, 2.60] nPa
#   p0 is the reference solar wind pressure
A = 0.14
B = 18.2
C = -217.2
p0 = 2.04  # nPa
# plot features
N = 20  # number of yz ellipses to plot
npe = 100  # number of points per yz ellipse
# end input---------------------------------------------------

imf_clock_angle = np.arctan2(Bimfy, Bimfz)
if imf_clock_angle < 0:
    imf_clock_angle = imf_clock_angle + 2*np.pi
imf_clock_angle = np.degrees(imf_clock_angle)
print('imf clock angle : ', imf_clock_angle)
vgse = [-vsw, 0, 0]  # solar wind velocity in GSE: T96 input.
m_p = 1.6726219e-27  # proton mass kg
# solar wind dynamic pressure nPa
pdyn = (0.85 * n * 1e6 * m_p * (vsw * 1000)**2)*1e9
# Standoff distance of the Earth’s bow shock (Jelinek)
Rbs = 15.02 * pdyn**(-1 / 6.55)
print("R_bsj  = ", Rbs, ' R_E')

Rmp = Sibeck_Rmp(pdyn, A, B, C, p0)
print("Rmp  (Sibeck) = ", Rmp, ' R_E')

# Plot -----------------------------------------------------------------------
fig = plt.figure()
ax = fig.gca(projection='3d')
fig.tight_layout()

# Create the triangulations of the T96 and KF94 magnetopauses
# Magnetopause  Kobela and Fluckiger, 1994
pts3d = np.zeros((N*npe, 3))
X = np.zeros((N, npe))
Y = np.zeros((N, npe))
Z = np.zeros((N, npe))

R = 0.
for i in range(0, N):
    R = R + 1
    # R=4.
    alpha = np.linspace(-1 * np.pi / 2, 3 * np.pi/2, npe)
    Z[i] = R * np.cos(alpha)
    # Y=np.zeros_like(Z)
    Y[i] = R * np.sin(alpha)
    X[i] = -(-2 * (Rmp**2) + Y[i]**2 + Z[i]**2) / (2 * Rmp)
    # print('X= ', X)
# ax0.plot(X)
# Magnetopause T96
# xmp, ymp, zmp : closest points of T96 mpause to KF94 mpause
xn_pd = pdyn
vel = -vsw
xmp = np.zeros_like(X)
ymp = np.zeros_like(Y)
zmp = np.zeros_like(Z)

for i in range(0, N):
    for j in range(0, npe):
        xgsm = X[i][j]
        ygsm = Y[i][j]
        zgsm = Z[i][j]
        xmgnp, ymgnp, zmgnp, dist, id = tsyg.t96_mgnp_08(xn_pd, vel,
                                                         xgsm, ygsm, zgsm)
        xmp[i][j] = xmgnp
        ymp[i][j] = ymgnp
        zmp[i][j] = zmgnp
        # print('GSM in: ', xgsm,ygsm,zgsm)
        # print('GSM MP: ', xmgnp,ymgnp,zmgnp, dist, id)

# ax.plot3D(X[i], Y[i], Z[i], 'red', label='KF94')
# ax.plot3D(xmp[i], ymp[i], zmp[i], 'blue', label='T96')
# #ax.scatter3D(xmp, ymp, zmp, c=zmp, cmap='Blues',label='T96')

ax.set_xlim(-20, 20)
ax.set_ylim(-20, 20)
ax.set_zlim(-20, 20)
# ax.set_aspect("equal")
ax.set_xlabel(r'$X_{GSM}$ ($R_E$)')
ax.set_ylabel(r'$Y_{GSM}$ ($R_E$)')
ax.set_zlabel(r'$Z_{GSM}$ ($R_E$)')
# ax.set_title('Magnetopause ')
fig.suptitle('Magnetopause' + ' ' + str(year)+'/'
             + "{:02d}".format(month) + '/'
             + "{:02d}".format(day) + ' ' + "{:02d}".format(hour) + ':'
             + "{:02d}".format(min)+':'+"{:02d}".format(sec)+' UT'
             + '\n' + 'IMF clock angle : '
             + "{:3.1f}".format(imf_clock_angle) + u"\N{DEGREE SIGN}")
ax.legend(fancybox=True, framealpha=0.5)

# Triangulation of T96 magnetopause
pts3d = np.zeros((N * npe, 3))
k = 0
for i in range(0, N):
    for j in range(0, npe):
        pts3d[k][0] = xmp[i][j]
        pts3d[k][1] = ymp[i][j]
        pts3d[k][2] = zmp[i][j]
        k = k + 1

hull = ConvexHull(pts3d, qhull_options='QJ')
indices = hull.simplices
faces = pts3d[indices]

# delete triangles closing the Mpause (convex hull) at Xmin
# faces1=triangles to delete
# faces1=np.delete(faces, np.where(faces[:,:,0]>-7.7), axis=0)
# triangles to keep
x_min = np.min(faces[:, :, 0] + 0.01)
tri_to_delete = []
for i in range(0, np.size(faces, axis=0)):
    if faces[i][0][0] < x_min or faces[i][1][0] < x_min or \
            faces[i][2][0] < x_min:
        tri_to_delete.append(i)
faces = np.delete(faces, tri_to_delete, axis=0)
# print('area: ', hull.area)
# print('volume: ', hull.volume)

# Compute the shear angle -------------------------------------------------
# Magnetosheath field
Bmsx, Bmsy, Bmsz = KF94_field(Rbs, Rmp, Bimfx, Bimfy, Bimfz, X, Y, Z)
bxgsm = np.zeros_like(Bmsx)
bygsm = np.zeros_like(Bmsx)
bzgsm = np.zeros_like(Bmsx)
t1 = datetime.datetime(year, month, day, hour, min, sec)
t0 = datetime.datetime(1970, 1, 1)
ut = (t1-t0).total_seconds()
# IGRF init
hxgsw, hygsw, hzgsw = [0.0, 0.0, 0.0]
doy = (t1 - datetime.datetime(t1.year, 1, 1)).days + 1
# print("doy = ", doy)
vgse = [-vsw, 0, 0]  # solar wind velocity in GSE.
tsyg.recalc_08(year, doy, hour, min, sec, *vgse)
# get the dipole tilt ps from the geopack1 fortran common block
# the common block geopack1 is mapped into 2 arrays aa[10] and bb[22]
# and 2 scalars sps et cps
ps = tsyg.geopack1.bb[3]
# T96 init
bx1, by1, bz1 = [0.0, 0.0, 0.0]
cmap = cm.rainbow   # set the colormap
# map the colormap cmap to shear angle range from 0 to 180 deg
norm = Normalize(vmin=0, vmax=180)
midp = centroid(faces)
facecolors = [find_color_for_point(pt, cmap, Rbs, Rmp,
                                   A, B, C,
                                   pdyn, dst, ps,
                                   Bimfx, Bimfy, Bimfz)
              for pt in midp]  # smooth gradient
coll = Poly3DCollection(faces, facecolors=facecolors,
                        edgecolors='none', alpha=1.0)
ax.add_collection(coll)
cbar = fig.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax)
cbar.ax.set_ylabel('Shear Angle (deg)')

plt.show()


if __name__ == "__main__":
    main()