#include "orbital_mechanics.h"
#include <cmath>
#include <cassert>
#include <cstdio>

void orbital_elements_to_cartesian(OrbitalElements elements, double parent_mass,
                                   Vec3* out_position, Vec3* out_velocity) {
    double a = elements.semi_major_axis;
    double e = elements.eccentricity;
    double nu = elements.true_anomaly;

    double mu = G * parent_mass;

    double sin_nu = sin(nu);
    double cos_nu = cos(nu);

    double p;
    if (fabs(e - 1.0) < PARABOLIC_TOLERANCE) {
        p = elements.semi_latus_rectum;
    } else {
        p = a * (1.0 - e * e); // Semi-latus rectum: p = a(1-e²)
    }

    double r = p / (1.0 + e * cos_nu); // Polar equation of orbit

    double x_orbital = r * cos_nu;
    double y_orbital = r * sin_nu;

    Vec3 position = {x_orbital, y_orbital, 0.0};

    double vx_orbital = -sqrt(mu / p) * sin_nu; // Velocity from vis-viva equation
    double vy_orbital = sqrt(mu / p) * (e + cos_nu);

    Vec3 velocity = {vx_orbital, vy_orbital, 0.0};

    double omega = elements.argument_of_periapsis;
    double i = elements.inclination;
    double Omega = elements.longitude_of_ascending_node;

    Mat3 rotation = mat3_rotation_orbital(omega, i, Omega); // z-x-z Euler angles: R_z(Ω) · R_x(i) · R_z(ω)

    *out_position = mat3_multiply_vec3(rotation, position);
    *out_velocity = mat3_multiply_vec3(rotation, velocity);
}

// Shared solver constants
static const double KEPLER_TOLERANCE = 1.0e-10;
static const int KEPLER_MAX_ITERATIONS = 50;

double get_initial_trial_value(double mean_anomaly, double eccentricity) {
    return mean_anomaly + eccentricity * sin(mean_anomaly)
        + ((pow(eccentricity, 2) / 2.0) * sin(2.0 * mean_anomaly));
}

double solve_kepler_elliptical(double mean_anomaly, double eccentricity) {
    double E = get_initial_trial_value(mean_anomaly, eccentricity);
    double E_prev = E + 2.0 * KEPLER_TOLERANCE;

    int iterations = 0;
    while (fabs(E - E_prev) > KEPLER_TOLERANCE && iterations < KEPLER_MAX_ITERATIONS) {
        E_prev = E;
        double sin_E = sin(E);
        E = E - (E - eccentricity * sin_E - mean_anomaly) / (1.0 - eccentricity * cos(E));
        iterations++;
    }

    return E;
}

double solve_kepler_hyperbolic(double mean_anomaly, double eccentricity) {
    // Initial guess for hyperbolic anomaly
    double H = mean_anomaly;
    if (eccentricity * sinh(mean_anomaly) > mean_anomaly) {
        H = log(2.0 * mean_anomaly / eccentricity);
    }

    double H_prev = H + 2.0 * KEPLER_TOLERANCE;
    int iterations = 0;

    while (fabs(H - H_prev) > KEPLER_TOLERANCE && iterations < KEPLER_MAX_ITERATIONS) {
        H_prev = H;
        double sinh_H = sinh(H);
        double cosh_H = cosh(H);
        H = H - (H - eccentricity * sinh_H - mean_anomaly) / (1.0 - eccentricity * cosh(H));
        iterations++;
    }

    return H;
}

double eccentric_to_true_anomaly(double eccentric_anomaly, double eccentricity) {
    if (fabs(1.0 - eccentricity) < 0.01) {
    // Near-parabolic: use cos/sin formulation to avoid numeric instability
        double E = eccentric_anomaly;
        double e = eccentricity;

        double cos_E = cos(E);
        double sin_E = sin(E);
        double denominator = 1.0 - e * cos_E;

        double cos_nu = (cos_E - e) / denominator;
        double sin_nu = sin_E * sqrt(1.0 - e * e) / denominator;

        cos_nu = fmax(-1.0, fmin(1.0, cos_nu));
        sin_nu = fmax(-1.0, fmin(1.0, sin_nu));

        return atan2(sin_nu, cos_nu);
    }

    double tan_half_E = tan(eccentric_anomaly / 2.0); // Tangent half-angle formula: tan(ν/2) = √((1+e)/(1-e)) · tan(E/2)
    double tan_half_nu = sqrt((1.0 + eccentricity) / (1.0 - eccentricity)) * tan_half_E;
    return 2.0 * atan(tan_half_nu);
}

