#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include "../src/physics.h"
#include "../src/orbital_mechanics.h"
#include "../src/simulation.h"
#include "../src/config_loader.h"
#include "../src/test_utilities.h"
#include <cmath>

using Catch::Matchers::WithinAbs;

// Fixture: tolerance constants
const double LONG_TERM_ANG_TOL = 1e-10;
const double SMALL_DT_VEL_CHANGE_TOL = 1.0;

// Helper functions
static OrbitalElements make_elements(double a, double e, double nu) {
    OrbitalElements el = {};
    el.semi_major_axis = a;
    el.eccentricity = e;
    el.true_anomaly = nu;
    return el;
}

static void get_state(double a, double e, double nu, double parent_mass, Vec3& pos, Vec3& vel) {
    OrbitalElements el = make_elements(a, e, nu);
    orbital_elements_to_cartesian(el, parent_mass, &pos, &vel);
}

static void propagate_and_get_state(double a, double e, double nu, double dt,
                                     double parent_mass, Vec3& pos, Vec3& vel) {
    OrbitalElements el = make_elements(a, e, nu);
    OrbitalElements final_el = propagate_orbital_elements(el, dt, parent_mass);
    orbital_elements_to_cartesian(final_el, parent_mass, &pos, &vel);
}

SCENARIO("Propagation through apsides (velocity extrema)",
         "[analytical][propagation][apsides]") {
    const double TIME_STEP = 60.0;
    SimulationState* sim = create_simulation(10, 2, 0, TIME_STEP);
    REQUIRE(load_system_config(sim, "tests/test_analytical_propagation.toml"));

    Spacecraft* apsides_craft = &sim->spacecraft[0];
    Spacecraft* timestep_craft = &sim->spacecraft[1];
    CelestialBody* earth = &sim->bodies[0];

    const double mu = G * earth->mass;
    const double A1_A = apsides_craft->orbit.semi_major_axis;
    const double A1_E = apsides_craft->orbit.eccentricity;
    const double A2_A = timestep_craft->orbit.semi_major_axis;
    const double A2_E = timestep_craft->orbit.eccentricity;

    auto check_apsides_radius = [&](double a, double e, double nu,
                                     double expected_r, const char* label) {
        Vec3 pos, vel;
        get_state(a, e, nu, earth->mass, pos, vel);
        const double r = vec3_magnitude(pos);
        INFO(label);
        INFO("  Expected r: " << expected_r << " m");
        INFO("  Calculated r: " << r << " m");
        REQUIRE_THAT(r, WithinAbs(expected_r, R_TOL));
    };

    SECTION("apsides spacecraft perigee radius = a*(1-e)") {
        check_apsides_radius(A1_A, A1_E, 0.0, A1_A * (1.0 - A1_E),
                             "Apsides spacecraft perigee");
    }

    SECTION("apsides spacecraft apogee radius = a*(1+e)") {
        check_apsides_radius(A1_A, A1_E, M_PI, A1_A * (1.0 + A1_E),
                             "Apsides spacecraft apogee");
    }

    SECTION("apsides spacecraft perigee velocity > apogee velocity") {
        Vec3 pos_peri, vel_peri, pos_apo, vel_apo;
        get_state(A1_A, A1_E, 0.0, earth->mass, pos_peri, vel_peri);
        get_state(A1_A, A1_E, M_PI, earth->mass, pos_apo, vel_apo);
        const double v_peri = vec3_magnitude(vel_peri);
        const double v_apo = vec3_magnitude(vel_apo);
        INFO("v_peri: " << v_peri << " m/s");
        INFO("v_apo: " << v_apo << " m/s");
        REQUIRE(v_peri > v_apo);
    }

    SECTION("apsides spacecraft perigee velocity > velocity at pi/4") {
        Vec3 pos_45, vel_45, pos_peri, vel_peri;
        get_state(A1_A, A1_E, M_PI / 4.0, earth->mass, pos_45, vel_45);
        get_state(A1_A, A1_E, 0.0, earth->mass, pos_peri, vel_peri);
        const double v_45 = vec3_magnitude(vel_45);
        const double v_peri = vec3_magnitude(vel_peri);
        INFO("v_peri: " << v_peri << " m/s");
        INFO("v_at_pi4: " << v_45 << " m/s");
        REQUIRE(v_peri > v_45);
    }

    SECTION("timestep spacecraft perigee radius = a*(1-e)") {
        check_apsides_radius(A2_A, A2_E, 0.0, A2_A * (1.0 - A2_E),
                             "Timestep spacecraft perigee");
    }

    destroy_simulation(sim);
}

