# -*- coding: utf-8 -*-
"""beamtools.py

   module for calculating geometry parameters and magnetic aspect
   angle of radar targets

   use aspect_elaz or aspect_txty to calculate aspect angles of targets
   specified by (el,az) or (tx,ty) angles

  Created by Erhan Kudeki on 11/29/08.
  Copyright (c) 2008 ECE, UIUC. All rights reserved.

  History
  Pablo Reyes on 12/15/2021
    - flattening has been changed from 1/298.257 to 1./298.257223563
      using the WGS84 reference in:
      https://gis.icao.int/eganp/webpdf/REF08-Doc9674.pdf
    - First crack at PEP 8 compliance
    - Using NumPy Style Python Docstrings
"""

import datetime
import numpy as np
from ..pyigrf import pyigrf

# module level variables:

eps = np.finfo(float).eps                  # float resolution
deg = np.pi / 180.                         # to express angles in degree values
a_igrf = 6371.2                            # mean earth radius (km)

a_WGS = 6378.137                           # equatorial radius WGS 84
flatness = 1/298.257223563
b_WGS = a_WGS * (1 - flatness)             # WGS polar radius
eccentricity = np.sqrt(a_WGS**2 - b_WGS**2) / a_WGS

def datetime2year(datetime0):
    """Converts datetime into year with decimals

    Parameters
    ----------
    datetime0: datetime
        Date and time to convert to float year.

    Returns
    -------
    fyear : float
        The date in year as a float number with fraction.

    """
    nyear = datetime0.year
    yeardays = (datetime.datetime(nyear+1,1,1) - datetime.datetime(nyear,1,1))
    return nyear + (datetime0 - datetime.datetime(nyear,1,1))/yeardays

def llh2xyz(latg, lon, h):
    """Converts a point in geodetic to ECEF coordinates.

    Parameters
    ----------
    latg : float
        The geodetic latitude in radians.
    lon : float
        The longitude in radians.
    h : float
        The height above the local ellipsoid

    Returns
    -------
    x : float
        ECEF x coordinate
    y : float
        ECEF y coordinate
    z : float
        ECEF z coordinate
    """

    n = a_WGS / np.sqrt(1 - flatness * (2 - flatness) * np.sin(latg)**2.)
    # cartesian geocentric coordinates wrt Greenwich
    x = (n + h) * np.cos(latg) * np.cos(lon)
    y = (n + h) * np.cos(latg) * np.sin(lon)
    z = (n * (1 - eccentricity**2.) + h) * np.sin(latg)

    return x, y, z

def xyz2llh(x, y, z):
    """Converts a point in  ECEF to geodetic coordinates.

    Returns longitude 'lon', geodetic latitude 'lat', and height 'h'
    of position (x,y,z) defined in geocentric coordinate system

    Parameters
    ----------
    x : float
        ECEF x coordinate
    y : float
        ECEF y coordinate
    z : float
        ECEF z coordinate

    Returns
    -------
    lat : float
        The geodetic latitude in radians.
    lon : float
        The longitude in radians.
    h : float
        The height above the local ellipsoid
    """

    p = np.sqrt(x**2. + y**2.)
    lon = np.arctan2(y, x)
    lat = np.arctan2(z, p)
    latp = lat
    for i in range(10):
        n = a_WGS / np.sqrt(1 - flatness * (2 - flatness) * np.sin(latp)**2.)
        h = p / np.cos(latp) - n
        lat = np.arctan(z / (p * (1 - n * eccentricity**2. / (n + h))))
        if abs(lat - latp) < 3 * eps:
            n = a_WGS / np.sqrt(1 - flatness * (2 - flatness) * np.sin(lat)**2.)
            h = p / np.cos(lat) - n
            break
        latp = lat
    return lat, lon, h

