#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/orbital_objects.h"
#include "../src/maneuver.h"
#include "../src/config_loader.h"
#include <cmath>
#include <vector>

const double TIME_STEP = 60.0;
const double POSITION_TOLERANCE = 1e-3;
const double VELOCITY_TOLERANCE = 1e-3;
const double ENERGY_TOLERANCE_RELATIVE = 1e-9;
const double ENERGY_TOLERANCE_ABSOLUTE = 1e-6;

const double RK4_CIRCULAR_TOLERANCE = 2e-7;
const double RK4_ELLIPTICAL_TOLERANCE = 6e-5;
const double RK4_HIGH_ECCENTRICITY_TOLERANCE = 5e-3;

double calculate_spacecraft_kinetic_energy(Vec3 velocity, double mass) {
    double v_squared = velocity.x * velocity.x +
                       velocity.y * velocity.y +
                       velocity.z * velocity.z;
    return 0.5 * mass * v_squared;
}

double calculate_spacecraft_potential_energy(Vec3 position, double craft_mass, double parent_mass) {
    double r = sqrt(position.x * position.x +
                   position.y * position.y +
                   position.z * position.z);
    if (r < 1.0) r = 1.0;
    return -G * craft_mass * parent_mass / r;
}

double calculate_spacecraft_total_energy(Vec3 position, Vec3 velocity,
                                         double craft_mass, double parent_mass) {
    double ke = calculate_spacecraft_kinetic_energy(velocity, craft_mass);
    double pe = calculate_spacecraft_potential_energy(position, craft_mass, parent_mass);
    return ke + pe;
}

double get_orbital_period(double semi_major_axis, double parent_mass) {
    return 2.0 * M_PI * sqrt(pow(semi_major_axis, 3) / (G * parent_mass));
}

TEST_CASE("Config loading for hybrid energy conservation", "[hybrid][energy][config]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    REQUIRE(sim->body_count == 2);
    REQUIRE(std::string(sim->bodies[0].name) == "Sun");
    REQUIRE(std::string(sim->bodies[1].name) == "Earth");

    REQUIRE(sim->craft_count == 6);

    REQUIRE(std::string(sim->spacecraft[0].name) == "Circular_Orbit");
    REQUIRE(sim->spacecraft[0].parent_index == 1);

    REQUIRE(std::string(sim->spacecraft[1].name) == "Elliptical_Orbit");
    REQUIRE(sim->spacecraft[1].parent_index == 1);

    REQUIRE(std::string(sim->spacecraft[2].name) == "High_Eccentricity_Orbit");
    REQUIRE(sim->spacecraft[2].parent_index == 1);

    REQUIRE(std::string(sim->spacecraft[3].name) == "Inclined_Orbit");
    REQUIRE(sim->spacecraft[3].parent_index == 1);

    REQUIRE(std::string(sim->spacecraft[4].name) == "Fast_Orbit");
    REQUIRE(sim->spacecraft[4].parent_index == 1);

    REQUIRE(std::string(sim->spacecraft[5].name) == "Slow_Orbit");
    REQUIRE(sim->spacecraft[5].parent_index == 1);

    destroy_simulation(sim);
}

TEST_CASE("Energy comparison for circular orbit", "[hybrid][energy][circular]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

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

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int steps = 100;
    double dt = orbital_period / steps;

    for (int i = 0; i < steps; i++) {
        analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
        rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
    }

    orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

    double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                 craft->mass, earth->mass);
    double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                               craft->mass, earth->mass);

    INFO("Analytical energy: " << analytical_energy << " J");
    INFO("Numerical energy: " << numerical_energy << " J");
    INFO("Energy difference (analytical): " <<
         fabs(analytical_energy - initial_energy) << " J");
    INFO("Energy difference (numerical): " <<
         fabs(numerical_energy - initial_energy) << " J");

    double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
    double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);
    INFO("Analytical drift: " << analytical_drift);
    INFO("Numerical drift: " << numerical_drift);

    REQUIRE(analytical_drift < 1e-12);
    REQUIRE(numerical_drift < RK4_CIRCULAR_TOLERANCE);

    Vec3 pos_diff = vec3_sub(analytical_pos, numerical_pos);
    double pos_error = vec3_magnitude(pos_diff);
    Vec3 vel_diff = vec3_sub(analytical_vel, numerical_vel);
    double vel_error = vec3_magnitude(vel_diff);

    INFO("Position error: " << pos_error << " m");
    INFO("Velocity error: " << vel_error << " m/s");

    double pos_tolerance_circular = POSITION_TOLERANCE * 1e5;
    double vel_tolerance_circular = VELOCITY_TOLERANCE * 1000;

    REQUIRE(pos_error < pos_tolerance_circular);
    REQUIRE(vel_error < vel_tolerance_circular);

    destroy_simulation(sim);
}

