"""Coordinate transformation functions for astronomical calculations.

This module provides functions to transform between different coordinate
systems used in astronomy: rectangular/spherical, ecliptic/equatorial, etc.
"""

import numpy as np
from typing import Tuple, Union


def spherical_to_rectangular(
    lon: Union[float, np.ndarray],
    lat: Union[float, np.ndarray],
    r: Union[float, np.ndarray]
) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray], Union[float, np.ndarray]]:
    """Convert spherical coordinates to rectangular coordinates (vectorized)

    Parameters
    ----------
    lon : float or np.ndarray
        Longitude in degrees (scalar or array).
    lat : float or np.ndarray
        Latitude in degrees (scalar or array).
    r : float or np.ndarray
        Distance (scalar or array).

    Returns
    -------
    tuple of (float or np.ndarray, float or np.ndarray, float or np.ndarray)
        Tuple of (x, y, z) in same units as r (scalar or array).

    Notes
    -----
    This function is vectorized and works with both scalars and arrays.
    """
    lon_rad = np.deg2rad(lon)
    lat_rad = np.deg2rad(lat)

    x = r * np.cos(lon_rad) * np.cos(lat_rad)
    y = r * np.sin(lon_rad) * np.cos(lat_rad)
    z = r * np.sin(lat_rad)

    return x, y, z


def rectangular_to_spherical(
    x: Union[float, np.ndarray],
    y: Union[float, np.ndarray],
    z: Union[float, np.ndarray]
) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray], Union[float, np.ndarray]]:
    """Convert rectangular coordinates to spherical coordinates (vectorized)

    Parameters
    ----------
    x : float or np.ndarray
        X coordinate (scalar or array).
    y : float or np.ndarray
        Y coordinate (scalar or array).
    z : float or np.ndarray
        Z coordinate (scalar or array).

    Returns
    -------
    tuple of (float or np.ndarray, float or np.ndarray, float or np.ndarray)
        Tuple of (lon, lat, r) where angles are in degrees (scalar or array).

    Notes
    -----
    This function is vectorized and works with both scalars and arrays.
    """
    r = np.sqrt(x**2 + y**2 + z**2)

    # Handle zero distance (works for both scalars and arrays)
    is_scalar = np.ndim(r) == 0
    if is_scalar:
        if r == 0:
            return 0.0, 0.0, 0.0
    else:
        # For arrays, avoid division by zero
        r = np.where(r == 0, 1.0, r)  # Temp fix for division

    lon = np.rad2deg(np.arctan2(y, x))
    lat = np.rad2deg(np.arcsin(z / r))

    # Normalize longitude to 0-360
    lon = np.where(lon < 0, lon + 360.0, lon)

    return lon, lat, r


def ecliptic_to_equatorial(
    x: Union[float, np.ndarray],
    y: Union[float, np.ndarray],
    z: Union[float, np.ndarray],
    obliquity: Union[float, np.ndarray]
) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray], Union[float, np.ndarray]]:
    """Convert ecliptic coordinates to equatorial coordinates (vectorized)

    Parameters
    ----------
    x : float or np.ndarray
        Ecliptic X coordinate (scalar or array).
    y : float or np.ndarray
        Ecliptic Y coordinate (scalar or array).
    z : float or np.ndarray
        Ecliptic Z coordinate (scalar or array).
    obliquity : float or np.ndarray
        Obliquity of ecliptic in degrees (scalar or array).

    Returns
    -------
    tuple of (float or np.ndarray, float or np.ndarray, float or np.ndarray)
        Tuple of (x_eq, y_eq, z_eq) equatorial rectangular coordinates (scalar or array).

    Notes
    -----
    This function is vectorized and works with both scalars and arrays.
    """
    obl_rad = np.deg2rad(obliquity)
    cos_obl = np.cos(obl_rad)
    sin_obl = np.sin(obl_rad)

    x_eq = x
    y_eq = y * cos_obl - z * sin_obl
    z_eq = y * sin_obl + z * cos_obl

    return x_eq, y_eq, z_eq


def equatorial_to_ecliptic(x: float, y: float, z: float, obliquity: float) -> Tuple[float, float, float]:
    """Convert equatorial coordinates to ecliptic coordinates

    Parameters
    ----------
    x : float
        Equatorial X coordinate.
    y : float
        Equatorial Y coordinate.
    z : float
        Equatorial Z coordinate.
    obliquity : float
        Obliquity of ecliptic in degrees.

    Returns
    -------
    tuple of (float, float, float)
        Tuple of (x_ecl, y_ecl, z_ecl) ecliptic rectangular coordinates.
    """
    obl_rad = np.deg2rad(obliquity)
    cos_obl = np.cos(obl_rad)
    sin_obl = np.sin(obl_rad)

    x_ecl = x
    y_ecl = y * cos_obl + z * sin_obl
    z_ecl = -y * sin_obl + z * cos_obl

    return x_ecl, y_ecl, z_ecl


