"""Functions for gnomonic projections"""

# This code originates from https://gitlab.in2p3.fr/betoule/imageproc by Marc
# Betoule.

import numpy as np

deg2rad = np.pi / 180.0


def dot3(x, M):
    """return the xMx vector where x is 2xN, M is 2x2xN"""
    xMx = np.matmul(x.T[:, None, :], np.matmul(M.T, x.T[:, :, None]))
    return np.squeeze(xMx)


def inverse_gnomonic_projection(
    x_rad, y_rad, rac_rad, decc_rad, vx_rad=None, vy_rad=None, vxy_rad=None
):
    """ """
    rho = np.sqrt(x_rad**2 + y_rad**2)
    c = np.arctan(rho)
    sinc_over_rho = np.where(rho == 0, 1, np.sin(c) / rho)
    cde, sde = np.cos(decc_rad[None, :]), np.sin(decc_rad[None, :])
    gamma = rho * cde * np.cos(c) - y_rad * sde * np.sin(c)

    dec_rad = np.arcsin(np.cos(c) * sde + y_rad * sinc_over_rho * cde)
    ra_rad = rac_rad[None, :] + np.arctan2(x_rad * np.sin(c), gamma)

    if vx_rad is not None and vy_rad is not None:
        dra, ddec = inverse_gnomonic_dradec(x_rad, y_rad, rac_rad, decc_rad)
        if vxy_rad is None:
            vxy_rad = np.zeros(len(vx_rad))
        era = np.sqrt(dot3(dra, np.array([(vx_rad, vxy_rad), (vxy_rad, vy_rad)])))
        edec = np.sqrt(dot3(ddec, np.array([(vx_rad, vxy_rad), (vxy_rad, vy_rad)])))
        return ra_rad.flatten(), dec_rad.flatten(), era, edec
    else:
        return ra_rad.flatten(), dec_rad.flatten()


def inverse_gnomonic_dcenter(x_rad, y_rad, rac_rad, decc_rad):
    """x, y in rad
    return drac, ddecc in rad
    """
    rho = np.sqrt(x_rad**2 + y_rad**2)
    c = np.arctan(rho)
    # sinc_over_rho = np.where(rho==0, 1,np.sin(c)/rho)
    cde, sde = np.cos(decc_rad[None, :]), np.sin(decc_rad[None, :])
    gamma = rho * cde * np.cos(c) - y_rad * sde * np.sin(c)
    dgammaddecc = -rho * sde * np.cos(c) - y_rad * cde * np.sin(c)

    v2 = np.sqrt(1 - (np.cos(c) * sde + (y_rad * np.sin(c) * cde) / rho) ** 2)
    g2 = 1 + (x_rad * np.sin(c) / gamma) ** 2

    ddecddecc = np.cos(c) * cde - (y_rad * np.sin(c) * sde) / rho
    ddecddecc /= v2
    draddecc = -x_rad * np.sin(c) * dgammaddecc / gamma**2
    draddecc /= g2
    ddecdrac = np.zeros((len(draddecc.flatten())), dtype=float)
    draddrac = np.ones((len(draddecc.flatten())), dtype=float)
    return np.array([draddrac, draddecc.flatten()]), np.array(
        [ddecdrac, ddecddecc.flatten()]
    )


def inverse_gnomonic_dradec(x_rad, y_rad, rac_rad, decc_rad):
    """x, y in rad
    return dra, ddec in rad
    """
    rho = np.sqrt(x_rad**2 + y_rad**2)
    c = np.arctan(rho)
    # sinc_over_rho = np.where(rho==0, 1,np.sin(c)/rho)
    cde, sde = np.cos(decc_rad[None, :]), np.sin(decc_rad[None, :])
    gamma = rho * cde * np.cos(c) - y_rad * sde * np.sin(c)
    drhodx, drhody = x_rad / rho, y_rad / rho
    dirhodx, dirhody = -x_rad / rho**3, -y_rad / rho**3
    dcdx, dcdy = 1 / (1 + rho**2) * drhodx, 1 / (1 + rho**2) * drhody
    v2 = np.sqrt(1 - (np.cos(c) * sde + (y_rad * np.sin(c) * cde) / rho) ** 2)

    ddecdx = -np.sin(c) * sde * dcdx
    ddecdx += (y_rad * np.cos(c) * cde / rho) * dcdx
    ddecdx += y_rad * np.sin(c) * cde * dirhodx
    ddecdy = -np.sin(c) * sde * dcdy
    ddecdy += cde / rho * (y_rad * np.cos(c) * dcdy + np.sin(c))
    ddecdy += y_rad * np.sin(c) * cde * dirhody
    ddec = np.array([(ddecdx / v2).flatten(), (ddecdy / v2).flatten()])

    dgammadx = drhodx * cde * np.cos(c)
    dgammadx -= rho * cde * np.sin(c) * dcdx
    dgammadx -= y_rad * sde * np.cos(c) * dcdx
    dgammady = drhody * cde * np.cos(c)
    dgammady -= rho * cde * np.sin(c) * dcdy
    dgammady -= y_rad * sde * np.cos(c) * dcdy
    dgammady -= sde * np.sin(c)
    g2 = 1 + (x_rad * np.sin(c) / gamma) ** 2

    dradx = (np.sin(c) + dcdx * x_rad * np.cos(c)) / gamma
    dradx -= x_rad * np.sin(c) * dgammadx / (gamma**2)
    drady = x_rad * np.cos(c) * dcdy / gamma
    drady -= x_rad * np.sin(c) * dgammady / (gamma**2)
    dra = np.array([(dradx / g2).flatten(), (drady / g2).flatten()])
    return dra, ddec


