const constants = @import("./constants.zig");
const vec3 = @import("./vec3.zig");
const drone = @import("./drone.zig");
const std = @import("std");

pub fn gravity(mass: f64) vec3.Vec3F64 {
    return .init(0, 0, mass * constants.g);
}

pub fn propThrust(attitude: vec3.QuatF64, p: drone.Propeller) struct { vec3.Vec3F64, vec3.Vec3F64 } {
    const body_thrust = vec3.Vec3F64.init(
        0,
        0,
        -p.thrust_coefficient * p.angular_velocity * p.angular_velocity,
    );

    const force = vec3.quatApply(attitude, body_thrust);

    // pitching/rolling moment induced by thrust force
    const force_moment = vec3.vec3Cross(p.cg_to_prop, body_thrust);

    // Yaw moment induced by roation of propeller
    const drag_moment = vec3.Vec3F64.init(
        0,
        0,
        p.moment_coefficient * p.angular_velocity * p.angular_velocity,
    );

    return .{ force, vec3.vec3Add(force_moment, drag_moment) };
}

pub fn drag(
    drag_coefficient: f64,
    velocity: vec3.Vec3F64,
) vec3.Vec3F64 {
    return .init(
        velocity.x() * drag_coefficient,
        velocity.y() * drag_coefficient,
        velocity.z() * drag_coefficient,
    );
}

test "gravity is correct" {
    const mass = 10;
    const fg = gravity(mass);
    try std.testing.expect(std.math.approxEqAbs(f64, fg.x(), 0, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, fg.y(), 0, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, fg.z(), constants.g * mass, 1e-12));
}

test "thrust is correct" {
    const attitude = vec3.yawPitchRollToQuat(0, 0, -std.math.pi / 2.0);
    const prop: drone.Propeller = .{
        .cg_to_prop = vec3.Vec3F64.init(1, 1, 0),
        .angular_velocity = 100,
        .thrust_coefficient = 1,
        .moment_coefficient = 2,
    };

    const expected_thrust_mag = -100 * 100;
    const thrust, const momentum = propThrust(attitude, prop);

    try std.testing.expectApproxEqAbs(0, thrust.x(), 1e-11);
    try std.testing.expectApproxEqAbs(expected_thrust_mag, thrust.y(), 1e-11);
    try std.testing.expectApproxEqAbs(0, thrust.z(), 1e-11);
    try std.testing.expectApproxEqAbs(expected_thrust_mag, momentum.x(), 1e-11);
    try std.testing.expectApproxEqAbs(-expected_thrust_mag, momentum.y(), 1e-11);
    try std.testing.expectApproxEqAbs(100 * 100 * 2, momentum.z(), 1e-11);
}

test "drag is correct" {
    const velocity = vec3.Vec3F64.init(0, 1, 1);
    const f_drag = drag(0.5, velocity);

    try std.testing.expectApproxEqAbs(0, f_drag.x(), 1e-11);
    try std.testing.expectApproxEqAbs(0.5, f_drag.y(), 1e-11);
    try std.testing.expectApproxEqAbs(0.5, f_drag.z(), 1e-11);
}