# Licensed under a 3-clause BSD style license - see LICENSE.rst

import math

import erfa
import numpy as np

from astropy import units as u
from astropy.coordinates.attributes import TimeAttribute
from astropy.coordinates.baseframe import base_doc, frame_transform_graph
from astropy.coordinates.matrix_utilities import rotation_matrix
from astropy.coordinates.representation import (
    CartesianRepresentation,
    UnitSphericalRepresentation,
)
from astropy.coordinates.transformations import (
    DynamicMatrixTransform,
    FunctionTransformWithFiniteDifference,
)
from astropy.time import Time
from astropy.utils.compat import COPY_IF_NEEDED
from astropy.utils.decorators import format_doc

from .baseradec import BaseRADecFrame, doc_components
from .utils import EQUINOX_B1950

__all__ = ["FK4", "FK4NoETerms"]

jd1950 = Time("B1950").jd

doc_footer_fk4 = """
    Other parameters
    ----------------
    equinox : `~astropy.time.Time`
        The equinox of this frame.
    obstime : `~astropy.time.Time`
        The time this frame was observed.  If ``None``, will be the same as
        ``equinox``.
"""


@format_doc(base_doc, components=doc_components, footer=doc_footer_fk4)
class FK4(BaseRADecFrame):
    """
    A coordinate or frame in the FK4 system.

    Note that this is a barycentric version of FK4 - that is, the origin for
    this frame is the Solar System Barycenter, *not* the Earth geocenter.

    The frame attributes are listed under **Other Parameters**.
    """

    equinox = TimeAttribute(default=EQUINOX_B1950, doc="The equinox time")
    obstime = TimeAttribute(
        default=None,
        secondary_attribute="equinox",
        doc="The reference time (e.g., time of observation)",
    )


# the "self" transform


@frame_transform_graph.transform(FunctionTransformWithFiniteDifference, FK4, FK4)
def fk4_to_fk4(fk4coord1, fk4frame2):
    # deceptively complicated: need to transform to No E-terms FK4, precess, and
    # then come back, because precession is non-trivial with E-terms
    fnoe_w_eqx1 = fk4coord1.transform_to(FK4NoETerms(equinox=fk4coord1.equinox))
    fnoe_w_eqx2 = fnoe_w_eqx1.transform_to(FK4NoETerms(equinox=fk4frame2.equinox))
    return fnoe_w_eqx2.transform_to(fk4frame2)


@format_doc(base_doc, components=doc_components, footer=doc_footer_fk4)
class FK4NoETerms(BaseRADecFrame):
    """
    A coordinate or frame in the FK4 system, but with the E-terms of aberration
    removed.

    The frame attributes are listed under **Other Parameters**.
    """

    equinox = TimeAttribute(default=EQUINOX_B1950, doc="The equinox time")
    obstime = TimeAttribute(
        default=None,
        secondary_attribute="equinox",
        doc="The reference time (e.g., time of observation)",
    )

    @staticmethod
    def _precession_matrix(oldequinox, newequinox):
        """
        Compute and return the precession matrix for FK4 using Newcomb's method.
        Used inside some of the transformation functions.

        Parameters
        ----------
        oldequinox : `~astropy.time.Time`
            The equinox to precess from.
        newequinox : `~astropy.time.Time`
            The equinox to precess to.

        Returns
        -------
        newcoord : array
            The precession matrix to transform to the new equinox
        """
        # tropical years
        t1 = (oldequinox.byear - 1850.0) / 1000.0
        dt = (newequinox.byear - 1850.0) / 1000.0 - t1
        dt_over_3600 = dt / 3600

        z1 = zeta1 = (0.060 * t1 + 139.720) * t1 + 23035.545
        zeta2 = -0.27 * t1 + 30.240
        zeta = ((17.995 * dt + zeta2) * dt + zeta1) * dt_over_3600

        z2 = 109.480 + 0.39 * t1
        z = ((18.325 * dt + z2) * dt + z1) * dt_over_3600

        theta1 = (-0.37 * t1 - 85.29) * t1 + 20051.12
        theta2 = -0.37 * t1 - 42.65
        theta = ((-41.8 * dt + theta2) * dt + theta1) * dt_over_3600

        return (
            rotation_matrix(-z, "z")
            @ rotation_matrix(theta, "y")
            @ rotation_matrix(-zeta, "z")
        )


# the "self" transform


@frame_transform_graph.transform(DynamicMatrixTransform, FK4NoETerms, FK4NoETerms)
def fk4noe_to_fk4noe(fk4necoord1, fk4neframe2):
    return fk4necoord1._precession_matrix(fk4necoord1.equinox, fk4neframe2.equinox)


