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

const double POSITION_TOLERANCE = 1.0e3;
const double VELOCITY_TOLERANCE = 1.0e-3;
const double ENERGY_TOLERANCE = 1.0e-6;
const double ORBITAL_ELEMENT_TOLERANCE = 1.0e-9;

double calculate_spacecraft_kinetic_energy(Spacecraft* craft) {
    double v_squared = craft->local_velocity.x * craft->local_velocity.x +
                       craft->local_velocity.y * craft->local_velocity.y +
                       craft->local_velocity.z * craft->local_velocity.z;
    return 0.5 * craft->mass * v_squared;
}

double calculate_spacecraft_potential_energy(Spacecraft* craft, CelestialBody* parent) {
    double distance = vec3_magnitude(craft->local_position);
    if (distance < 1.0) distance = 1.0;
    return -G * craft->mass * parent->mass / distance;
}

double calculate_spacecraft_total_energy(Spacecraft* craft, CelestialBody* parent) {
    return calculate_spacecraft_kinetic_energy(craft) +
           calculate_spacecraft_potential_energy(craft, parent);
}

OrbitalElements simulate_continuous_burn(OrbitalElements initial_orbit, double parent_mass,
                                         double total_dv, double burn_duration,
                                         int num_steps, BurnDirection direction) {
    OrbitalElements current_orbit = initial_orbit;
    double dt_burn_step = burn_duration / num_steps;
    double dv_per_step = total_dv / num_steps;

    for (int i = 0; i < num_steps; i++) {
        Vec3 pos;
        Vec3 vel;
        orbital_elements_to_cartesian(current_orbit, parent_mass, &pos, &vel);

        Vec3 dir = get_burn_direction_vector(direction, pos, vel);
        Vec3 dv_vec = vec3_scale(dir, dv_per_step);
        vel = vec3_add(vel, dv_vec);

        current_orbit = cartesian_to_orbital_elements(pos, vel, parent_mass);
        current_orbit = propagate_orbital_elements(current_orbit, dt_burn_step, parent_mass);
    }

    return current_orbit;
}

TEST_CASE("Config loading for continuous thrust tests", "[hybrid][continuous][config]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

    REQUIRE(sim->body_count == 2);
    REQUIRE(sim->craft_count == 4);

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

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

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

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

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

    destroy_simulation(sim);
}

TEST_CASE("Continuous low-thrust burns (ion engines)", "[hybrid][continuous][low_thrust]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    double initial_semi_major = craft->orbit.semi_major_axis;
    double initial_eccentricity = craft->orbit.eccentricity;

    INFO("Initial semi-major axis: " << initial_semi_major << " m");
    INFO("Initial eccentricity: " << initial_eccentricity);

    double burn_duration = 5000.0;
    double total_dv = 100.0;
    int num_steps = 100;

    OrbitalElements final_orbit = simulate_continuous_burn(craft->orbit, earth->mass,
                                                           total_dv, burn_duration,
                                                           num_steps, BURN_PROGRADE);

    INFO("Final semi-major axis: " << final_orbit.semi_major_axis << " m");
    INFO("Final eccentricity: " << final_orbit.eccentricity);

    REQUIRE(final_orbit.semi_major_axis > initial_semi_major);

    double a_before = initial_semi_major;
    double a_after = final_orbit.semi_major_axis;
    double mu = G * earth->mass;
    double v_circular_initial = sqrt(mu / a_before);
    double v_circular_final = sqrt(mu / a_after);

    double epsilon_initial = -mu / (2.0 * a_before);
    double epsilon_final = -mu / (2.0 * a_after);
    double delta_epsilon = epsilon_final - epsilon_initial;

    INFO("Initial circular velocity: " << v_circular_initial << " m/s");
    INFO("Final circular velocity: " << v_circular_final << " m/s");
    INFO("Initial specific energy: " << epsilon_initial << " J/kg");
    INFO("Final specific energy: " << epsilon_final << " J/kg");
    INFO("Energy change: " << delta_epsilon << " J/kg");
    INFO("Applied delta-v: " << total_dv << " m/s");

    double expected_dv_from_energy = delta_epsilon / v_circular_initial;

    INFO("Expected delta-v from energy: " << expected_dv_from_energy << " m/s");

    double relative_error = fabs(expected_dv_from_energy - total_dv) / total_dv;
    INFO("Relative error: " << relative_error * 100 << "%");

    REQUIRE(relative_error < 0.01);

    REQUIRE(final_orbit.eccentricity < 0.01);

    destroy_simulation(sim);
}

