lox-space 0.1.0-alpha.48

// SPDX-FileCopyrightText: 2025 Helge Eichhorn <git@helgeeichhorn.de>
//
// SPDX-License-Identifier: MPL-2.0

use crate::bodies::dynamic::{DynOrigin, UnknownOriginId, UnknownOriginName};
use crate::bodies::{
    Origin, TryMeanRadius, TryPointMass, TryRotationalElements, TrySpheroid, TryTriaxialEllipsoid,
};
use crate::units::python::{PyAngle, PyAngularRate, PyDistance, PyGravitationalParameter};
use lox_core::types::units::Seconds;
use lox_units::{Angle, AngularRate};
use pyo3::create_exception;
use pyo3::exceptions::{PyException, PyTypeError, PyValueError};
use pyo3::prelude::*;
use std::str::FromStr;

create_exception!(
    lox_space,
    UndefinedOriginPropertyError,
    PyException,
    "Python exception raised when a body property is not defined for a given origin."
);

/// PyO3 error wrapper for [`lox_bodies::UndefinedOriginPropertyError`].
pub struct PyUndefinedOriginPropertyError(pub crate::bodies::UndefinedOriginPropertyError);

impl From<PyUndefinedOriginPropertyError> for PyErr {
    fn from(err: PyUndefinedOriginPropertyError) -> Self {
        UndefinedOriginPropertyError::new_err(err.0.to_string())
    }
}

/// PyO3 error wrapper for [`lox_bodies::dynamic::UnknownOriginId`].
pub struct PyUnknownOriginId(pub UnknownOriginId);

impl From<PyUnknownOriginId> for PyErr {
    fn from(err: PyUnknownOriginId) -> Self {
        PyValueError::new_err(err.0.to_string())
    }
}

struct PyUnknownOriginName(UnknownOriginName);

impl From<PyUnknownOriginName> for PyErr {
    fn from(err: PyUnknownOriginName) -> Self {
        PyValueError::new_err(err.0.to_string())
    }
}

/// Represents a celestial body (planet, moon, barycenter, etc.).
///
/// Origin objects represent celestial bodies using NAIF/SPICE identifiers.
/// They provide access to physical properties such as gravitational parameters,
/// radii, and rotational elements.
///
/// Args:
///     origin: Body name (e.g., "Earth", "Moon") or NAIF ID (e.g., 399 for Earth).
///
/// Raises:
///     ValueError: If the origin name or ID is not recognized.
///     TypeError: If the argument is neither a string nor an integer.
#[pyclass(name = "Origin", module = "lox_space", frozen, eq, from_py_object)]
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct PyOrigin(pub DynOrigin);

#[pymethods]
impl PyOrigin {
    #[new]
    fn new(origin: &Bound<'_, PyAny>) -> PyResult<Self> {
        if let Ok(origin) = origin.extract::<i32>() {
            return Ok(Self(origin.try_into().map_err(PyUnknownOriginId)?));
        }
        if let Ok(origin) = origin.extract::<&str>() {
            return Ok(Self(
                DynOrigin::from_str(origin).map_err(PyUnknownOriginName)?,
            ));
        }
        Err(PyTypeError::new_err(
            "`origin` must be either a string or an integer",
        ))
    }

    /// Return the string representation of this origin.
    pub fn __repr__(&self) -> String {
        format!("Origin(\"{}\")", self.name())
    }

    fn __str__(&self) -> &str {
        self.name()
    }

    fn __getnewargs__(&self) -> (&str,) {
        (self.name(),)
    }

    /// Return the NAIF ID of this body.
    pub fn id(&self) -> i32 {
        self.0.id().0
    }

