# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Contains the transformation functions for getting to/from ITRS, TEME, GCRS, and CIRS.
These are distinct from the ICRS and AltAz functions because they are just
rotations without aberration corrections or offsets.
"""

import numpy as np
import erfa

import astropy.units as u
from astropy.coordinates.baseframe import frame_transform_graph
from astropy.coordinates.transformations import FunctionTransformWithFiniteDifference
from astropy.coordinates.matrix_utilities import matrix_transpose
from astropy.coordinates.representation import CartesianDifferential

from .gcrs import GCRS, PrecessedGeocentric
from .cirs import CIRS
from .itrs import ITRS
from .equatorial import TEME, TETE
from .utils import get_polar_motion, get_jd12, EARTH_CENTER

# # first define helper functions


def teme_to_itrs_mat(time):
    # Sidereal time, rotates from ITRS to mean equinox
    # Use 1982 model for consistency with Vallado et al (2006)
    # http://www.celestrak.com/publications/aiaa/2006-6753/AIAA-2006-6753.pdf
    gst = erfa.gmst82(*get_jd12(time, 'ut1'))

    # Polar Motion
    # Do not include TIO locator s' because it is not used in Vallado 2006
    xp, yp = get_polar_motion(time)
    pmmat = erfa.pom00(xp, yp, 0)

    # rotation matrix
    # c2tcio expects a GCRS->CIRS matrix as it's first argument.
    # Here, we just set that to an I-matrix, because we're already
    # in TEME and the difference between TEME and CIRS is just the
    # rotation by the sidereal time rather than the Earth Rotation Angle
    return erfa.c2tcio(np.eye(3), gst, pmmat)


def gcrs_to_cirs_mat(time):
    # celestial-to-intermediate matrix
    return erfa.c2i06a(*get_jd12(time, 'tt'))


def cirs_to_itrs_mat(time):
    # compute the polar motion p-matrix
    xp, yp = get_polar_motion(time)
    sp = erfa.sp00(*get_jd12(time, 'tt'))
    pmmat = erfa.pom00(xp, yp, sp)

    # now determine the Earth Rotation Angle for the input obstime
    # era00 accepts UT1, so we convert if need be
    era = erfa.era00(*get_jd12(time, 'ut1'))

    # c2tcio expects a GCRS->CIRS matrix, but we just set that to an I-matrix
    # because we're already in CIRS
    return erfa.c2tcio(np.eye(3), era, pmmat)


def tete_to_itrs_mat(time):
    # compute the polar motion p-matrix
    xp, yp = get_polar_motion(time)
    sp = erfa.sp00(*get_jd12(time, 'tt'))
    pmmat = erfa.pom00(xp, yp, sp)

    # now determine the greenwich apparent siderial time for the input obstime
    # we use the 2006A model for consistency with RBPN matrix use in GCRS <-> TETE
    ujd1, ujd2 = get_jd12(time, 'ut1')
    jd1, jd2 = get_jd12(time, 'tt')
    gast = erfa.gst06a(ujd1, ujd2, jd1, jd2)

    # c2tcio expects a GCRS->CIRS matrix, but we just set that to an I-matrix
    # because we're already in CIRS equivalent frame
    return erfa.c2tcio(np.eye(3), gast, pmmat)


def gcrs_precession_mat(equinox):
    gamb, phib, psib, epsa = erfa.pfw06(*get_jd12(equinox, 'tt'))
    return erfa.fw2m(gamb, phib, psib, epsa)


def get_location_gcrs(reference, gcrs_to_reference_matrix):
    """Create a GCRS frame at the location and obstime of a given reference.

    Helper function that avoids location.get_gcrs (which would trigger
    infinite recursion in some cases), and uses an already calculated
    GCRS to Reference matrix to calculate obsgeoloc and obsgeovel
    required for GCRS.
    """
    # TODO: ideally, GCRS and other frames would share how they describe
    # the location.  See gh-10996.
    obstime = reference.obstime
    loc_itrs = reference.location.get_itrs(obstime)
    zeros = np.broadcast_to(0. * u.km / u.s, (3,) + loc_itrs.shape, subok=True)
    loc_itrs.data.differentials['s'] = CartesianDifferential(zeros)
    gcrs_cart = (loc_itrs.transform_to(reference.__class__(obstime=obstime))
                 .cartesian.transform(matrix_transpose(gcrs_to_reference_matrix)))
    return GCRS(obstime=obstime,
                obsgeoloc=gcrs_cart.without_differentials(),
                obsgeovel=gcrs_cart.differentials['s'].to_cartesian())


# now the actual transforms

@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, TETE)
def gcrs_to_tete(gcrs_coo, tete_frame):
    # Classical NPB matrix, IAU 2006/2000A
    # (same as in builtin_frames.utils.get_cip).
    rbpn = erfa.pnm06a(*get_jd12(tete_frame.obstime, 'tt'))
    # Get GCRS coordinates for the target observer location and time.
    loc_gcrs = get_location_gcrs(tete_frame, rbpn)
    gcrs_coo2 = gcrs_coo.transform_to(loc_gcrs)
    # Now we are relative to the correct observer, do the transform to TETE.
    # These rotations are defined at the geocenter, but can be applied to
    # topocentric positions as well, assuming rigid Earth. See p57 of
    # https://www.usno.navy.mil/USNO/astronomical-applications/publications/Circular_179.pdf
    crepr = gcrs_coo2.cartesian.transform(rbpn)
    return tete_frame.realize_frame(crepr)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, TETE, GCRS)
def tete_to_gcrs(tete_coo, gcrs_frame):
    # Compute the pn matrix, and then multiply by its transpose.
    rbpn = erfa.pnm06a(*get_jd12(tete_coo.obstime, 'tt'))
    newrepr = tete_coo.cartesian.transform(matrix_transpose(rbpn))
    # We now have a GCRS vector for the input location and obstime.
    # Turn it into a GCRS frame instance.
    loc_gcrs = get_location_gcrs(tete_coo, rbpn)
    gcrs = loc_gcrs.realize_frame(newrepr)
    # Finally, do any needed offsets (no-op if same obstime and location)
    return gcrs.transform_to(gcrs_frame)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, TETE, ITRS)
def tete_to_itrs(tete_coo, itrs_frame):
    # first get us to TETE at the target obstime, and geocentric position
    tete_coo2 = tete_coo.transform_to(TETE(obstime=itrs_frame.obstime,
                                           location=EARTH_CENTER))

    # now get the pmatrix
    pmat = tete_to_itrs_mat(itrs_frame.obstime)
    crepr = tete_coo2.cartesian.transform(pmat)
    return itrs_frame.realize_frame(crepr)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ITRS, TETE)
def itrs_to_tete(itrs_coo, tete_frame):
    # compute the pmatrix, and then multiply by its transpose
    pmat = tete_to_itrs_mat(itrs_coo.obstime)
    newrepr = itrs_coo.cartesian.transform(matrix_transpose(pmat))
    tete = TETE(newrepr, obstime=itrs_coo.obstime)

    # now do any needed offsets (no-op if same obstime)
    return tete.transform_to(tete_frame)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, CIRS)
def gcrs_to_cirs(gcrs_coo, cirs_frame):
    # first get us to a 0 pos/vel GCRS at the target obstime
    gcrs_coo2 = gcrs_coo.transform_to(GCRS(obstime=cirs_frame.obstime))

    # now get the pmatrix
    pmat = gcrs_to_cirs_mat(cirs_frame.obstime)
    crepr = gcrs_coo2.cartesian.transform(pmat)
    return cirs_frame.realize_frame(crepr)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, CIRS, GCRS)
def cirs_to_gcrs(cirs_coo, gcrs_frame):
    # compute the pmatrix, and then multiply by its transpose
    pmat = gcrs_to_cirs_mat(cirs_coo.obstime)
    newrepr = cirs_coo.cartesian.transform(matrix_transpose(pmat))
    gcrs = GCRS(newrepr, obstime=cirs_coo.obstime)

    # now do any needed offsets (no-op if same obstime and 0 pos/vel)
    return gcrs.transform_to(gcrs_frame)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, CIRS, ITRS)
def cirs_to_itrs(cirs_coo, itrs_frame):
    # first get us to CIRS at the target obstime
    cirs_coo2 = cirs_coo.transform_to(CIRS(obstime=itrs_frame.obstime))

    # now get the pmatrix
    pmat = cirs_to_itrs_mat(itrs_frame.obstime)
    crepr = cirs_coo2.cartesian.transform(pmat)
    return itrs_frame.realize_frame(crepr)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ITRS, CIRS)
def itrs_to_cirs(itrs_coo, cirs_frame):
    # compute the pmatrix, and then multiply by its transpose
    pmat = cirs_to_itrs_mat(itrs_coo.obstime)
    newrepr = itrs_coo.cartesian.transform(matrix_transpose(pmat))
    cirs = CIRS(newrepr, obstime=itrs_coo.obstime)

    # now do any needed offsets (no-op if same obstime)
    return cirs.transform_to(cirs_frame)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ITRS, ITRS)
def itrs_to_itrs(from_coo, to_frame):
    # this self-transform goes through CIRS right now, which implicitly also
    # goes back to ICRS
    return from_coo.transform_to(CIRS()).transform_to(to_frame)

# TODO: implement GCRS<->CIRS if there's call for it.  The thing that's awkward
# is that they both have obstimes, so an extra set of transformations are necessary.
# so unless there's a specific need for that, better to just have it go through the above
# two steps anyway


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, GCRS, PrecessedGeocentric)
def gcrs_to_precessedgeo(from_coo, to_frame):
    # first get us to GCRS with the right attributes (might be a no-op)
    gcrs_coo = from_coo.transform_to(GCRS(obstime=to_frame.obstime,
                                          obsgeoloc=to_frame.obsgeoloc,
                                          obsgeovel=to_frame.obsgeovel))

    # now precess to the requested equinox
    pmat = gcrs_precession_mat(to_frame.equinox)
    crepr = gcrs_coo.cartesian.transform(pmat)
    return to_frame.realize_frame(crepr)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, PrecessedGeocentric, GCRS)
def precessedgeo_to_gcrs(from_coo, to_frame):
    # first un-precess
    pmat = gcrs_precession_mat(from_coo.equinox)
    crepr = from_coo.cartesian.transform(matrix_transpose(pmat))
    gcrs_coo = GCRS(crepr,
                    obstime=to_frame.obstime,
                    obsgeoloc=to_frame.obsgeoloc,
                    obsgeovel=to_frame.obsgeovel)

    # then move to the GCRS that's actually desired
    return gcrs_coo.transform_to(to_frame)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, TEME, ITRS)
def teme_to_itrs(teme_coo, itrs_frame):
    # first get us to TEME at the target obstime
    # TODO: self transform?
    teme_coo2 = teme_coo.transform_to(TEME(obstime=itrs_frame.obstime))

    # now get the pmatrix
    pmat = teme_to_itrs_mat(itrs_frame.obstime)
    crepr = teme_coo2.cartesian.transform(pmat)
    return itrs_frame.realize_frame(crepr)


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, ITRS, TEME)
def itrs_to_teme(itrs_coo, teme_frame):
    # compute the pmatrix, and then multiply by its transpose
    pmat = teme_to_itrs_mat(itrs_coo.obstime)
    newrepr = itrs_coo.cartesian.transform(matrix_transpose(pmat))
    teme = TEME(newrepr, obstime=itrs_coo.obstime)

    # now do any needed offsets (no-op if same obstime)
    return teme.transform_to(teme_frame)