SCENARIO("Full orbit propagation returns to initial state",
         "[analytical][propagation][period]") {
    const double TIME_STEP = 60.0;
    SimulationState* sim = create_simulation(10, 2, 0, TIME_STEP);
    REQUIRE(load_system_config(sim, "tests/test_analytical_propagation.toml"));

    Spacecraft* apsides_craft = &sim->spacecraft[0];
    Spacecraft* timestep_craft = &sim->spacecraft[1];
    CelestialBody* earth = &sim->bodies[0];

    const double mu = G * earth->mass;
    const double A1_A = apsides_craft->orbit.semi_major_axis;
    const double A1_E = apsides_craft->orbit.eccentricity;
    const double A2_A = timestep_craft->orbit.semi_major_axis;
    const double A2_E = timestep_craft->orbit.eccentricity;

    const double A1_PERIOD = 2.0 * M_PI * sqrt(A1_A * A1_A * A1_A / mu);
    const double A2_PERIOD = 2.0 * M_PI * sqrt(A2_A * A2_A * A2_A / mu);

    auto check_period_return = [&](double a, double e, double period,
                                    const char* label) {
        OrbitalElements el = make_elements(a, e, 0.0);
        Vec3 pos_initial, vel_initial;
        orbital_elements_to_cartesian(el, earth->mass, &pos_initial, &vel_initial);

        OrbitalElements final_el = propagate_orbital_elements(el, period, earth->mass);
        Vec3 pos_final, vel_final;
        orbital_elements_to_cartesian(final_el, earth->mass, &pos_final, &vel_final);

        const double pos_error = vec3_distance(pos_initial, pos_final);
        const double vel_error = vec3_distance(vel_initial, vel_final);
        const double r_initial = vec3_magnitude(pos_initial);
        const double v_initial = vec3_magnitude(vel_initial);
        const double rel_pos_error = pos_error / r_initial * 100.0;
        const double rel_vel_error = vel_error / v_initial * 100.0;

        INFO(label);
        INFO("  Relative position error: " << rel_pos_error << "%");
        INFO("  Relative velocity error: " << rel_vel_error << "%");

        REQUIRE_THAT(rel_pos_error, WithinAbs(0.0, REL_TOL * 100.0));
        REQUIRE_THAT(rel_vel_error, WithinAbs(0.0, REL_TOL * 100.0));
    };

    auto check_anomaly_return = [&](double a, double e, double period,
                                     double initial_nu, const char* label) {
        OrbitalElements el = make_elements(a, e, initial_nu);
        OrbitalElements final_el = propagate_orbital_elements(el, period, earth->mass);
        const double final_nu = final_el.true_anomaly;
        const double expected_nu = std::fmod(initial_nu + 2.0 * M_PI, 2.0 * M_PI);
        double anomaly_error = std::abs(final_nu - expected_nu);
        if (anomaly_error > M_PI) {
            anomaly_error = 2.0 * M_PI - anomaly_error;
        }

        INFO(label);
        INFO("  Initial nu: " << initial_nu << " rad");
        INFO("  Final nu: " << final_nu << " rad");
        INFO("  Anomaly error: " << anomaly_error << " rad");

        REQUIRE_THAT(anomaly_error, WithinAbs(0.0, ANG_TOL));
    };

    SECTION("apsides spacecraft position returns after one period") {
        check_period_return(A1_A, A1_E, A1_PERIOD,
                            "Apsides spacecraft position");
    }

    SECTION("apsides spacecraft velocity returns after one period") {
        check_period_return(A1_A, A1_E, A1_PERIOD,
                            "Apsides spacecraft velocity");
    }

    SECTION("apsides spacecraft true anomaly returns after one period") {
        check_anomaly_return(A1_A, A1_E, A1_PERIOD, 0.0,
                             "Apsides spacecraft nu");
    }

    SECTION("timestep spacecraft position returns after one period") {
        check_period_return(A2_A, A2_E, A2_PERIOD,
                            "Timestep spacecraft position");
    }

    SECTION("timestep spacecraft velocity returns after one period") {
        check_period_return(A2_A, A2_E, A2_PERIOD,
                            "Timestep spacecraft velocity");
    }

    SECTION("timestep spacecraft true anomaly returns after one period") {
        check_anomaly_return(A2_A, A2_E, A2_PERIOD, 0.0,
                             "Timestep spacecraft nu");
    }

    destroy_simulation(sim);
}