    /// Return the name of this body.
    pub fn name(&self) -> &'static str {
        self.0.name()
    }

    /// Return the gravitational parameter (GM).
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn gravitational_parameter(&self) -> PyResult<PyGravitationalParameter> {
        let gp = self
            .0
            .try_gravitational_parameter()
            .map_err(PyUndefinedOriginPropertyError)?;
        Ok(PyGravitationalParameter(gp))
    }

    /// Return the mean radius.
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn mean_radius(&self) -> PyResult<PyDistance> {
        Ok(PyDistance(
            self.0
                .try_mean_radius()
                .map_err(PyUndefinedOriginPropertyError)?,
        ))
    }

    /// Return the triaxial radii (x, y, z).
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn radii(&self) -> PyResult<(PyDistance, PyDistance, PyDistance)> {
        let (a, b, c) = self.0.try_radii().map_err(PyUndefinedOriginPropertyError)?;
        Ok((PyDistance(a), PyDistance(b), PyDistance(c)))
    }

    /// Return the equatorial radius.
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn equatorial_radius(&self) -> PyResult<PyDistance> {
        Ok(PyDistance(
            self.0
                .try_equatorial_radius()
                .map_err(PyUndefinedOriginPropertyError)?,
        ))
    }

    /// Return the polar radius.
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn polar_radius(&self) -> PyResult<PyDistance> {
        Ok(PyDistance(
            self.0
                .try_polar_radius()
                .map_err(PyUndefinedOriginPropertyError)?,
        ))
    }

    /// Return rotational elements (right ascension, declination, rotation angle).
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    ///
    /// Returns:
    ///     Tuple of (right_ascension, declination, rotation_angle) as Angles.
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn rotational_elements(&self, et: Seconds) -> PyResult<(PyAngle, PyAngle, PyAngle)> {
        let (ra, dec, rot) = self
            .0
            .try_rotational_elements(et)
            .map_err(PyUndefinedOriginPropertyError)?;
        Ok((
            PyAngle(Angle::radians(ra)),
            PyAngle(Angle::radians(dec)),
            PyAngle(Angle::radians(rot)),
        ))
    }

    /// Return rotational element rates.
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    ///
    /// Returns:
    ///     Tuple of (ra_rate, dec_rate, rotation_rate) as AngularRates.
    ///
    /// Raises:
    ///     UndefinedOriginPropertyError: If not defined for this body.
    pub fn rotational_element_rates(
        &self,
        et: Seconds,
    ) -> PyResult<(PyAngularRate, PyAngularRate, PyAngularRate)> {
        let (ra_rate, dec_rate, rot_rate) = self
            .0
            .try_rotational_element_rates(et)
            .map_err(PyUndefinedOriginPropertyError)?;
        Ok((
            PyAngularRate(AngularRate::radians_per_second(ra_rate)),
            PyAngularRate(AngularRate::radians_per_second(dec_rate)),
            PyAngularRate(AngularRate::radians_per_second(rot_rate)),
        ))
    }

    /// Return the right ascension of the pole.
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    pub fn right_ascension(&self, et: Seconds) -> PyResult<PyAngle> {
        Ok(PyAngle(Angle::radians(
            self.0
                .try_right_ascension(et)
                .map_err(PyUndefinedOriginPropertyError)?,
        )))
    }

    /// Return the rate of change of right ascension.
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    pub fn right_ascension_rate(&self, et: Seconds) -> PyResult<PyAngularRate> {
        Ok(PyAngularRate(AngularRate::radians_per_second(
            self.0
                .try_right_ascension_rate(et)
                .map_err(PyUndefinedOriginPropertyError)?,
        )))
    }

    /// Return the declination of the pole.
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    pub fn declination(&self, et: Seconds) -> PyResult<PyAngle> {
        Ok(PyAngle(Angle::radians(
            self.0
                .try_declination(et)
                .map_err(PyUndefinedOriginPropertyError)?,
        )))
    }

    /// Return the rate of change of declination.
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    pub fn declination_rate(&self, et: Seconds) -> PyResult<PyAngularRate> {
        Ok(PyAngularRate(AngularRate::radians_per_second(
            self.0
                .try_declination_rate(et)
                .map_err(PyUndefinedOriginPropertyError)?,
        )))
    }

    /// Return the rotation angle (prime meridian).
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    pub fn rotation_angle(&self, et: Seconds) -> PyResult<PyAngle> {
        Ok(PyAngle(Angle::radians(
            self.0
                .try_rotation_angle(et)
                .map_err(PyUndefinedOriginPropertyError)?,
        )))
    }

    /// Return the rotation rate.
    ///
    /// Args:
    ///     et: Ephemeris time in seconds from J2000.
    pub fn rotation_rate(&self, et: Seconds) -> PyResult<PyAngularRate> {
        Ok(PyAngularRate(AngularRate::radians_per_second(
            self.0
                .try_rotation_rate(et)
                .map_err(PyUndefinedOriginPropertyError)?,
        )))
    }
}

impl TryFrom<&Bound<'_, PyAny>> for PyOrigin {
    type Error = PyErr;

    fn try_from(value: &Bound<'_, PyAny>) -> Result<Self, Self::Error> {
        if let Ok(origin) = value.extract::<PyOrigin>() {
            return Ok(origin);
        }
        if let Ok(name) = value.extract::<&str>() {
            return Ok(Self(
                DynOrigin::from_str(name).map_err(PyUnknownOriginName)?,
            ));
        }
        if let Ok(id) = value.extract::<i32>() {
            return Ok(Self(id.try_into().map_err(PyUnknownOriginId)?));
        }
        Err(PyTypeError::new_err(
            "'origin' argument must be a string, an integer, or an 'Origin' instance",
        ))
    }
}