TEST_CASE("Multi-burn sequences", "[hybrid][continuous][multi_burn]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    double initial_semi_major = craft->orbit.semi_major_axis;

    INFO("Initial semi-major axis: " << initial_semi_major << " m");

    double burn_duration_1 = 2000.0;
    double total_dv_1 = 50.0;
    int num_steps_1 = 20;

    OrbitalElements orbit_after_burn1 = simulate_continuous_burn(craft->orbit, earth->mass,
                                                                  total_dv_1, burn_duration_1,
                                                                  num_steps_1, BURN_PROGRADE);

    INFO("Semi-major axis after burn 1: " << orbit_after_burn1.semi_major_axis << " m");

    REQUIRE(orbit_after_burn1.semi_major_axis > initial_semi_major);

    double burn_duration_2 = 3000.0;
    double total_dv_2 = 75.0;
    int num_steps_2 = 30;

    OrbitalElements final_orbit = simulate_continuous_burn(orbit_after_burn1, earth->mass,
                                                            total_dv_2, burn_duration_2,
                                                            num_steps_2, BURN_PROGRADE);

    INFO("Final semi-major axis: " << final_orbit.semi_major_axis << " m");

    REQUIRE(final_orbit.semi_major_axis > orbit_after_burn1.semi_major_axis);

    double a_before = initial_semi_major;
    double a_after = final_orbit.semi_major_axis;
    double mu = G * earth->mass;
    double v_circular_initial = sqrt(mu / a_before);
    double v_circular_final = sqrt(mu / a_after);

    double epsilon_initial = -mu / (2.0 * a_before);
    double epsilon_final = -mu / (2.0 * a_after);
    double delta_epsilon = epsilon_final - epsilon_initial;
    double total_dv_applied = total_dv_1 + total_dv_2;

    INFO("Total applied delta-v: " << total_dv_applied << " m/s");
    INFO("Initial specific energy: " << epsilon_initial << " J/kg");
    INFO("Final specific energy: " << epsilon_final << " J/kg");
    INFO("Energy change: " << delta_epsilon << " J/kg");

    double expected_dv_from_energy = delta_epsilon / v_circular_initial;

    INFO("Expected delta-v from energy: " << expected_dv_from_energy << " m/s");

    double relative_error = fabs(expected_dv_from_energy - total_dv_applied) / total_dv_applied;
    INFO("Relative error: " << relative_error * 100 << "%");

    REQUIRE(relative_error < 0.01);

    destroy_simulation(sim);
}