double true_anomaly_to_eccentric_anomaly(double true_anomaly, double eccentricity) {
    if (fabs(1.0 - eccentricity) < 0.01) {
    // Near-parabolic: use cos/sin formulation to avoid numeric instability
        double nu = true_anomaly;
        double e = eccentricity;

        double cos_nu = cos(nu);
        double sin_nu = sin(nu);
        double denominator = 1.0 + e * cos_nu;

        double cos_E = (cos_nu + e) / denominator;
        double sin_E = sin_nu * sqrt(1.0 - e * e) / denominator;

        cos_E = fmax(-1.0, fmin(1.0, cos_E));
        sin_E = fmax(-1.0, fmin(1.0, sin_E));

        return atan2(sin_E, cos_E);
    }

    double tan_half_nu = tan(true_anomaly / 2.0);
    double tan_half_E = sqrt((1.0 - eccentricity) / (1.0 + eccentricity)) * tan_half_nu;
    return 2.0 * atan(tan_half_E);
}

double hyperbolic_to_true_anomaly(double hyperbolic_anomaly, double eccentricity) {
    // Hyperbolic H to true anomaly: tan(ν/2) = √((e+1)/(e-1)) · tanh(H/2)
    double tanh_half_H = tanh(hyperbolic_anomaly / 2.0);
    double factor = sqrt((eccentricity + 1.0) / (eccentricity - 1.0));  // Inverted
    double tan_half_nu = factor * tanh_half_H;

    // Clamp for numeric stability
    if (tan_half_nu >= 1e10) {
        tan_half_nu = 1e10;
    } else if (tan_half_nu <= -1e10) {
        tan_half_nu = -1e10;
    }

    return 2.0 * atan(tan_half_nu);  // Use atan, not atanh
}

int is_near_hyperbolic_asymptote(double true_anomaly, double eccentricity) {
    // Check if true anomaly is close to asymptote
    // For hyperbolic orbit, asymptotes are at ν = ± acos(-1/e)
    double asymptote = acos(-1.0 / eccentricity);
    double distance_from_asymptote = fabs(fabs(true_anomaly) - asymptote);
    return distance_from_asymptote < 0.01;
}

double true_anomaly_to_hyperbolic(double true_anomaly, double eccentricity) {
    // True anomaly to hyperbolic anomaly: tanh(H/2) = √((e-1)/(e+1)) · tan(ν/2)
    // Solving for H: H = 2 · atanh(√((e-1)/(e+1)) · tan(ν/2))

    if (is_near_hyperbolic_asymptote(true_anomaly, eccentricity)) {
        return -1e10;
    }

    double tan_half_nu = tan(true_anomaly / 2.0);
    double factor = sqrt((eccentricity - 1.0) / (eccentricity + 1.0));
    double tanh_half_H = tan_half_nu * factor;  // Multiply, not divide

    if (tanh_half_H >= 1.0) {
        tanh_half_H = 0.999999999999999;
    } else if (tanh_half_H <= -1.0) {
        tanh_half_H = -0.999999999999999;
    }

    return 2.0 * atanh(tanh_half_H);
}

// Conversion chain: M → E/H → ν
double mean_anomaly_to_true_anomaly(double mean_anomaly, double eccentricity) {
    if (eccentricity < 1.0) {
        double E = solve_kepler_elliptical(mean_anomaly, eccentricity);
        return eccentric_to_true_anomaly(E, eccentricity);
    } else {
        double H = solve_kepler_hyperbolic(mean_anomaly, eccentricity);
        return hyperbolic_to_true_anomaly(H, eccentricity);
    }
}

double solve_barker_equation(double mean_anomaly) {
    if (fabs(mean_anomaly) < 1e-15) {
        return 0.0;
    }

    double c = 1.5 * mean_anomaly;
    double discriminant = c * c + 1.0;
    double sqrt_discriminant = sqrt(discriminant);
    double D = cbrt(c + sqrt_discriminant) + cbrt(c - sqrt_discriminant);
    double nu = 2.0 * atan(D);

    return nu;
}