def gnomonic_dxy(ra_rad, dec_rad, rac_rad, decc_rad):
    """Compute the Jacobian of the gnomonic projection w.r.t. (ra, dec).

    For each source position, evaluates the partial derivatives of the
    tangent-plane coordinates (x, y) with respect to right ascension and
    declination. The result can be used to propagate positional uncertainties
    or proper motions into the tangent plane via a linear transformation.

    Parameters
    ----------
    ra_rad : array_like
        Right ascension of the source(s), in radians.
    dec_rad : array_like
        Declination of the source(s), in radians.
    rac_rad : array_like
        Right ascension of the tangent point(s) (projection center), in radians.
    decc_rad : array_like
        Declination of the tangent point(s) (projection center), in radians.

    Returns
    -------
    dx : ndarray, shape (2, N)
        Jacobian rows for the x coordinate. ``dx[0]`` = ∂x/∂ra and
        ``dx[1]`` = ∂x/∂dec, both flattened over all (source, center) pairs.
        Units are rad / rad (dimensionless).
    dy : ndarray, shape (2, N)
        Jacobian rows for the y coordinate. ``dy[0]`` = ∂y/∂ra and
        ``dy[1]`` = ∂y/∂dec, both flattened over all (source, center) pairs.
        Units are rad / rad (dimensionless).
    """
    s, c = np.sin(decc_rad[None, :]), np.cos(decc_rad[None, :])
    c2 = np.cos(ra_rad - rac_rad[None, :])
    d = np.sin(dec_rad) * s + np.cos(dec_rad) * c * c2

    ddr = -np.cos(dec_rad) * c * np.sin(ra_rad - rac_rad[None, :])
    ddd = s * np.cos(dec_rad) - c * np.sin(dec_rad) * c2

    dxdr = (
        np.cos(dec_rad) * c2 * d
        - (np.cos(dec_rad) * np.sin(ra_rad - rac_rad[None, :])) * ddr
    )

    dxdd = -np.sin(dec_rad) * np.sin(ra_rad - rac_rad[None, :]) * d
    dxdd -= (np.cos(dec_rad) * np.sin(ra_rad - rac_rad[None, :])) * ddd
    dx = np.array([(dxdr / d**2).flatten(), (dxdd / d**2).flatten()])

    dydr = np.cos(dec_rad) * s * np.sin(ra_rad - rac_rad[None, :]) * d
    dydr -= (np.sin(dec_rad) * c - np.cos(dec_rad) * s * c2) * ddr
    dydd = (np.cos(dec_rad) * c + np.sin(dec_rad) * s * c2) * d
    dydd -= (np.sin(dec_rad) * c - np.cos(dec_rad) * s * c2) * ddd
    dy = np.array([(dydr / d**2).flatten(), (dydd / d**2).flatten()])
    return dx, dy