TEST_CASE("Mode transitions during burns", "[hybrid][continuous][mode_transition]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    double initial_semi_major = craft->orbit.semi_major_axis;
    double initial_eccentricity = craft->orbit.eccentricity;

    INFO("Initial semi-major axis: " << initial_semi_major << " m");
    INFO("Initial eccentricity: " << initial_eccentricity);

    double burn_duration = 4000.0;
    double total_dv = 200.0;
    int num_steps = 80;

    OrbitalElements current_orbit = craft->orbit;
    double dt_burn_step = burn_duration / num_steps;
    double dv_per_step = total_dv / num_steps;

    for (int i = 0; i < num_steps; i++) {
        Vec3 pos;
        Vec3 vel;
        orbital_elements_to_cartesian(current_orbit, earth->mass, &pos, &vel);

        Vec3 dir = get_burn_direction_vector(BURN_PROGRADE, pos, vel);
        Vec3 dv_vec = vec3_scale(dir, dv_per_step);
        vel = vec3_add(vel, dv_vec);

        OrbitalElements orbit_from_cart = cartesian_to_orbital_elements(pos, vel, earth->mass);

        current_orbit = propagate_orbital_elements(orbit_from_cart, dt_burn_step, earth->mass);
    }

    INFO("Final semi-major axis: " << current_orbit.semi_major_axis << " m");
    INFO("Final eccentricity: " << current_orbit.eccentricity);

    REQUIRE(current_orbit.semi_major_axis > initial_semi_major);

    double mu = G * earth->mass;
    double energy_before = -mu / (2.0 * initial_semi_major);
    double energy_after = -mu / (2.0 * current_orbit.semi_major_axis);
    double energy_change = energy_after - energy_before;

    double expected_energy_change = 0.5 * (sqrt(mu / initial_semi_major) + sqrt(mu / current_orbit.semi_major_axis)) * total_dv;

    INFO("Energy change: " << energy_change << " J/kg");
    INFO("Expected energy change: " << expected_energy_change << " J/kg");

    REQUIRE(fabs(energy_change) > 0);

    destroy_simulation(sim);
}

TEST_CASE("Energy conservation during burns", "[hybrid][continuous][energy]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    double initial_energy = calculate_spacecraft_total_energy(craft, earth);

    INFO("Initial total energy: " << initial_energy << " J");

    double burn_duration = 6000.0;
    double total_dv = 150.0;
    int num_steps = 120;

    OrbitalElements current_orbit = craft->orbit;
    double dt_burn_step = burn_duration / num_steps;
    double dv_per_step = total_dv / num_steps;

    std::vector<double> energy_history;
    double max_energy_jump = 0.0;

    for (int i = 0; i < num_steps; i++) {
        Vec3 pos;
        Vec3 vel;
        orbital_elements_to_cartesian(current_orbit, earth->mass, &pos, &vel);

        Vec3 dir = get_burn_direction_vector(BURN_PROGRADE, pos, vel);
        Vec3 dv_vec = vec3_scale(dir, dv_per_step);
        vel = vec3_add(vel, dv_vec);

        current_orbit = cartesian_to_orbital_elements(pos, vel, earth->mass);
        current_orbit = propagate_orbital_elements(current_orbit, dt_burn_step, earth->mass);

        orbital_elements_to_cartesian(current_orbit, earth->mass, &pos, &vel);
        Spacecraft temp_craft = *craft;
        temp_craft.local_position = pos;
        temp_craft.local_velocity = vel;

        double current_energy = calculate_spacecraft_total_energy(&temp_craft, earth);
        energy_history.push_back(current_energy);

        if (i > 0) {
            double energy_jump = fabs(current_energy - energy_history[i - 1]);
            max_energy_jump = fmax(max_energy_jump, energy_jump);
        }
    }

    double final_energy = energy_history[num_steps - 1];
    double total_energy_change = final_energy - initial_energy;

    INFO("Final total energy: " << final_energy << " J");
    INFO("Total energy change: " << total_energy_change << " J");
    INFO("Max energy jump between steps: " << max_energy_jump << " J");

    REQUIRE(total_energy_change > 0);

    double expected_energy_change_approx = craft->mass * sqrt(G * earth->mass / craft->orbit.semi_major_axis) * total_dv;
    double relative_error = fabs(total_energy_change - expected_energy_change_approx) / expected_energy_change_approx;

    INFO("Expected approximate energy change: " << expected_energy_change_approx << " J");
    INFO("Relative error: " << relative_error * 100 << "%");

    REQUIRE(relative_error < 0.1);

    double average_step_energy_change = fabs(total_energy_change) / num_steps;
    double max_jump_ratio = max_energy_jump / average_step_energy_change;

    INFO("Average energy change per step: " << average_step_energy_change << " J");
    INFO("Max jump / average: " << max_jump_ratio);

    REQUIRE(max_jump_ratio < 10.0);

    destroy_simulation(sim);
}