// FIXME: refactor for readability and sanity
OrbitalElements cartesian_to_orbital_elements(Vec3 position, Vec3 velocity, double parent_mass) {
    double mu = G * parent_mass;

    Vec3 h_vec = vec3_cross(position, velocity);
    Vec3 r_vec = position;
    Vec3 v_vec = velocity;

    double r = vec3_magnitude(r_vec);
    double v = vec3_magnitude(v_vec);
    double v_squared = v * v;

    double specific_energy = -mu / r + v_squared / 2.0; // Specific orbital energy: ε = v²/2 - μ/r
    double h = vec3_magnitude(h_vec);
    // Eccentricity vector: e_vec = (v² - μ/r)r_vec - (r_vec·v_vec)v_vec

    double e_vec_x = ((v_squared - mu / r) * r_vec.x - (vec3_dot(r_vec, v_vec)) * v_vec.x) / mu;
    double e_vec_y = ((v_squared - mu / r) * r_vec.y - (vec3_dot(r_vec, v_vec)) * v_vec.y) / mu;
    double e_vec_z = ((v_squared - mu / r) * r_vec.z - (vec3_dot(r_vec, v_vec)) * v_vec.z) / mu;

    Vec3 e_vec = {e_vec_x, e_vec_y, e_vec_z};
    double e = vec3_magnitude(e_vec);

    double a;
    // Near-parabolic: energy too close to zero, use large value
    if (fabs(specific_energy) < 1e-10) {
        a = 1e10;
    } else {
        a = -mu / (2.0 * specific_energy); // Semi-major axis: a = -μ/(2ε)
    }

    double r_mag = vec3_magnitude(r_vec);
    double r_dot_e = vec3_dot(r_vec, e_vec);

    // Ascending node vector: n = k × h_vec (k is unit Z vector)
    Vec3 n_vec = {0.0, 0.0, 1.0};
    Vec3 n = vec3_cross(n_vec, h_vec);
    double n_mag = vec3_magnitude(n);

    // True anomaly: angle from periapsis to position
    double true_anomaly;
    if (e < 1e-10) {
        // Circular orbit: no periapsis direction. Use argument of latitude
        // (position angle in orbital plane) as true anomaly, with omega=0.
        // For nearly-coplanar orbits, the ascending node is numerically
        // unstable. Use the inclination to decide which reference to use.
        double true_anomaly_from_position;
        double sin_i = (h > 1e-10) ? n_mag / h : 1.0;
        if (sin_i > 1e-6 && n_mag > 1e-10) {
            // Well-defined ascending node: compute argument of latitude
            double x_AN = n.x / n_mag;
            double y_AN = n.y / n_mag;
            // y_AN in orbital plane = (h × n) / |h × n|
            double h_cross_n_x = h_vec.y * 0.0 - h_vec.z * n.y;
            double h_cross_n_y = h_vec.z * n.x - h_vec.x * 0.0;
            double h_cross_n_z = h_vec.x * n.y - h_vec.y * n.x;
            double hcn_mag = sqrt(h_cross_n_x*h_cross_n_x + h_cross_n_y*h_cross_n_y + h_cross_n_z*h_cross_n_z);
            if (hcn_mag > 1e-10) {
                h_cross_n_x /= hcn_mag;
                h_cross_n_y /= hcn_mag;
                h_cross_n_z /= hcn_mag;
            }
            double r_xAN = r_vec.x * x_AN + r_vec.y * y_AN;
            double r_yAN = r_vec.x * h_cross_n_x + r_vec.y * h_cross_n_y + r_vec.z * h_cross_n_z;
            true_anomaly_from_position = atan2(r_yAN, r_xAN);
        } else {
            // Nearly coplanar: ascending node is numerically unstable.
            // Use X-axis as reference. For coplanar orbits this gives
            // the argument of latitude = atan2(y, x).
            true_anomaly_from_position = atan2(r_vec.y, r_vec.x);
        }
        true_anomaly = normalize_angle(true_anomaly_from_position);
    } else {
        double cos_nu = r_dot_e / (r_mag * e);
        cos_nu = fmax(-1.0, fmin(1.0, cos_nu));
        double sin_nu;

        if (fabs(cos_nu) > 1.0 - 1e-10) {
            Vec3 h_cross_e = vec3_cross(h_vec, e_vec);
            double denom = r_mag * e * h;
            if (denom > 1e-10) {
                sin_nu = vec3_dot(r_vec, h_cross_e) / denom;
            } else {
                sin_nu = 0.0;
            }
        } else {
            Vec3 r_cross_h = vec3_cross(r_vec, h_vec);
            double denom = r_mag * e * h;
            sin_nu = (denom > 1e-10) ? vec3_dot(r_cross_h, e_vec) / denom : 0.0;
        }

        true_anomaly = atan2(sin_nu, cos_nu);
        if (true_anomaly == -M_PI) {
            true_anomaly = M_PI;
        }
        true_anomaly = normalize_angle(true_anomaly);
    }

    double i;
    double h_z = h_vec.z;
    if (h > 1e-10) {
        i = acos(h_z / h); // Inclination: i = acos(h_z / h)
    } else {
        i = 0.0;
    }

    // Longitude of ascending node from n vector
    double Omega;
    if (n_mag > 1e-10) {
        Omega = acos(n.x / n_mag);
        if (n.y < 0.0) {
            Omega = 2.0 * M_PI - Omega;
        }
    } else {
        Omega = 0.0;
    }
    // Argument of periapsis: ω = atan2(n×e·h, e·n)
    // For coplanar orbits, use longitude of periapsis (angle of eccentricity vector)

    double omega;
    double inclination_threshold = 0.01;

    if (e > 1e-10 && n_mag > 1e-10 && i > inclination_threshold) {
        double cos_omega = vec3_dot(e_vec, n) / (e * n_mag);
        Vec3 n_cross_e = vec3_cross(n, e_vec);
        double sin_omega = vec3_dot(n_cross_e, h_vec) / (e * n_mag * h);
        omega = atan2(sin_omega, cos_omega);

        if (omega < 0.0) {
            omega += 2.0 * M_PI;
        }
    } else if (e > 1e-10) {
        // Coplanar or near-circular: use longitude of periapsis
        omega = atan2(e_vec.y, e_vec.x);
        if (omega < 0.0) {
            omega += 2.0 * M_PI;
        }
    } else {
        omega = 0.0;
    }

    OrbitalElements elements;
    if (fabs(e - 1.0) < 1e-3) {
        elements.semi_latus_rectum = (h * h) / mu;
    } else {
        elements.semi_major_axis = a;
    }
    elements.eccentricity = e;
    elements.true_anomaly = true_anomaly;
    elements.inclination = i;
    elements.longitude_of_ascending_node = Omega;
    elements.argument_of_periapsis = omega;

    return elements;
}