TEST_CASE("Energy comparison for elliptical orbit", "[hybrid][energy][elliptical]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

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

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int steps = 100;
    double dt = orbital_period / steps;

    for (int i = 0; i < steps; i++) {
        analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
        rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
    }

    orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

    double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                 craft->mass, earth->mass);
    double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                               craft->mass, earth->mass);

    INFO("Analytical energy: " << analytical_energy << " J");
    INFO("Numerical energy: " << numerical_energy << " J");
    INFO("Energy difference (analytical): " <<
         fabs(analytical_energy - initial_energy) << " J");
    INFO("Energy difference (numerical): " <<
         fabs(numerical_energy - initial_energy) << " J");

    double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
    double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);
    INFO("Analytical drift: " << analytical_drift);
    INFO("Numerical drift: " << numerical_drift);

    REQUIRE(analytical_drift < 1e-12);
    REQUIRE(numerical_drift < RK4_ELLIPTICAL_TOLERANCE);

    Vec3 pos_diff = vec3_sub(analytical_pos, numerical_pos);
    double pos_error = vec3_magnitude(pos_diff);
    Vec3 vel_diff = vec3_sub(analytical_vel, numerical_vel);
    double vel_error = vec3_magnitude(vel_diff);

    INFO("Position error: " << pos_error << " m");
    INFO("Velocity error: " << vel_error << " m/s");

    double pos_tolerance_elliptical = POSITION_TOLERANCE * 1e7;
    double vel_tolerance_elliptical = VELOCITY_TOLERANCE * 1e4;

    REQUIRE(pos_error < pos_tolerance_elliptical);
    REQUIRE(vel_error < vel_tolerance_elliptical);

    destroy_simulation(sim);
}

TEST_CASE("Energy comparison for high eccentricity orbit", "[hybrid][energy][high_eccentricity]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    Spacecraft* craft = &sim->spacecraft[2];
    CelestialBody* earth = &sim->bodies[1];

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int steps = 200;
    double dt = orbital_period / steps;

    for (int i = 0; i < steps; i++) {
        analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
        rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
    }

    orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

    double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                 craft->mass, earth->mass);
    double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                               craft->mass, earth->mass);

    INFO("Analytical energy: " << analytical_energy << " J");
    INFO("Numerical energy: " << numerical_energy << " J");
    INFO("Energy difference (analytical): " <<
         fabs(analytical_energy - initial_energy) << " J");
    INFO("Energy difference (numerical): " <<
         fabs(numerical_energy - initial_energy) << " J");

    double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
    double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);
    INFO("Analytical drift: " << analytical_drift);
    INFO("Numerical drift: " << numerical_drift);

    REQUIRE(analytical_drift < 1e-12);
    REQUIRE(numerical_drift < RK4_HIGH_ECCENTRICITY_TOLERANCE);

    Vec3 pos_diff = vec3_sub(analytical_pos, numerical_pos);
    double pos_error = vec3_magnitude(pos_diff);
    Vec3 vel_diff = vec3_sub(analytical_vel, numerical_vel);
    double vel_error = vec3_magnitude(vel_diff);

    INFO("Position error: " << pos_error << " m");
    INFO("Velocity error: " << vel_error << " m/s");

    double pos_tolerance_high_ecc = POSITION_TOLERANCE * 1e10;
    double vel_tolerance_high_ecc = VELOCITY_TOLERANCE * 1e7;

    REQUIRE(pos_error < pos_tolerance_high_ecc);
    REQUIRE(vel_error < vel_tolerance_high_ecc);

    destroy_simulation(sim);
}

