use crate::bodies::{self, CelestialBody};
use crate::constants::*;
use crate::kepler::solve_kepler;
use nalgebra::Vector3;

/// Compute the heliocentric ecliptic position of a body at a given Julian Date.
///
/// Returns position in AU in the ecliptic coordinate frame:
/// - X toward vernal equinox
/// - Y in ecliptic plane, 90° from X
/// - Z toward ecliptic north pole
///
/// For moons, returns the heliocentric position (parent position + moon offset).
pub fn position_at_epoch(bodies: &[CelestialBody], body_id: usize, jd: f64) -> Vector3<f64> {
    let body = &bodies[body_id];

    // Sun is at origin
    if body_id == bodies::id::SUN {
        return Vector3::zeros();
    }

    // Centuries since J2000.0
    let t = (jd - J2000_JD) / DAYS_PER_CENTURY;

    // Compute current elements
    let a = body.elements.a + body.rates.a * t;
    let e = body.elements.e + body.rates.e * t;
    let i = (body.elements.i + body.rates.i * t) * DEG_TO_RAD;
    let l = (body.elements.l + body.rates.l * t) * DEG_TO_RAD;
    let w_bar = (body.elements.w_bar + body.rates.w_bar * t) * DEG_TO_RAD;
    let omega = (body.elements.omega + body.rates.omega * t) * DEG_TO_RAD;

    // Argument of perihelion
    let w = w_bar - omega;

    // Mean anomaly
    let m = l - w_bar;

    // Solve Kepler's equation for eccentric anomaly
    let ea = solve_kepler(m, e);

    // True anomaly
    let nu = 2.0 * ((1.0 + e).sqrt() * (ea / 2.0).sin())
        .atan2((1.0 - e).sqrt() * (ea / 2.0).cos());

    // Distance from focus
    let r = a * (1.0 - e * ea.cos());

    // Position in orbital plane
    let x_orb = r * nu.cos();
    let y_orb = r * nu.sin();

    // Rotate to ecliptic coordinates
    let cos_w = w.cos();
    let sin_w = w.sin();
    let cos_o = omega.cos();
    let sin_o = omega.sin();
    let cos_i = i.cos();
    let sin_i = i.sin();

    let x_ecl = (cos_w * cos_o - sin_w * sin_o * cos_i) * x_orb
        + (-sin_w * cos_o - cos_w * sin_o * cos_i) * y_orb;
    let y_ecl = (cos_w * sin_o + sin_w * cos_o * cos_i) * x_orb
        + (-sin_w * sin_o + cos_w * cos_o * cos_i) * y_orb;
    let z_ecl = (sin_w * sin_i) * x_orb + (cos_w * sin_i) * y_orb;

    let pos = Vector3::new(x_ecl, y_ecl, z_ecl);

    // For moons, add parent body position
    if let Some(parent_id) = bodies::parent_body(body_id) {
        let parent_pos = position_at_epoch(bodies, parent_id, jd);
        parent_pos + pos
    } else {
        pos
    }
}

/// Compute positions of all bodies at a given epoch.
/// Returns a flat Vec of [x, y, z, x, y, z, ...] in AU.
pub fn all_positions_at_epoch(bodies: &[CelestialBody], jd: f64) -> Vec<f64> {
    let mut result = Vec::with_capacity(bodies.len() * 3);
    for i in 0..bodies.len() {
        let pos = position_at_epoch(bodies, i, jd);
        result.push(pos.x);
        result.push(pos.y);
        result.push(pos.z);
    }
    result
}

/// Sample points around a body's orbit for drawing orbit ellipses.
/// For planets, samples one full heliocentric orbit.
/// For moons, samples one full orbit around the parent and adds parent position.
/// Returns flat [x, y, z, x, y, z, ...] in AU.
pub fn orbit_points(
    bodies: &[CelestialBody],
    body_id: usize,
    jd: f64,
    samples: usize,
) -> Vec<f64> {
    let body = &bodies[body_id];
    let t = (jd - J2000_JD) / DAYS_PER_CENTURY;

    let a = body.elements.a + body.rates.a * t;
    let e = body.elements.e + body.rates.e * t;
    let i = (body.elements.i + body.rates.i * t) * DEG_TO_RAD;
    let w_bar = (body.elements.w_bar + body.rates.w_bar * t) * DEG_TO_RAD;
    let omega = (body.elements.omega + body.rates.omega * t) * DEG_TO_RAD;
    let w = w_bar - omega;

    let cos_w = w.cos();
    let sin_w = w.sin();
    let cos_o = omega.cos();
    let sin_o = omega.sin();
    let cos_i = i.cos();
    let sin_i = i.sin();

    // Parent offset for moons
    let parent_pos = if let Some(parent_id) = bodies::parent_body(body_id) {
        position_at_epoch(bodies, parent_id, jd)
    } else {
        Vector3::zeros()
    };

    let mut result = Vec::with_capacity(samples * 3);
    for s in 0..samples {
        let nu = TAU * (s as f64) / (samples as f64);
        let r = a * (1.0 - e * e) / (1.0 + e * nu.cos());

        let x_orb = r * nu.cos();
        let y_orb = r * nu.sin();

        let x = (cos_w * cos_o - sin_w * sin_o * cos_i) * x_orb
            + (-sin_w * cos_o - cos_w * sin_o * cos_i) * y_orb
            + parent_pos.x;
        let y = (cos_w * sin_o + sin_w * cos_o * cos_i) * x_orb
            + (-sin_w * sin_o + cos_w * cos_o * cos_i) * y_orb
            + parent_pos.y;
        let z = (sin_w * sin_i) * x_orb + (cos_w * sin_i) * y_orb + parent_pos.z;

        result.push(x);
        result.push(y);
        result.push(z);
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bodies::all_bodies;

    #[test]
    fn test_sun_at_origin() {
        let bodies = all_bodies();
        let pos = position_at_epoch(&bodies, bodies::id::SUN, J2000_JD);
        assert!(pos.norm() < 1e-15);
    }

    #[test]
    fn test_earth_distance_from_sun() {
        let bodies = all_bodies();
        let pos = position_at_epoch(&bodies, bodies::id::EARTH, J2000_JD);
        let distance_au = pos.norm();
        // Earth should be ~1 AU from the Sun
        assert!(
            (distance_au - 1.0).abs() < 0.02,
            "Earth distance: {} AU",
            distance_au
        );
    }

    #[test]
    fn test_jupiter_distance() {
        let bodies = all_bodies();
        let pos = position_at_epoch(&bodies, bodies::id::JUPITER, J2000_JD);
        let distance_au = pos.norm();
        // Jupiter should be ~5.2 AU
        assert!(
            (distance_au - 5.2).abs() < 0.3,
            "Jupiter distance: {} AU",
            distance_au
        );
    }

    #[test]
    fn test_all_positions_length() {
        let bodies = all_bodies();
        let positions = all_positions_at_epoch(&bodies, J2000_JD);
        assert_eq!(positions.len(), bodies.len() * 3);
    }
}