OrbitalElements propagate_orbital_elements(const OrbitalElements& elements, double dt, double parent_mass) {
    double a = elements.semi_major_axis;
    double e = elements.eccentricity;
    double nu = elements.true_anomaly;
    double mu = G * parent_mass;

    if (fabs(e - 1.0) < PARABOLIC_TOLERANCE) {
        double p = elements.semi_latus_rectum;
        double D = tan(nu / 2.0);
        double M = D + (D * D * D) / 3.0;
        double n = sqrt(mu / pow(p, 3.0));
        M = M + n * dt;
        double nu_new = solve_barker_equation(M);

        OrbitalElements result = elements;
        result.true_anomaly = nu_new;
        return result;
    } else if (e < 1.0) {
        double n = sqrt(mu / pow(a, 3.0));

        double E = 2.0 * atan(sqrt((1.0 - e) / (1.0 + e)) * tan(nu / 2.0));

        double M = E - e * sin(E);

        M = M + n * dt;

        double E_new = get_initial_trial_value(M, e);

        const double CONVERGENCE_TOLERANCE = 1.0e-10;
        const int MAX_ITERATIONS = 50;

        int iterations = 0;
        double E_prev = E_new + 2.0 * CONVERGENCE_TOLERANCE;
        while (fabs(E_new - E_prev) > CONVERGENCE_TOLERANCE && iterations < MAX_ITERATIONS) {
            E_prev = E_new;
            double sin_E = sin(E_new);
            E_new = E_new - (E_new - e * sin_E - M) / (1.0 - e * cos(E_new));
            iterations++;
        }

        OrbitalElements result = elements;
        result.true_anomaly = 2.0 * atan(sqrt((1.0 + e) / (1.0 - e)) * tan(E_new / 2.0));
        return result;
    } else {  // e >= 1.0 (hyperbolic)
        double n = sqrt(mu / pow(-a, 3.0));

        // Convert true anomaly to hyperbolic anomaly
        double H = true_anomaly_to_hyperbolic(nu, e);

        // Compute mean anomaly from hyperbolic anomaly
        double M = e * sinh(H) - H;
        double M_new = M + n * dt;

        // Newton-Raphson iteration for convergence
        const double HYPERBOLIC_TOLERANCE = 1.0e-10;
        const int MAX_HYPERBOLIC_ITERATIONS = 50;
        int iterations = 0;
        double H_new = H;
        double H_prev = H_new + 2.0 * HYPERBOLIC_TOLERANCE;

        while (fabs(H_new - H_prev) > HYPERBOLIC_TOLERANCE && iterations < MAX_HYPERBOLIC_ITERATIONS) {
            H_prev = H_new;
            double sinh_H = sinh(H_new);
            double cosh_H = cosh(H_new);
            H_new = H_new - (e * sinh_H - H_new - M_new) / (e * cosh_H - 1.0);
            iterations++;
        }

        OrbitalElements result = elements;
        result.true_anomaly = hyperbolic_to_true_anomaly(H_new, e);

        return result;
    }
}