def heliocentric_to_geocentric(
    x_planet: Union[float, np.ndarray],
    y_planet: Union[float, np.ndarray],
    z_planet: Union[float, np.ndarray],
    x_earth: Union[float, np.ndarray],
    y_earth: Union[float, np.ndarray],
    z_earth: Union[float, np.ndarray]
) -> Tuple[Union[float, np.ndarray], Union[float, np.ndarray], Union[float, np.ndarray]]:
    """Convert heliocentric coordinates to geocentric coordinates (vectorized)

    Parameters
    ----------
    x_planet : float or np.ndarray
        Planet heliocentric X coordinate (scalar or array).
    y_planet : float or np.ndarray
        Planet heliocentric Y coordinate (scalar or array).
    z_planet : float or np.ndarray
        Planet heliocentric Z coordinate (scalar or array).
    x_earth : float or np.ndarray
        Earth heliocentric X coordinate (scalar or array).
    y_earth : float or np.ndarray
        Earth heliocentric Y coordinate (scalar or array).
    z_earth : float or np.ndarray
        Earth heliocentric Z coordinate (scalar or array).

    Returns
    -------
    tuple of (float or np.ndarray, float or np.ndarray, float or np.ndarray)
        Tuple of (x_geo, y_geo, z_geo) geocentric coordinates (scalar or array).
    """
    x_geo = x_planet - x_earth
    y_geo = y_planet - y_earth
    z_geo = z_planet - z_earth

    return x_geo, y_geo, z_geo


def geocentric_to_topocentric(
    lon: float, lat: float, dist: float, observer_lat: float, observer_lon: float, lst: float, height: float = 0.0
) -> Tuple[float, float, float]:
    """Convert geocentric coordinates to topocentric coordinates

    Parameters
    ----------
    lon : float
        Geocentric longitude in degrees.
    lat : float
        Geocentric latitude in degrees.
    dist : float
        Distance in AU.
    observer_lat : float
        Observer's latitude in degrees.
    observer_lon : float
        Observer's longitude in degrees.
    lst : float
        Local sidereal time in degrees.
    height : float, optional
        Observer's height above sea level in meters.

    Returns
    -------
    tuple of (float, float, float)
        Tuple of (az, alt, dist) where
        az is azimuth in degrees (0 = North, 90 = East),
        alt is altitude in degrees above horizon,
        dist is topocentric distance in AU.
    """
    # Earth's equatorial radius in AU
    earth_radius_au = 4.26352e-5

    # Observer's position
    obs_lat_rad = np.deg2rad(observer_lat)

    # Parallax factors
    rho_cos = np.cos(obs_lat_rad) + height / 6378140.0 * np.cos(obs_lat_rad)
    rho_sin = np.sin(obs_lat_rad) + height / 6378140.0 * np.sin(obs_lat_rad)

    # Convert to rectangular equatorial
    x, y, z = spherical_to_rectangular(lon, lat, dist)

    # Hour angle
    ha = lst - lon
    ha_rad = np.deg2rad(ha)

    # Topocentric rectangular coordinates
    dx = -earth_radius_au * rho_cos * np.cos(ha_rad)
    dy = -earth_radius_au * rho_cos * np.sin(ha_rad)
    dz = -earth_radius_au * rho_sin

    x_topo = x + dx
    y_topo = y + dy
    z_topo = z + dz

    # Convert to horizontal coordinates
    lat_rad = np.deg2rad(lat)

    # Rotate to horizon system
    x_hor = -x_topo * np.sin(ha_rad) + y_topo * np.cos(ha_rad)
    y_hor = (
        -x_topo * np.cos(ha_rad) * np.sin(obs_lat_rad)
        - y_topo * np.sin(ha_rad) * np.sin(obs_lat_rad)
        + z_topo * np.cos(obs_lat_rad)
    )
    z_hor = (
        x_topo * np.cos(ha_rad) * np.cos(obs_lat_rad)
        + y_topo * np.sin(ha_rad) * np.cos(obs_lat_rad)
        + z_topo * np.sin(obs_lat_rad)
    )

    # Convert to azimuth and altitude
    az = np.rad2deg(np.arctan2(x_hor, y_hor))
    alt = np.rad2deg(np.arcsin(z_hor / np.sqrt(x_hor**2 + y_hor**2 + z_hor**2)))

    # Normalize azimuth to 0-360
    if az < 0:
        az += 360.0

    # Topocentric distance
    dist_topo = np.sqrt(x_topo**2 + y_topo**2 + z_topo**2)

    return az, alt, dist_topo