TEST_CASE("Energy comparison for inclined orbit", "[hybrid][energy][inclined]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    Spacecraft* craft = &sim->spacecraft[3];
    CelestialBody* earth = &sim->bodies[1];

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int steps = 100;
    double dt = orbital_period / steps;

    for (int i = 0; i < steps; i++) {
        analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
        rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
    }

    orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

    double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                 craft->mass, earth->mass);
    double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                               craft->mass, earth->mass);

    INFO("Analytical energy: " << analytical_energy << " J");
    INFO("Numerical energy: " << numerical_energy << " J");
    INFO("Energy difference (analytical): " <<
         fabs(analytical_energy - initial_energy) << " J");
    INFO("Energy difference (numerical): " <<
         fabs(numerical_energy - initial_energy) << " J");

    double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
    double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);
    INFO("Analytical drift: " << analytical_drift);
    INFO("Numerical drift: " << numerical_drift);

    REQUIRE(analytical_drift < 1e-12);
    REQUIRE(numerical_drift < RK4_CIRCULAR_TOLERANCE);

    Vec3 pos_diff = vec3_sub(analytical_pos, numerical_pos);
    double pos_error = vec3_magnitude(pos_diff);
    Vec3 vel_diff = vec3_sub(analytical_vel, numerical_vel);
    double vel_error = vec3_magnitude(vel_diff);

    INFO("Position error: " << pos_error << " m");
    INFO("Velocity error: " << vel_error << " m/s");

    double pos_tolerance_inclined = POSITION_TOLERANCE * 1e5;
    double vel_tolerance_inclined = VELOCITY_TOLERANCE * 1000;

    REQUIRE(pos_error < pos_tolerance_inclined);
    REQUIRE(vel_error < vel_tolerance_inclined);

    INFO("Inclination test: analytical.z = " << analytical_pos.z <<
         ", numerical.z = " << numerical_pos.z);

    destroy_simulation(sim);
}

TEST_CASE("Energy comparison for fast orbit", "[hybrid][energy][fast]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    Spacecraft* craft = &sim->spacecraft[4];
    CelestialBody* earth = &sim->bodies[1];

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int orbits = 10;
    int steps_per_orbit = 100;
    double dt = orbital_period / steps_per_orbit;

    double analytical_drift_max = 0.0;
    double numerical_drift_max = 0.0;

    for (int orbit = 0; orbit < orbits; orbit++) {
        for (int step = 0; step < steps_per_orbit; step++) {
            analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
            rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
        }

        orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

        double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                     craft->mass, earth->mass);
        double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                                   craft->mass, earth->mass);

        double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
        double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);

        analytical_drift_max = std::max(analytical_drift_max, analytical_drift);
        numerical_drift_max = std::max(numerical_drift_max, numerical_drift);

        INFO("Orbit " << orbit + 1 << ": analytical drift = " << analytical_drift
             << ", numerical drift = " << numerical_drift);
    }

    INFO("Maximum analytical drift: " << analytical_drift_max);
    INFO("Maximum numerical drift: " << numerical_drift_max);

    REQUIRE(analytical_drift_max < 1e-12);
    REQUIRE(numerical_drift_max < RK4_CIRCULAR_TOLERANCE * 10);
    REQUIRE(analytical_drift_max < numerical_drift_max);

    destroy_simulation(sim);
}

TEST_CASE("Energy comparison for slow orbit", "[hybrid][energy][slow]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    Spacecraft* craft = &sim->spacecraft[5];
    CelestialBody* earth = &sim->bodies[1];

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int steps = 100;
    double dt = orbital_period / steps;

    for (int i = 0; i < steps; i++) {
        analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
        rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
    }

    orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

    double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                 craft->mass, earth->mass);
    double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                               craft->mass, earth->mass);

    INFO("Analytical energy: " << analytical_energy << " J");
    INFO("Numerical energy: " << numerical_energy << " J");
    INFO("Energy difference (analytical): " <<
         fabs(analytical_energy - initial_energy) << " J");
    INFO("Energy difference (numerical): " <<
         fabs(numerical_energy - initial_energy) << " J");

    double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
    double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);
    INFO("Analytical drift: " << analytical_drift);
    INFO("Numerical drift: " << numerical_drift);

    REQUIRE(analytical_drift < 1e-12);
    REQUIRE(numerical_drift < RK4_CIRCULAR_TOLERANCE);

    Vec3 pos_diff = vec3_sub(analytical_pos, numerical_pos);
    double pos_error = vec3_magnitude(pos_diff);
    Vec3 vel_diff = vec3_sub(analytical_vel, numerical_vel);
    double vel_error = vec3_magnitude(vel_diff);

    INFO("Position error: " << pos_error << " m");
    INFO("Velocity error: " << vel_error << " m/s");

    double pos_tolerance_slow = POSITION_TOLERANCE * 1e5;
    double vel_tolerance_slow = VELOCITY_TOLERANCE * 100;

    REQUIRE(pos_error < pos_tolerance_slow);
    REQUIRE(vel_error < vel_tolerance_slow);

    destroy_simulation(sim);
}