// Normalize angle to [0, 2π) range
double normalize_angle(double angle) {
    while (angle < 0.0) angle += 2.0 * M_PI;
    while (angle >= 2.0 * M_PI) angle -= 2.0 * M_PI;
    return angle;
}

// Calculate shortest angular distance between two angles (always positive, range [0, π])
double angular_distance(double a, double b) {
    double diff = fabs(normalize_angle(a) - normalize_angle(b));
    return (diff > M_PI) ? (2.0 * M_PI - diff) : diff;
}

// Calculate eccentricity vector from state vectors
Vec3 calculate_eccentricity_vector(Vec3 r, Vec3 v, Vec3 h, double mu) {
    Vec3 v_cross_h = vec3_cross(v, h);
    Vec3 v_cross_h_over_mu = vec3_scale(v_cross_h, 1.0 / mu);
    double r_mag = vec3_magnitude(r);
    Vec3 r_over_mag = vec3_scale(r, 1.0 / r_mag);
    return vec3_sub(v_cross_h_over_mu, r_over_mag);
}

// Calculate true anomaly from position and velocity vectors
double calculate_true_anomaly(Vec3 r, Vec3 v, Vec3 e_vec, double e_mag, double r_mag) {
    // For near-circular orbits, eccentricity vector is near-zero
    // Compute true anomaly as the angle in the orbital plane
    if (e_mag < 1e-10) {
        Vec3 h = vec3_cross(r, v);
        double h_mag = vec3_magnitude(h);
        if (h_mag < 1e-10) return 0.0;

        // Create a coordinate system in the orbital plane
        Vec3 z_hat = vec3_scale(h, 1.0 / h_mag);
        // Choose x-axis as cross product of Z (world up) and orbit normal
        // This gives a consistent reference direction in the orbital plane
        Vec3 world_z = {0.0, 0.0, 1.0};
        Vec3 x_hat = vec3_cross(world_z, z_hat);
        double x_hat_mag = vec3_magnitude(x_hat);
        if (x_hat_mag < 1e-10) {
            // Orbit is equatorial, use world X as reference
            x_hat = (Vec3){1.0, 0.0, 0.0};
        } else {
            x_hat = vec3_scale(x_hat, 1.0 / x_hat_mag);
        }
        Vec3 y_hat = vec3_cross(z_hat, x_hat);

        // Project position onto this orbital plane coordinate system
        double x_proj = vec3_dot(r, x_hat);
        double y_proj = vec3_dot(r, y_hat);

        // True anomaly is the angle in the orbital plane
        double nu = atan2(y_proj, x_proj);
        if (nu < 0) nu += 2.0 * M_PI;
        return nu;
    }

    // Standard calculation using eccentricity vector
    double cos_nu = vec3_dot(e_vec, r) / (e_mag * r_mag);
    cos_nu = fmax(-1.0, fmin(1.0, cos_nu));
    double nu = acos(cos_nu);

    // Determine correct quadrant using cross product
    Vec3 r_cross_v = vec3_cross(r, v);
    double r_cross_v_dot_e = vec3_dot(r_cross_v, e_vec);
    if (r_cross_v_dot_e < 0) {
        nu = 2.0 * M_PI - nu;
    }

    return nu;
}