#include "physics.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};
}

// Cross product
Vec3 vec3_cross(Vec3 a, Vec3 b) {
    return {
        a.y * b.z - a.z * b.y,
        a.z * b.x - a.x * b.z,
        a.x * b.y - a.y * b.x
    };
}

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

// Dot product
double vec3_dot(Vec3 a, Vec3 b) {
    return a.x * b.x + a.y * b.y + a.z * b.z;
}

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

void rk4_step(Vec3* position, Vec3* velocity, double dt,
              double body_mass, double parent_mass) {
    Vec3 k1_vel, k2_vel, k3_vel, k4_vel;
    Vec3 k1_pos, k2_pos, k3_pos, k4_pos;

    Vec3 pos0 = *position;
    Vec3 vel0 = *velocity;

    k1_vel = evaluate_acceleration(pos0, body_mass, parent_mass);
    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, body_mass, parent_mass);
    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, body_mass, parent_mass);
    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, body_mass, parent_mass);
    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));

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

Vec3 evaluate_acceleration(Vec3 relative_pos, double body_mass, double parent_mass) {
    Vec3 total_force = {0.0, 0.0, 0.0};

    double distance = vec3_magnitude(relative_pos);
    if (distance < 1.0) {
        distance = 1.0;
    }

    double force_magnitude = G * body_mass * parent_mass / (distance * distance);
    Vec3 direction = vec3_normalize(vec3_scale(relative_pos, -1.0));
    total_force = vec3_scale(direction, force_magnitude);

    return calculate_acceleration(total_force, body_mass);
}

Mat3 mat3_identity() {
    return {1.0, 0.0, 0.0,
            0.0, 1.0, 0.0,
            0.0, 0.0, 1.0};
}

Mat3 mat3_multiply(Mat3 a, Mat3 b) {
    return {
        a.m00 * b.m00 + a.m01 * b.m10 + a.m02 * b.m20,
        a.m00 * b.m01 + a.m01 * b.m11 + a.m02 * b.m21,
        a.m00 * b.m02 + a.m01 * b.m12 + a.m02 * b.m22,
        a.m10 * b.m00 + a.m11 * b.m10 + a.m12 * b.m20,
        a.m10 * b.m01 + a.m11 * b.m11 + a.m12 * b.m21,
        a.m10 * b.m02 + a.m11 * b.m12 + a.m12 * b.m22,
        a.m20 * b.m00 + a.m21 * b.m10 + a.m22 * b.m20,
        a.m20 * b.m01 + a.m21 * b.m11 + a.m22 * b.m21,
        a.m20 * b.m02 + a.m21 * b.m12 + a.m22 * b.m22
    };
}

Vec3 mat3_multiply_vec3(Mat3 m, Vec3 v) {
    return {
        m.m00 * v.x + m.m01 * v.y + m.m02 * v.z,
        m.m10 * v.x + m.m11 * v.y + m.m12 * v.z,
        m.m20 * v.x + m.m21 * v.y + m.m22 * v.z
    };
}

Mat3 mat3_rotation_x(double angle) {
    double c = cos(angle);
    double s = sin(angle);
    return {
        1.0, 0.0, 0.0,
        0.0, c, -s,
        0.0, s, c
    };
}

Mat3 mat3_rotation_z(double angle) {
    double c = cos(angle);
    double s = sin(angle);
    return {
        c, -s, 0.0,
        s, c, 0.0,
        0.0, 0.0, 1.0
    };
}

Mat3 mat3_rotation_orbital(double omega, double i, double Omega) {
    Mat3 Rz_omega = mat3_rotation_z(omega);
    Mat3 Rx_i = mat3_rotation_x(i);
    Mat3 Rz_Omega = mat3_rotation_z(Omega);

    Mat3 temp = mat3_multiply(Rx_i, Rz_omega);
    return mat3_multiply(Rz_Omega, temp);
}