TEST_CASE("Accuracy of continuous vs. impulsive burns", "[hybrid][continuous][accuracy]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    double initial_semi_major = craft->orbit.semi_major_axis;

    double burn_duration = 5000.0;
    double total_dv = 100.0;
    int num_steps_continuous = 100;

    OrbitalElements orbit_continuous = simulate_continuous_burn(craft->orbit, earth->mass,
                                                                total_dv, burn_duration,
                                                                num_steps_continuous, BURN_PROGRADE);

    OrbitalElements orbit_impulsive = simulate_continuous_burn(craft->orbit, earth->mass,
                                                              total_dv, burn_duration,
                                                              1, BURN_PROGRADE);

    INFO("Initial semi-major axis: " << initial_semi_major << " m");
    INFO("Continuous burn semi-major axis: " << orbit_continuous.semi_major_axis << " m");
    INFO("Impulsive burn semi-major axis: " << orbit_impulsive.semi_major_axis << " m");

    double difference_semi_major = fabs(orbit_continuous.semi_major_axis - orbit_impulsive.semi_major_axis);
    double relative_difference = difference_semi_major / orbit_continuous.semi_major_axis * 100.0;

    INFO("Semi-major axis difference: " << difference_semi_major << " m");
    INFO("Relative difference: " << relative_difference << "%");

    REQUIRE(relative_difference < 1.0);

    double mu = G * earth->mass;
    double v_continuous = sqrt(mu / orbit_continuous.semi_major_axis);
    double v_impulsive = sqrt(mu / orbit_impulsive.semi_major_axis);
    double v_difference = fabs(v_continuous - v_impulsive);

    INFO("Continuous burn velocity: " << v_continuous << " m/s");
    INFO("Impulsive burn velocity: " << v_impulsive << " m/s");
    INFO("Velocity difference: " << v_difference << " m/s");

    REQUIRE(v_difference < 2.0);

    destroy_simulation(sim);
}

TEST_CASE("Propagation during continuous burn", "[hybrid][continuous][propagation]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    double burn_duration = 5000.0;
    double total_dv = 100.0;
    int num_steps = 100;

    OrbitalElements current_orbit = craft->orbit;
    double dt_burn_step = burn_duration / num_steps;
    double dv_per_step = total_dv / num_steps;

    std::vector<Vec3> positions;
    std::vector<double> times;

    for (int i = 0; i <= num_steps; i++) {
        Vec3 pos;
        Vec3 vel;
        orbital_elements_to_cartesian(current_orbit, earth->mass, &pos, &vel);

        positions.push_back(pos);
        times.push_back(i * dt_burn_step);

        if (i < num_steps) {
            Vec3 dir = get_burn_direction_vector(BURN_PROGRADE, pos, vel);
            Vec3 dv_vec = vec3_scale(dir, dv_per_step);
            vel = vec3_add(vel, dv_vec);

            current_orbit = cartesian_to_orbital_elements(pos, vel, earth->mass);
            current_orbit = propagate_orbital_elements(current_orbit, dt_burn_step, earth->mass);
        }
    }

    double total_path_length = 0.0;
    for (size_t i = 1; i < positions.size(); i++) {
        total_path_length += vec3_distance(positions[i - 1], positions[i]);
    }

    INFO("Total path length during burn: " << total_path_length << " m");

    Vec3 pos_start = positions[0];
    Vec3 pos_end = positions[num_steps];
    double straight_line_distance = vec3_distance(pos_start, pos_end);

    INFO("Straight-line distance: " << straight_line_distance << " m");

    REQUIRE(total_path_length > straight_line_distance);

    double initial_radius = vec3_magnitude(pos_start);
    double final_radius = vec3_magnitude(pos_end);

    INFO("Initial radius: " << initial_radius << " m");
    INFO("Final radius: " << final_radius << " m");

    REQUIRE(final_radius > initial_radius);

    double mu = G * earth->mass;
    double v_initial = sqrt(mu / craft->orbit.semi_major_axis);
    double epsilon_initial = -mu / (2.0 * craft->orbit.semi_major_axis);
    double epsilon_final = epsilon_initial + v_initial * total_dv;
    double a_expected = -mu / (2.0 * epsilon_final);

    INFO("Expected final semi-major axis: " << a_expected << " m");

    double r_at_periapsis = a_expected * (1.0 - craft->orbit.eccentricity);
    double r_at_apoapsis = a_expected * (1.0 + craft->orbit.eccentricity);

    INFO("Expected radius at periapsis: " << r_at_periapsis << " m");
    INFO("Expected radius at apoapsis: " << r_at_apoapsis << " m");

    REQUIRE(final_radius >= r_at_periapsis - 1.0e5);
    REQUIRE(final_radius <= r_at_apoapsis + 1.0e5);

    destroy_simulation(sim);
}

