#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 hyperbolic_to_true_anomaly(double hyperbolic_anomaly, double eccentricity) {
    // Hyperbolic E to true anomaly: tanh(ν/2) = √((e-1)/(e+1)) · tanh(H/2)
    double tanh_half_H = tanh(hyperbolic_anomaly / 2.0);
    double tanh_half_nu = sqrt((eccentricity - 1.0) / (eccentricity + 1.0)) * tanh_half_H;
    return 2.0 * atanh(tanh_half_nu);
}

// 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;
}

// TODO: refactor for readability
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);

    // Near-circular: e ≈ 0, true anomaly undefined
    double true_anomaly;
    if (e < 1e-10) {
        true_anomaly = 0.0;
    } else {
        double cos_nu = r_dot_e / (r_mag * e * mu);
        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 * mu;
            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;
        }
    }

    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;
    }

    // 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);

    // 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)

    double omega;
    if (e > 1e-10 && n_mag > 1e-10) {
        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 {
        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 {
        double p = a * (1.0 - e * e);
        double n = sqrt(mu / pow(fabs(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;
    }
}