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

    double r = p / (1.0 + e * cos_nu);

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

    *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) {
        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);
    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);
}

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

// 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;
    double h = vec3_magnitude(h_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;
    if (fabs(specific_energy) < 1e-10) {
        a = 1e10;
    } else {
        a = -mu / (2.0 * specific_energy);
    }

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

    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));
        true_anomaly = acos(cos_nu);
        if (vec3_dot(r_vec, v_vec) < 0.0) {
            true_anomaly = 2.0 * M_PI - true_anomaly;
        }
        // Normalize to (-π, π] range
        if (true_anomaly > M_PI) {
            true_anomaly -= 2.0 * M_PI;
        }
    }

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

    Vec3 n_vec = {0.0, 0.0, 1.0};
    Vec3 n = vec3_cross(n_vec, h_vec);
    double n_mag = vec3_magnitude(n);

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

    double e_dot_n = vec3_dot(e_vec, n) / (mu * n_mag);

    double omega;
    if (e > 1e-10 && n_mag > 1e-10) {
        omega = acos(e_dot_n / e);
        if (e_vec.z < 0.0) {
            omega = 2.0 * M_PI - omega;
        }
    } 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;
    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;
}