TEST_CASE("Numerical stability during many burn/conversion cycles", "[hybrid][continuous][stability]") {
    const double TIME_STEP = 60.0;

    SimulationState* sim = create_simulation(2, 4, 0, TIME_STEP);

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

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

    OrbitalElements initial_orbit = craft->orbit;
    double burn_duration = 5000.0;
    double total_dv = 100.0;
    int num_steps = 100;
    double dt_burn_step = burn_duration / num_steps;
    double dv_per_step = total_dv / num_steps;

    std::vector<double> semi_major_history;
    std::vector<double> eccentricity_history;

    OrbitalElements current_orbit = craft->orbit;

    for (int i = 0; i < num_steps; i++) {
        Vec3 pos;
        Vec3 vel;
        orbital_elements_to_cartesian(current_orbit, earth->mass, &pos, &vel);

        Vec3 dir = get_burn_direction_vector(BURN_PROGRADE, pos, vel);
        Vec3 dv_vec = vec3_scale(dir, dv_per_step);
        vel = vec3_add(vel, dv_vec);

        current_orbit = cartesian_to_orbital_elements(pos, vel, earth->mass);
        current_orbit = propagate_orbital_elements(current_orbit, dt_burn_step, earth->mass);

        semi_major_history.push_back(current_orbit.semi_major_axis);
        eccentricity_history.push_back(current_orbit.eccentricity);
    }

    bool monotonic_increase = true;
    for (size_t i = 1; i < semi_major_history.size(); i++) {
        if (semi_major_history[i] < semi_major_history[i - 1]) {
            monotonic_increase = false;
            break;
        }
    }

    INFO("Monotonic semi-major axis increase: " << (monotonic_increase ? "yes" : "no"));

    REQUIRE(monotonic_increase);

    double max_eccentricity = 0.0;
    double min_eccentricity = 1.0;

    for (size_t i = 0; i < eccentricity_history.size(); i++) {
        max_eccentricity = fmax(max_eccentricity, eccentricity_history[i]);
        min_eccentricity = fmin(min_eccentricity, eccentricity_history[i]);
    }

    INFO("Max eccentricity during burn: " << max_eccentricity);
    INFO("Min eccentricity during burn: " << min_eccentricity);

    REQUIRE(max_eccentricity < 0.1);

    double initial_semi_major = initial_orbit.semi_major_axis;
    double final_semi_major = semi_major_history[num_steps - 1];
    double total_change = final_semi_major - initial_semi_major;
    double average_change_per_step = total_change / num_steps;

    INFO("Total semi-major axis change: " << total_change << " m");
    INFO("Average change per step: " << average_change_per_step << " m");

    double max_deviation = 0.0;
    for (size_t i = 0; i < semi_major_history.size(); i++) {
        double expected = initial_semi_major + (i + 1) * average_change_per_step;
        double deviation = fabs(semi_major_history[i] - expected);
        max_deviation = fmax(max_deviation, deviation);
    }

    INFO("Max deviation from linear trend: " << max_deviation << " m");
    INFO("Relative deviation: " << (max_deviation / total_change * 100) << "%");

    REQUIRE(max_deviation < total_change * 0.5);

    destroy_simulation(sim);
}