TEST_CASE("Pre/post burn energy validation", "[hybrid][energy][burn]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    SECTION("Circular orbit burn") {
        Spacecraft* craft = &sim->spacecraft[0];
        CelestialBody* earth = &sim->bodies[1];

        Vec3 pos, vel;
        orbital_elements_to_cartesian(craft->orbit, earth->mass, &pos, &vel);
        craft->local_position = pos;
        craft->local_velocity = vel;

        double initial_energy = calculate_spacecraft_total_energy(pos, vel, craft->mass, earth->mass);
        INFO("Initial energy: " << initial_energy << " J");

        double delta_v = 100.0;
        Vec3 v_initial = craft->local_velocity;

        apply_impulsive_burn(craft, BURN_PROGRADE, delta_v);

        Vec3 v_final = craft->local_velocity;
        Vec3 dv = vec3_sub(v_final, v_initial);

        double expected_energy_change = vec3_dot(v_initial, dv) * craft->mass +
                                          0.5 * craft->mass * vec3_dot(dv, dv);

        double final_energy = calculate_spacecraft_total_energy(craft->local_position,
                                                                craft->local_velocity,
                                                                craft->mass, earth->mass);
        double actual_energy_change = final_energy - initial_energy;

        INFO("Final energy: " << final_energy << " J");
        INFO("Expected ΔE: " << expected_energy_change << " J");
        INFO("Actual ΔE: " << actual_energy_change << " J");

        double energy_error = fabs(actual_energy_change - expected_energy_change) / fabs(expected_energy_change);
        REQUIRE(energy_error < 1e-9);

        Vec3 analytical_pos, analytical_vel;
        Vec3 numerical_pos, numerical_vel;
        analytical_pos = craft->local_position;
        analytical_vel = craft->local_velocity;
        numerical_pos = craft->local_position;
        numerical_vel = craft->local_velocity;

        OrbitalElements analytical_elements = cartesian_to_orbital_elements(analytical_pos, analytical_vel, earth->mass);

        double orbital_period = get_orbital_period(analytical_elements.semi_major_axis, earth->mass);
        int steps = 100;
        double dt = orbital_period / steps;

        for (int i = 0; i < steps; i++) {
            analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
            rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
        }

        orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

        double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                     craft->mass, earth->mass);
        double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                                   craft->mass, earth->mass);

        double analytical_drift = fabs(analytical_energy - final_energy) / fabs(final_energy);
        double numerical_drift = fabs(numerical_energy - final_energy) / fabs(final_energy);

        INFO("Post-burn analytical drift: " << analytical_drift);
        INFO("Post-burn numerical drift: " << numerical_drift);

        REQUIRE(analytical_drift < 1e-12);
        REQUIRE(numerical_drift < RK4_CIRCULAR_TOLERANCE);
    }

    SECTION("Elliptical orbit burn") {
        Spacecraft* craft = &sim->spacecraft[1];
        CelestialBody* earth = &sim->bodies[1];

        Vec3 pos, vel;
        orbital_elements_to_cartesian(craft->orbit, earth->mass, &pos, &vel);
        craft->local_position = pos;
        craft->local_velocity = vel;

        double initial_energy = calculate_spacecraft_total_energy(pos, vel, craft->mass, earth->mass);
        INFO("Initial energy: " << initial_energy << " J");

        double delta_v = 100.0;
        Vec3 v_initial = craft->local_velocity;

        apply_impulsive_burn(craft, BURN_PROGRADE, delta_v);

        Vec3 v_final = craft->local_velocity;
        Vec3 dv = vec3_sub(v_final, v_initial);

        double expected_energy_change = vec3_dot(v_initial, dv) * craft->mass +
                                          0.5 * craft->mass * vec3_dot(dv, dv);

        double final_energy = calculate_spacecraft_total_energy(craft->local_position,
                                                                craft->local_velocity,
                                                                craft->mass, earth->mass);
        double actual_energy_change = final_energy - initial_energy;

        INFO("Final energy: " << final_energy << " J");
        INFO("Expected ΔE: " << expected_energy_change << " J");
        INFO("Actual ΔE: " << actual_energy_change << " J");

        double energy_error = fabs(actual_energy_change - expected_energy_change) / fabs(expected_energy_change);
        REQUIRE(energy_error < 1e-9);

        Vec3 analytical_pos, analytical_vel;
        Vec3 numerical_pos, numerical_vel;
        analytical_pos = craft->local_position;
        analytical_vel = craft->local_velocity;
        numerical_pos = craft->local_position;
        numerical_vel = craft->local_velocity;

        OrbitalElements analytical_elements = cartesian_to_orbital_elements(analytical_pos, analytical_vel, earth->mass);

        double orbital_period = get_orbital_period(analytical_elements.semi_major_axis, earth->mass);
        int steps = 100;
        double dt = orbital_period / steps;

        for (int i = 0; i < steps; i++) {
            analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
            rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
        }

        orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

        double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                     craft->mass, earth->mass);
        double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                                   craft->mass, earth->mass);

        double analytical_drift = fabs(analytical_energy - final_energy) / fabs(final_energy);
        double numerical_drift = fabs(numerical_energy - final_energy) / fabs(final_energy);

        INFO("Post-burn analytical drift: " << analytical_drift);
        INFO("Post-burn numerical drift: " << numerical_drift);

        REQUIRE(analytical_drift < 1e-12);
        REQUIRE(numerical_drift < RK4_ELLIPTICAL_TOLERANCE);
    }

    destroy_simulation(sim);
}