def mean_obliquity(jd: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
    """Calculate mean obliquity of the ecliptic (vectorized)

    Parameters
    ----------
    jd : float or np.ndarray
        Julian Date (scalar or array).

    Returns
    -------
    float or np.ndarray
        Mean obliquity in degrees (scalar or array).

    Notes
    -----
    This function is vectorized and works with both scalars and arrays.
    """
    # Centuries since J2000.0
    T = (jd - 2451545.0) / 36525.0

    # IAU 2006 formula
    obliquity = (
        23.0
        + 26.0 / 60.0
        + 21.448 / 3600.0
        - 46.8150 / 3600.0 * T
        - 0.00059 / 3600.0 * T**2
        + 0.001813 / 3600.0 * T**3
    )

    return obliquity


def nutation(jd: float) -> Tuple[float, float]:
    """Calculate nutation in longitude and obliquity

    Parameters
    ----------
    jd : float
        Julian Date.

    Returns
    -------
    tuple of (float, float)
        Tuple of (nut_lon, nut_obl) in degrees.
    """
    # Days since J2000.0
    d = jd - 2451545.0

    # Mean elongation of Moon from Sun
    D = np.deg2rad(297.85036 + 445267.111480 * d / 36525.0)

    # Sun's mean anomaly
    M = np.deg2rad(357.52772 + 35999.050340 * d / 36525.0)

    # Moon's mean anomaly
    Mm = np.deg2rad(134.96298 + 477198.867398 * d / 36525.0)

    # Moon's argument of latitude
    F = np.deg2rad(93.27191 + 483202.017538 * d / 36525.0)

    # Longitude of ascending node
    omega = np.deg2rad(125.04452 - 1934.136261 * d / 36525.0)

    # Nutation in longitude (arcseconds)
    nut_lon_as = (
        -17.20 * np.sin(omega)
        - 1.32 * np.sin(2 * F - 2 * D + 2 * omega)
        - 0.23 * np.sin(2 * F + 2 * omega)
        + 0.21 * np.sin(2 * omega)
    )

    # Nutation in obliquity (arcseconds)
    nut_obl_as = (
        9.20 * np.cos(omega)
        + 0.57 * np.cos(2 * F - 2 * D + 2 * omega)
        + 0.10 * np.cos(2 * F + 2 * omega)
        - 0.09 * np.cos(2 * omega)
    )

    # Convert to degrees
    nut_lon = nut_lon_as / 3600.0
    nut_obl = nut_obl_as / 3600.0

    return nut_lon, nut_obl


def true_obliquity(jd: float) -> float:
    """Calculate true obliquity of the ecliptic

    Parameters
    ----------
    jd : float
        Julian Date.

    Returns
    -------
    float
        True obliquity in degrees.
    """
    mean_obl = mean_obliquity(jd)
    _, nut_obl = nutation(jd)

    return mean_obl + nut_obl


def aberration_correction(lon: float, lat: float, jd: float) -> Tuple[float, float]:
    """Apply stellar aberration correction

    Parameters
    ----------
    lon : float
        Longitude in degrees.
    lat : float
        Latitude in degrees.
    jd : float
        Julian Date.

    Returns
    -------
    tuple of (float, float)
        Tuple of (dlon, dlat) corrections in degrees.
    """
    # Days since J2000.0
    d = jd - 2451545.0

    # Earth's orbital eccentricity
    e = 0.016709 - 1.151e-9 * d

    # Sun's mean longitude
    L = np.deg2rad(280.460 + 0.9856474 * d)

    # Aberration constant (arcseconds)
    k = 20.49552

    # Corrections
    lon_rad = np.deg2rad(lon)
    lat_rad = np.deg2rad(lat)

    dlon_as = -k * np.cos(L - lon_rad) / np.cos(lat_rad) + k * e * np.cos(
        np.deg2rad(102.93735 + 0.71953 * d) - lon_rad
    ) / np.cos(lat_rad)

    dlat_as = -k * np.sin(lat_rad) * (np.sin(L - lon_rad) - e * np.sin(np.deg2rad(102.93735 + 0.71953 * d) - lon_rad))

    # Convert to degrees
    dlon = dlon_as / 3600.0
    dlat = dlat_as / 3600.0

    return dlon, dlat