class TargetGeometry:
    """Will contain radar coordinates and coordinate conversions"""

    def __init__(self,lat0, lon0, h0, IGRFmodel=None):

        self.lat0 = lat0
        self.lon0 = lon0
        self.h0 = h0

        # None argument instantiates the latest
        self.igrf = pyigrf(IGRFmodel)
        # igrf to be instantiated once and to be remembered  between calls

        n0 = a_WGS / np.sqrt(1 - flatness * (2 - flatness) * np.sin(self.lat0)**2.)
        # cartesian geocentric coordinates wrt Greenwich
        x0 = (n0 + self.h0) * np.cos(self.lat0) * np.cos(self.lon0)
        y0 = (n0 + self.h0) * np.cos(self.lat0) * np.sin(self.lon0)
        z0 = (n0 * (1 - eccentricity**2) + self.h0) * np.sin(self.lat0)
        xyz0 = np.array([x0, y0, z0])
        xy0 = np.array([x0, y0])
        r0 = np.sqrt(np.dot(xyz0, xyz0))
        p0 = np.sqrt(np.dot(xy0, xy0))

        self.xyz0 = xyz0 # to be remembered between calls
        # unit vectors to be remembered between calls, hence self.xxx
        self.east0 = np.array([-y0, x0, 0]) / p0
        # zenith and north directions wrt local ellipsoid
        self.zenith0 = np.array([np.cos(self.lat0) * np.cos(self.lon0),
                                 np.cos(self.lat0) * np.sin(self.lon0),
                                 np.sin(self.lat0)])
        self.north0 = np.cross(self.zenith0, self.east0)


    def xyz2dec_ha(self, vec):
        """Declination and hour angle in target direction used to describe radar beam
        direction at JRO, corresponding to latitude and relative longitude of the
        beam-spot on the celestial sphere, corresponds to rr->\infty, in which case:
        """
        vec = vec / np.sqrt(np.dot(vec, vec))
        p = np.sqrt(vec[0]**2 + vec[1] ** 2)
        dec = np.arctan2(vec[2], p) / deg                                  # in degrees
        ha = (np.arctan2(vec[1], vec[0]) - self.lon0) * (24 / (2.*np.pi)) * 60  # in minutes

        return dec,ha

    def dec_ha2el_az(self, dec, ha):
        """Returns elevation and azimuth angles of a radar beam with respect to
        local tangent plane.

        the beam is specified by:
                declination dec (deg)
                hour angle  ha  (min)
        with respect to radar location at longitude lon0 and height h0
        above reference ellipsiod at geodetic latitude lat0
        """

        lat = dec * deg                                # on celestial sphere
        lon = 2 * np.pi * (ha / (24 * 60))
        lon = lon + self.lon0                    # on celestial sphere
        vec = np.array([np.cos(lat) * np.cos(lon),
                        np.cos(lat) * np.sin(lon),
                        np.sin(lat)])
        hor = vec - np.dot(vec, self.zenith0) * self.zenith0
        hor = hor / np.sqrt(np.dot(hor, hor))
        el = np.arccos(np.dot(hor, vec)) / deg
        north = np.dot(hor, self.north0)
        east = np.dot(hor, self.east0)
        az = np.arctan2(east, north) / deg

        return el,az

    def aspect_angle(self, year, xyz):
        """Returns the magnetic aspect angle (rad) of a target with
        geocentric vector xyz defined in geocentric coordinates and local B
        """

        r = np.sqrt(np.dot(xyz, xyz))
        p = np.sqrt(xyz[0] ** 2 + xyz[1] ** 2)
        lat = np.arctan2(xyz[2], p)
        lon = np.arctan2(xyz[1], xyz[0])
        radial = xyz / r;                        # directions from target
        east = np.array([-xyz[1], xyz[0], 0]) / p
        north = -np.cross(east, radial)
        rr = xyz - self.xyz0
        u_rr = rr / np.sqrt(np.dot(rr, rr))      # unit vector from radar to target

        [bX, bY, bZ, bB] = self.igrf.igrf_B(year, r-a_igrf, lon / deg, lat / deg)
        bfield = np.array([bX, bY, bZ])
        B = bX * north + bY * east - bZ * radial
        u_B = B / np.sqrt(np.dot(B, B))
        aspect = np.arccos(np.dot(u_B, u_rr))
        return r, lat, lon, aspect, B

    def aspect_elaz(self, year, rr, el, az):
        """Returns magnetic aspect angle and geocentric coordinates of a target

        Parameters
        ----------
        year: float
          The date expressed as year. e.g.  Jan 31, 2025 is 2025. + 31 / 365
        rr : float
          Radius in km from the radar to the target.
        el : float
          Elevation in rad above local tangent plane to ellipsoid.
        az : float
          Azimuth in rad east of local north.

        Returns
        -------
        float
          Distance in km of the target from the center of the Earth.
        float
          Longitude in rad of the target in geodetic coordinates.
        float
          Latitude in rad  of the target in geodetic coordinates.
        float
          Declination in deg of target on the celestial coordinate system.
        float
          Hour angle in minutes of target on the celestial coordinate system.
        float
          Magnetic aspect angle between the radar pointing direction and the
          geomagnetic field B at the target.
        array_like
          Magnetic field B at the target.
        """

        tx = np.cos(el) * np.sin(az)                    # direction cosines wrt east and north
        ty = np.cos(el) * np.cos(az)
        tz = np.sin(el)
        xyz = self.xyz0 + rr * (tx * self.east0 + ty * self.north0 + tz * self.zenith0)
        #geocentric coordinates of target

        [r, lat, lon, aspect,B] = self.aspect_angle(year, xyz)
        [dec,ha] = self.xyz2dec_ha(xyz - self.xyz0)
        return r, lon, lat, dec, ha, aspect, B

    def cosBs(self, year, rr, el, az):
        """Decomposes the radial unit vector to the target to direction cosines
        of magnetic North, East, and Up
        """

        tx = np.cos(el) * np.sin(az)                              # direction cosines wrt east and north
        ty = np.cos(el) * np.cos(az)
        tz = np.sin(el)
        xyz = self.xyz0 + rr * (tx * self.east0 + ty * self.north0 + tz * self.zenith0)     # target vector
        r = np.sqrt(np.dot(xyz, xyz))
        lat, lon, h = self.xyz2llh(xyz[0], xyz[1], xyz[2])         # target lat, lon, height

        radial = xyz / r; # unit vector to target
        p = np.sqrt(xyz[0] ** 2 + xyz[1] ** 2)
        east = np.array([-xyz[1], xyz[0], 0]) / p # unit vector to east from target
        north = -np.cross(east, radial) # unit vector to north from target
        rr_ = xyz - self.xyz0 # vector from radar to target
        rr_u = rr_ / np.sqrt(np.dot(rr_, rr_)) # unit vector from radar to target

        [bX, bY, bZ, bB] = self.igrf.igrf_B(year, r - a_igrf, lon / deg, lat / deg)
        bfield = np.array([bX, bY, bZ])
        B = bX * north + bY * east - bZ * radial # magnetic field vector B
        bn = B / np.sqrt(np.dot(B, B))
        # "magnetic north" unit vector since B points by definition in "magnetic north" direction

        be = np.cross(bn, radial)
        be = be / np.sqrt(np.dot(be, be)) # magnetic east unit vector
        bu = np.cross(be, bn) # magnetic up unit vector

        cosBn = np.dot(bn, rr_u) # magnetic north direction-cosine of rr_u
        aspect_angle = np.arccos(cosBn)
        cosBe = np.dot(be, rr_u) # magnetic east direction-cosine of rr_u
        cosBu = np.dot(bu, rr_u) # magnetic up direction-cosine of rr_u

        """
        uLOS=cosBe*U(h)+cosBn*V(h)+cosBu*W(h) ... LOS wind model in terms of wind components to calculate and direction cosines
        """

        return r, lat, lon, h, xyz, B, aspect, cosBn, cosBe, cosBu

    def xyz2elaz(self, xyz):
        """Returns elevation and azimuth of a target at xyz ECEF coordinates

        Parameters
        ----------
        xyz : [float, float, float]
          ECEF [x, y, z] coordinates

        Returns
        -------
        rr : float
          The range in km.
        el : float
          The elevation in radians. (rad above local tangent plane to ellipsoid)
        az : float
          The azimuth in radians. (rad east of local north)

        """

        vec = xyz - self.xyz0
        rr = np.linalg.norm(vec)

        east = np.dot(vec, self.east0)
        north = np.dot(vec, self.north0)
        el = 90 - np.arccos(np.dot(vec, self.zenith0) / rr) * 180 / np.pi
        az = np.arctan2(east, north) * 180 / np.pi

        return rr,el,az

    def elaz2xyz(self, rr, el, az):
        """Returns the xyz ECEF coordinates of a target

        The target is specified by elevation, azimuth, and range from the radar.

        Parameters
        ----------
        rr : float
          The range in km.
        el : float
          The elevation in radians. (rad above local tangent plane to ellipsoid)
        az : float
          The azimuth in radians. (rad east of local north)

        Returns
        -------
        xyz : [float, float, float]
          ECEF [x, y, z] coordinates
        """
        tx = np.cos(el) * np.sin(az)        # direction cosines wrt east and north
        ty = np.cos(el) * np.cos(az)
        tz = np.sin(el)
        #geocentric coordinates of target :
        xyz = self.xyz0 + rr * (tx * self.east0 + ty * self.north0 + tz * self.zenith0)

        return xyz