# FK4-NO-E to/from FK4 ----------------------------->
# Unlike other frames, this module include *two* frame classes for FK4
# coordinates - one including the E-terms of aberration (FK4), and
# one not including them (FK4NoETerms). The following functions
# implement the transformation between these two.
def fk4_e_terms(equinox):
    """
    Return the e-terms of aberration vector.

    Parameters
    ----------
    equinox : Time object
        The equinox for which to compute the e-terms
    """
    # Constant of aberration at J2000; from Explanatory Supplement to the
    # Astronomical Almanac (Seidelmann, 2005).
    k = 0.0056932  # in degrees (v_earth/c ~ 1e-4 rad ~ 0.0057 deg)

    # Explanatory Supplement to the Astronomical Almanac: P. Kenneth
    #  Seidelmann (ed), University Science Books (1992).
    T = (equinox.jd - jd1950) / 36525.0
    # Eccentricity of the Earth's orbit
    ek = math.radians(k) * ((-0.000000126 * T - 0.00004193) * T + 0.01673011)
    # Mean longitude of perigee of the solar orbit
    g = np.radians((((0.012 * T + 1.65) * T + 6190.67) * T + 1015489.951) / 3600.0)
    minus_ek_cos_g = -ek * np.cos(g)
    # Obliquity of the ecliptic
    o = erfa.obl80(equinox.jd, 0)

    return (ek * np.sin(g), minus_ek_cos_g * np.cos(o), minus_ek_cos_g * np.sin(o))


@frame_transform_graph.transform(
    FunctionTransformWithFiniteDifference, FK4, FK4NoETerms
)
def fk4_to_fk4_no_e(fk4coord, fk4noeframe):
    # Extract cartesian vector
    rep = fk4coord.cartesian

    # Find distance (for re-normalization)
    d_orig = rep.norm()
    rep /= d_orig

    # Apply E-terms of aberration. Note that this depends on the equinox (not
    # the observing time/epoch) of the coordinates. See issue #1496 for a
    # discussion of this.
    eterms_a = CartesianRepresentation(
        u.Quantity(
            fk4_e_terms(fk4coord.equinox),
            u.dimensionless_unscaled,
            copy=COPY_IF_NEEDED,
        ),
        copy=False,
    )
    rep = rep - eterms_a + eterms_a.dot(rep) * rep

    # Find new distance (for re-normalization)
    d_new = rep.norm()

    # Renormalize
    rep *= d_orig / d_new

    # now re-cast into an appropriate Representation, and precess if need be
    if isinstance(fk4coord.data, UnitSphericalRepresentation):
        rep = rep.represent_as(UnitSphericalRepresentation)

    # if no obstime was given in the new frame, use the old one for consistency
    newobstime = (
        fk4coord._obstime if fk4noeframe._obstime is None else fk4noeframe._obstime
    )

    fk4noe = FK4NoETerms(rep, equinox=fk4coord.equinox, obstime=newobstime)
    if fk4coord.equinox != fk4noeframe.equinox:
        # precession
        fk4noe = fk4noe.transform_to(fk4noeframe)
    return fk4noe


@frame_transform_graph.transform(
    FunctionTransformWithFiniteDifference, FK4NoETerms, FK4
)
def fk4_no_e_to_fk4(fk4noecoord, fk4frame):
    # first precess, if necessary
    if fk4noecoord.equinox != fk4frame.equinox:
        fk4noe_w_fk4equinox = FK4NoETerms(
            equinox=fk4frame.equinox, obstime=fk4noecoord.obstime
        )
        fk4noecoord = fk4noecoord.transform_to(fk4noe_w_fk4equinox)

    # Extract cartesian vector
    rep = fk4noecoord.cartesian

    # Find distance (for re-normalization)
    d_orig = rep.norm()
    rep /= d_orig

    # Apply E-terms of aberration. Note that this depends on the equinox (not
    # the observing time/epoch) of the coordinates. See issue #1496 for a
    # discussion of this.
    eterms_a = CartesianRepresentation(
        u.Quantity(
            fk4_e_terms(fk4noecoord.equinox),
            u.dimensionless_unscaled,
            copy=COPY_IF_NEEDED,
        ),
        copy=False,
    )

    eterms_a_plus_rep_ini = eterms_a + rep
    for _ in range(10):
        rep = eterms_a_plus_rep_ini / (1.0 + eterms_a.dot(rep))

    # Find new distance (for re-normalization)
    d_new = rep.norm()

    # Renormalize
    rep *= d_orig / d_new

    # now re-cast into an appropriate Representation, and precess if need be
    if isinstance(fk4noecoord.data, UnitSphericalRepresentation):
        rep = rep.represent_as(UnitSphericalRepresentation)

    return fk4frame.realize_frame(rep)