def gnomonic_projection(
    ra_rad, dec_rad, rac_rad, decc_rad, era_rad=None, edec_rad=None, astropy=False
):
    """Project celestial coordinates onto the tangent plane (gnomonic projection).

    Maps spherical coordinates (ra, dec) to flat tangent-plane coordinates
    (x, y) using the gnomonic projection centered at (rac, decc):

        d = sin(dec)*sin(decc) + cos(dec)*cos(decc)*cos(ra - rac)
        x = cos(dec)*sin(ra - rac) / d
        y = (sin(dec)*cos(decc) - cos(dec)*sin(decc)*cos(ra - rac)) / d

    Optionally propagates uncorrelated position uncertainties into the
    tangent plane using the projection Jacobian from `gnomonic_dxy`.

    Parameters
    ----------
    ra_rad : array_like
        Right ascension of the source(s), in radians.
    dec_rad : array_like
        Declination of the source(s), in radians.
    rac_rad : array_like
        Right ascension of the tangent point(s) (projection center), in radians.
    decc_rad : array_like
        Declination of the tangent point(s) (projection center), in radians.
    era_rad : array_like, optional
        1-sigma uncertainty on right ascension, in radians. Must be provided
        together with `edec_rad` to trigger uncertainty propagation.
        Ignored when `astropy=True`.
    edec_rad : array_like, optional
        1-sigma uncertainty on declination, in radians. Must be provided
        together with `era_rad` to trigger uncertainty propagation.
        Ignored when `astropy=True`.
    astropy : bool, optional
        If ``True``, use the sign convention consistent with astropy
        (x = cos(decc)*sin(rac - ra) / d), which reverses the x-axis
        relative to the default. Uncertainty propagation is disabled in
        this mode. Default is ``False``.

    Returns
    -------
    x : ndarray
        Tangent-plane x-coordinate, in radians.
    y : ndarray
        Tangent-plane y-coordinate, in radians.
    ex : ndarray
        1-sigma uncertainty on x, in radians. Only returned when both
        `era_rad` and `edec_rad` are provided and `astropy` is ``False``.
    ey : ndarray
        1-sigma uncertainty on y, in radians. Only returned when both
        `era_rad` and `edec_rad` are provided and `astropy` is ``False``.
    """
    if astropy:
        s, c = np.sin(dec_rad), np.cos(dec_rad)
        c2 = np.cos(rac_rad[None, :] - ra_rad)
        d = np.sin(decc_rad[None, :]) * s + np.cos(decc_rad[None, :]) * c * c2
        x = (np.cos(decc_rad[None, :]) * np.sin(rac_rad[None, :] - ra_rad)) / d
        y = (np.sin(decc_rad[None, :]) * c - np.cos(decc_rad[None, :]) * s * c2) / d
    else:
        if len(dec_rad.shape) == 2 and dec_rad.shape[1] == 4:
            de, re = decc_rad[:, None], rac_rad[:, None]
        else:
            de, re = decc_rad[None, :], rac_rad[None, :]

        s, c = np.sin(de), np.cos(de)
        c2 = np.cos(ra_rad - re)
        d = np.sin(dec_rad) * s + np.cos(dec_rad) * c * c2
        x = (np.cos(dec_rad) * np.sin(ra_rad - re)) / d
        y = (np.sin(dec_rad) * c - np.cos(dec_rad) * s * c2) / d

    if era_rad is not None and edec_rad is not None and not astropy:
        dx, dy = gnomonic_dxy(ra_rad, dec_rad, rac_rad, decc_rad)
        z = np.zeros(len(era_rad))
        ex = np.sqrt(dot3(dx, np.array([(era_rad**2, z), (z, edec_rad**2)])))
        ey = np.sqrt(dot3(dy, np.array([(era_rad**2, z), (z, edec_rad**2)])))
        return x, y, ex, ey
    else:
        return x, y  # x.flatten(), y.flatten()


def tangent_plane_projection(ra, dec, ra_center, dec_center, pmra, pmdec):
    """Project celestial coordinates and proper motions onto the tangent plane.

    Performs a gnomonic (tangent plane) projection of sky positions centered at
    (ra_center, dec_center), and propagates proper motion vectors into the
    tangent plane coordinate system using the Jacobian of the projection.

    Parameters
    ----------
    ra : array_like
        Right ascension of the source(s) to project, in degrees.
    dec : array_like
        Declination of the source(s) to project, in degrees.
    ra_center : array_like
        Right ascension of the tangent point (projection center), in degrees.
    dec_center : array_like
        Declination of the tangent point (projection center), in degrees.
    pmra : array_like
        Proper motion in right ascension, in degrees. Should already include
        the cos(dec) factor if applicable.
    pmdec : array_like
        Proper motion in declination, in degrees.

    Returns
    -------
    tpx : ndarray
        Tangent plane x-coordinate of the projected position, in radians.
    tpy : ndarray
        Tangent plane y-coordinate of the projected position, in radians.
    pmtpx : ndarray
        Proper motion projected onto the tangent plane x-axis, in radians.
        Computed as the dot product of the gnomonic Jacobian row dx/d(ra,dec)
        with the proper motion vector (pmra, pmdec).
    pmtpy : ndarray
        Proper motion projected onto the tangent plane y-axis, in radians.
        Computed as the dot product of the gnomonic Jacobian row dy/d(ra,dec)
        with the proper motion vector (pmra, pmdec).
    """
    ra = np.deg2rad(np.asarray(ra))
    dec = np.deg2rad(np.asarray(dec))
    ra_center = np.deg2rad(np.atleast_1d(ra_center))
    dec_center = np.deg2rad(np.atleast_1d(dec_center))
    pmra = np.deg2rad(np.asarray(pmra))
    pmdec = np.deg2rad(np.asarray(pmdec))

    tpx, tpy = gnomonic_projection(ra, dec, ra_center, dec_center)
    tpdx, tpdy = gnomonic_dxy(ra, dec, ra_center, dec_center)

    tpdx = tpdx.T.reshape(-1, 1, 2)
    tpdy = tpdy.T.reshape(-1, 1, 2)
    pm = np.array([pmra, pmdec])
    pm = pm.T.reshape(-1, 2, 1)
    pmtpx = tpdx @ pm
    pmtpy = tpdy @ pm

    return tpx[0], tpy[0], pmtpx.squeeze(), pmtpy.squeeze()