TEST_CASE("Long-term energy drift comparison", "[hybrid][energy][long_term]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

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

    Vec3 analytical_pos, analytical_vel;
    Vec3 numerical_pos, numerical_vel;

    orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
    numerical_pos = analytical_pos;
    numerical_vel = analytical_vel;

    double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                             craft->mass, earth->mass);
    INFO("Initial energy: " << initial_energy << " J");

    double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);
    INFO("Orbital period: " << orbital_period << " s");

    OrbitalElements analytical_elements = craft->orbit;

    int orbits = 10;
    int steps_per_orbit = 100;
    double dt = orbital_period / steps_per_orbit;

    std::vector<double> analytical_energies;
    std::vector<double> numerical_energies;

    for (int orbit = 0; orbit < orbits; orbit++) {
        for (int step = 0; step < steps_per_orbit; step++) {
            analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
            rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);
        }

        orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

        double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                     craft->mass, earth->mass);
        double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                                   craft->mass, earth->mass);

        analytical_energies.push_back(analytical_energy);
        numerical_energies.push_back(numerical_energy);

        INFO("Orbit " << orbit + 1 << ": analytical = " << analytical_energy
             << ", numerical = " << numerical_energy);
    }

    double analytical_drift_final = fabs(analytical_energies.back() - initial_energy) / fabs(initial_energy);
    double numerical_drift_final = fabs(numerical_energies.back() - initial_energy) / fabs(initial_energy);

    INFO("Final analytical drift: " << analytical_drift_final);
    INFO("Final numerical drift: " << numerical_drift_final);

    REQUIRE(analytical_drift_final < 1e-12);
    REQUIRE(numerical_drift_final < RK4_ELLIPTICAL_TOLERANCE * 10);
    REQUIRE(analytical_drift_final < numerical_drift_final);

    double analytical_drift_max = 0.0;
    double numerical_drift_max = 0.0;

    for (size_t i = 0; i < analytical_energies.size(); i++) {
        double analytical_drift = fabs(analytical_energies[i] - initial_energy) / fabs(initial_energy);
        double numerical_drift = fabs(numerical_energies[i] - initial_energy) / fabs(initial_energy);
        analytical_drift_max = std::max(analytical_drift_max, analytical_drift);
        numerical_drift_max = std::max(numerical_drift_max, numerical_drift);
    }

    INFO("Maximum analytical drift: " << analytical_drift_max);
    INFO("Maximum numerical drift: " << numerical_drift_max);

    REQUIRE(analytical_drift_max < 1e-12);
    REQUIRE(analytical_drift_max < numerical_drift_max);

    destroy_simulation(sim);
}