SCENARIO("Vis-viva equation holds at multiple orbital positions",
         "[analytical][propagation][vis_viva]") {
    const double TIME_STEP = 60.0;
    SimulationState* sim = create_simulation(10, 2, 0, TIME_STEP);
    REQUIRE(load_system_config(sim, "tests/test_analytical_propagation.toml"));

    Spacecraft* apsides_craft = &sim->spacecraft[0];
    CelestialBody* earth = &sim->bodies[0];

    const double mu = G * earth->mass;
    const double A1_A = apsides_craft->orbit.semi_major_axis;
    const double A1_E = apsides_craft->orbit.eccentricity;

    const double true_anomalies[] = {0.0, M_PI / 4.0, M_PI / 2.0,
                                      3.0 * M_PI / 4.0, M_PI};
    const char* labels[] = {
        "nu = 0", "nu = pi/4", "nu = pi/2", "nu = 3pi/4", "nu = pi"
    };

    for (int i = 0; i < 5; i++) {
        SECTION(labels[i]) {
            const double nu = true_anomalies[i];
            Vec3 pos, vel;
            get_state(A1_A, A1_E, nu, earth->mass, pos, vel);
            const double r = vec3_magnitude(pos);
            const double v = vec3_magnitude(vel);
            const double expected_v = std::sqrt(mu * (2.0 / r - 1.0 / A1_A));
            const double v_error = std::abs(v - expected_v);
            const double rel_error = v_error / expected_v * 100.0;

            INFO("nu: " << nu << " rad (" << nu * 180.0 / M_PI << " deg)");
            INFO("r: " << r << " m");
            INFO("v: " << v << " m/s");
            INFO("expected_v: " << expected_v << " m/s");
            INFO("rel_error: " << rel_error << "%");

            REQUIRE_THAT(rel_error, WithinAbs(0.0, REL_TOL * 100.0));
        }
    }

    destroy_simulation(sim);
}

SCENARIO("Accuracy across different timestep sizes",
         "[analytical][timestep][accuracy]") {
    const double TIME_STEP = 60.0;
    SimulationState* sim = create_simulation(10, 2, 0, TIME_STEP);
    REQUIRE(load_system_config(sim, "tests/test_analytical_propagation.toml"));

    Spacecraft* timestep_craft = &sim->spacecraft[1];
    CelestialBody* earth = &sim->bodies[0];

    const double mu = G * earth->mass;
    const double A2_A = timestep_craft->orbit.semi_major_axis;
    const double A2_E = timestep_craft->orbit.eccentricity;
    const double A2_PERIOD = 2.0 * M_PI * sqrt(A2_A * A2_A * A2_A / mu);

    Vec3 pos_init, vel_init;
    get_state(A2_A, A2_E, 0.0, earth->mass, pos_init, vel_init);
    const double r_init = vec3_magnitude(pos_init);
    const double v_init = vec3_magnitude(vel_init);

    SECTION("large timestep (2x period) preserves state") {
        Vec3 pos_final, vel_final;
        propagate_and_get_state(A2_A, A2_E, 0.0, A2_PERIOD * 2.0, earth->mass,
                                pos_final, vel_final);
        const double r_final = vec3_magnitude(pos_final);
        const double v_final = vec3_magnitude(vel_final);
        const double rel_r_error = std::abs(r_final - r_init) / r_init * 100.0;
        const double rel_v_error = std::abs(v_final - v_init) / v_init * 100.0;

        INFO("Relative radius error: " << rel_r_error << "%");
        INFO("Relative velocity error: " << rel_v_error << "%");

        REQUIRE_THAT(rel_r_error, WithinAbs(0.0, REL_TOL * 100.0));
        REQUIRE_THAT(rel_v_error, WithinAbs(0.0, REL_TOL * 100.0));
    }

    SECTION("very small timestep (0.1 s) produces expected displacement") {
        const double dt = 0.1;
        Vec3 pos_final, vel_final;
        propagate_and_get_state(A2_A, A2_E, 0.0, dt, earth->mass,
                                pos_final, vel_final);
        const double pos_change = vec3_distance(pos_init, pos_final);
        const double vel_change = vec3_distance(vel_init, vel_final);
        const double expected_pos_change = v_init * dt;
        const double pos_error = std::abs(pos_change - expected_pos_change);
        const double rel_pos_error = pos_error / expected_pos_change * 100.0;

        INFO("dt: " << dt << " s");
        INFO("pos_change: " << pos_change << " m");
        INFO("expected_pos_change: " << expected_pos_change << " m");
        INFO("pos_error: " << pos_error << " m");
        INFO("rel_pos_error: " << rel_pos_error << "%");
        INFO("vel_change: " << vel_change << " m/s");

        REQUIRE_THAT(rel_pos_error, WithinAbs(0.0, REL_TOL * 100.0));
        REQUIRE_THAT(vel_change, WithinAbs(0.0, SMALL_DT_VEL_CHANGE_TOL));
    }

    SECTION("accuracy at 1x period") {
        const double dt = A2_PERIOD;
        Vec3 pos_final, vel_final;
        propagate_and_get_state(A2_A, A2_E, 0.0, dt, earth->mass,
                                pos_final, vel_final);
        const double pos_error = vec3_distance(pos_init, pos_final);
        const double vel_error = vec3_distance(vel_init, vel_final);

        INFO("dt: " << dt << " s (1x period)");
        INFO("pos_error: " << pos_error << " m");
        INFO("vel_error: " << vel_error << " m/s");

        REQUIRE_THAT(pos_error, WithinAbs(0.0, R_TOL));
        REQUIRE_THAT(vel_error, WithinAbs(0.0, V_TOL));
    }

    SECTION("accuracy at 10x period") {
        const double dt = A2_PERIOD * 10.0;
        Vec3 pos_final, vel_final;
        propagate_and_get_state(A2_A, A2_E, 0.0, dt, earth->mass,
                                pos_final, vel_final);
        const double pos_error = vec3_distance(pos_init, pos_final);
        const double vel_error = vec3_distance(vel_init, vel_final);

        INFO("dt: " << dt << " s (10x period)");
        INFO("pos_error: " << pos_error << " m");
        INFO("vel_error: " << vel_error << " m/s");

        REQUIRE_THAT(pos_error, WithinAbs(0.0, R_TOL));
        REQUIRE_THAT(vel_error, WithinAbs(0.0, V_TOL));
    }

    destroy_simulation(sim);
}

