#include "physics.h"
#include "simulation.h"
#include <cmath>

// Vector addition
Vec3 vec3_add(Vec3 a, Vec3 b) {
    return {a.x + b.x, a.y + b.y, a.z + b.z};
}

// Vector subtraction
Vec3 vec3_sub(Vec3 a, Vec3 b) {
    return {a.x - b.x, a.y - b.y, a.z - b.z};
}

// Scalar multiplication
Vec3 vec3_scale(Vec3 v, double s) {
    return {v.x * s, v.y * s, v.z * s};
}

// Vector magnitude
double vec3_magnitude(Vec3 v) {
    return sqrt(v.x * v.x + v.y * v.y + v.z * v.z);
}

// Distance between two points
double vec3_distance(Vec3 a, Vec3 b) {
    Vec3 diff = vec3_sub(a, b);
    return vec3_magnitude(diff);
}

// Normalize vector to unit length
Vec3 vec3_normalize(Vec3 v) {
    double mag = vec3_magnitude(v);
    if (mag > 0.0) {
        return vec3_scale(v, 1.0 / mag);
    }
    return {0.0, 0.0, 0.0};
}

// Calculate gravitational force using Newton's law: F = G * m1 * m2 / r^2
Vec3 calculate_gravity_force(CelestialBody* body, CelestialBody* parent) {
    Vec3 r = vec3_sub(parent->position, body->position);
    double distance = vec3_magnitude(r);

    // Avoid division by zero
    if (distance < 1.0) {
        distance = 1.0;
    }

    double force_magnitude = G * body->mass * parent->mass / (distance * distance);
    Vec3 direction = vec3_normalize(r);

    return vec3_scale(direction, force_magnitude);
}

// Calculate acceleration from force: a = F / m
Vec3 calculate_acceleration(Vec3 force, double mass) {
    if (mass > 0.0) {
        return vec3_scale(force, 1.0 / mass);
    }
    return {0.0, 0.0, 0.0};
}

Vec3 evaluate_acceleration(Vec3 pos, Vec3 vel, AccelerationContext* ctx) {
    CelestialBody temp_body = *ctx->current_body;
    temp_body.position = pos;
    temp_body.velocity = vel;

    Vec3 total_force = {0.0, 0.0, 0.0};

    if (temp_body.parent_index == -1) {
        for (int j = 0; j < ctx->sim->body_count; j++) {
            if (j == ctx->body_index) continue;
            CelestialBody* other = &ctx->sim->bodies[j];
            if (other->parent_index == -1) {
                Vec3 force = calculate_gravity_force(&temp_body, other);
                total_force = vec3_add(total_force, force);
            }
        }
    } else {
        if (temp_body.parent_index >= 0 && temp_body.parent_index < ctx->sim->body_count) {
            CelestialBody* parent = &ctx->sim->bodies[temp_body.parent_index];
            total_force = calculate_gravity_force(&temp_body, parent);
        }
    }

    return calculate_acceleration(total_force, temp_body.mass);
}

void rk4_step(CelestialBody* body, AccelerationContext* ctx, double dt) {
    Vec3 k1_vel, k2_vel, k3_vel, k4_vel;
    Vec3 k1_pos, k2_pos, k3_pos, k4_pos;

    Vec3 pos0 = body->position;
    Vec3 vel0 = body->velocity;

    k1_vel = evaluate_acceleration(pos0, vel0, ctx);
    k1_pos = vel0;

    Vec3 pos1 = vec3_add(pos0, vec3_scale(k1_pos, dt * 0.5));
    Vec3 vel1 = vec3_add(vel0, vec3_scale(k1_vel, dt * 0.5));
    k2_vel = evaluate_acceleration(pos1, vel1, ctx);
    k2_pos = vel1;

    Vec3 pos2 = vec3_add(pos0, vec3_scale(k2_pos, dt * 0.5));
    Vec3 vel2 = vec3_add(vel0, vec3_scale(k2_vel, dt * 0.5));
    k3_vel = evaluate_acceleration(pos2, vel2, ctx);
    k3_pos = vel2;

    Vec3 pos3 = vec3_add(pos0, vec3_scale(k3_pos, dt));
    Vec3 vel3 = vec3_add(vel0, vec3_scale(k3_vel, dt));
    k4_vel = evaluate_acceleration(pos3, vel3, ctx);
    k4_pos = vel3;

    Vec3 k_vel_sum = vec3_add(vec3_add(k1_vel, vec3_scale(k2_vel, 2.0)),
                              vec3_add(vec3_scale(k3_vel, 2.0), k4_vel));
    Vec3 k_pos_sum = vec3_add(vec3_add(k1_pos, vec3_scale(k2_pos, 2.0)),
                              vec3_add(vec3_scale(k3_pos, 2.0), k4_pos));

    body->velocity = vec3_add(vel0, vec3_scale(k_vel_sum, dt / 6.0));
    body->position = vec3_add(pos0, vec3_scale(k_pos_sum, dt / 6.0));
}