TEST_CASE("Energy accuracy across orbit types", "[hybrid][energy][accuracy]") {
    SimulationState* sim = create_simulation(10, 10, 0, TIME_STEP);

    REQUIRE(load_system_config(sim, "tests/test_hybrid_energy_conservation.toml"));

    const char* craft_names[] = {
        "Circular_Orbit",
        "Elliptical_Orbit",
        "High_Eccentricity_Orbit",
        "Inclined_Orbit",
        "Fast_Orbit",
        "Slow_Orbit"
    };

    for (int craft_idx = 0; craft_idx < 6; craft_idx++) {
        INFO("Testing spacecraft: " << craft_names[craft_idx]);

        Spacecraft* craft = &sim->spacecraft[craft_idx];
        CelestialBody* earth = &sim->bodies[1];

        Vec3 analytical_pos, analytical_vel;
        Vec3 numerical_pos, numerical_vel;

        orbital_elements_to_cartesian(craft->orbit, earth->mass, &analytical_pos, &analytical_vel);
        numerical_pos = analytical_pos;
        numerical_vel = analytical_vel;

        double initial_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                 craft->mass, earth->mass);

        double orbital_period = get_orbital_period(craft->orbit.semi_major_axis, earth->mass);

        OrbitalElements analytical_elements = craft->orbit;

        int time_points = 100;
        double dt = orbital_period / time_points;

        double max_energy_diff = 0.0;
        double max_analytical_drift = 0.0;
        double max_numerical_drift = 0.0;

        for (int i = 0; i < time_points; i++) {
            analytical_elements = propagate_orbital_elements(analytical_elements, dt, earth->mass);
            rk4_step(&numerical_pos, &numerical_vel, dt, craft->mass, earth->mass);

            orbital_elements_to_cartesian(analytical_elements, earth->mass, &analytical_pos, &analytical_vel);

            double analytical_energy = calculate_spacecraft_total_energy(analytical_pos, analytical_vel,
                                                                         craft->mass, earth->mass);
            double numerical_energy = calculate_spacecraft_total_energy(numerical_pos, numerical_vel,
                                                                       craft->mass, earth->mass);

            double energy_diff = fabs(analytical_energy - numerical_energy);
            double analytical_drift = fabs(analytical_energy - initial_energy) / fabs(initial_energy);
            double numerical_drift = fabs(numerical_energy - initial_energy) / fabs(initial_energy);

            max_energy_diff = std::max(max_energy_diff, energy_diff);
            max_analytical_drift = std::max(max_analytical_drift, analytical_drift);
            max_numerical_drift = std::max(max_numerical_drift, numerical_drift);

            INFO("  Time point " << i + 1 << ": energy diff = " << energy_diff
                 << ", analytical drift = " << analytical_drift
                 << ", numerical drift = " << numerical_drift);
        }

        double relative_energy_diff = max_energy_diff / fabs(initial_energy);
        INFO("Max relative energy difference: " << relative_energy_diff);
        INFO("Max analytical drift: " << max_analytical_drift);
        INFO("Max numerical drift: " << max_numerical_drift);

        double energy_diff_tolerance = (craft_idx == 2) ? RK4_HIGH_ECCENTRICITY_TOLERANCE * 30 :
                                    (craft_idx == 1) ? RK4_ELLIPTICAL_TOLERANCE :
                                    RK4_CIRCULAR_TOLERANCE;

        REQUIRE(relative_energy_diff < energy_diff_tolerance);
        REQUIRE(max_analytical_drift < 1e-12);
        REQUIRE(max_analytical_drift < max_numerical_drift);
    }

    destroy_simulation(sim);
}