SCENARIO("Long-term propagation stability",
         "[analytical][timestep][long_term]") {
    const double TIME_STEP = 60.0;
    SimulationState* sim = create_simulation(10, 2, 0, TIME_STEP);
    REQUIRE(load_system_config(sim, "tests/test_analytical_propagation.toml"));

    Spacecraft* timestep_craft = &sim->spacecraft[1];
    CelestialBody* earth = &sim->bodies[0];

    const double mu = G * earth->mass;
    const double A2_A = timestep_craft->orbit.semi_major_axis;
    const double A2_E = timestep_craft->orbit.eccentricity;
    const double A2_PERIOD = 2.0 * M_PI * sqrt(A2_A * A2_A * A2_A / mu);

    const double propagation_time = A2_PERIOD * 100.0;
    const double mean_motion = std::sqrt(mu / (A2_A * A2_A * A2_A));
    const double initial_nu = 0.0;

    OrbitalElements el = make_elements(A2_A, A2_E, initial_nu);
    OrbitalElements propagated = propagate_orbital_elements(el, propagation_time, earth->mass);
    const double final_nu = propagated.true_anomaly;

    // Compute expected final anomaly
    const double expected_delta_anomaly = mean_motion * propagation_time;
    double expected_final_nu = initial_nu + expected_delta_anomaly;
    while (expected_final_nu < 0.0) {
        expected_final_nu += 2.0 * M_PI;
    }
    while (expected_final_nu >= 2.0 * M_PI) {
        expected_final_nu -= 2.0 * M_PI;
    }

    // Compute shortest angular distance
    const double raw_error = std::abs(final_nu - expected_final_nu);
    const double anomaly_error = std::fmin(raw_error, 2.0 * M_PI - raw_error);

    INFO("Propagation time: " << propagation_time << " s ("
         << propagation_time / A2_PERIOD << " periods)");
    INFO("Initial nu: " << initial_nu << " rad");
    INFO("Final nu: " << final_nu << " rad");
    INFO("Expected nu: " << expected_final_nu << " rad");
    INFO("Anomaly error: " << anomaly_error << " rad ("
         << anomaly_error * 180.0 / M_PI << " deg)");

    REQUIRE_THAT(anomaly_error, WithinAbs(0.0, LONG_TERM_ANG_TOL));

    destroy_simulation(sim);
}