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

pub fn timestep(dt: f64, d: drone.Drone, props: []const drone.Propeller, s: *drone.State) void {
    var force = vec3.Vec3F64.init(0, 0, 0);
    var moment = vec3.Vec3F64.init(0, 0, 0);

    for (props) |prop| {
        const prop_force, const prop_moment = forces.propThrust(s.q, prop);
        force = vec3.vec3Add(force, prop_force);
        moment = vec3.vec3Add(moment, prop_moment);
    }

    force = vec3.vec3Add(force, forces.drag(d.drag_coeff, s.v));
    force = vec3.vec3Add(force, forces.gravity(d.mass_kg));

    const accel = vec3.vec3Scale(force, 1.0 / d.mass_kg);
    s.v = vec3.vec3Add(s.v, vec3.vec3Scale(accel, dt));
    s.x_cg = vec3.vec3Add(s.x_cg, vec3.vec3Scale(s.v, dt));

    const I_omega = vec3.vec3MulMat3(d.moment_of_inertia, s.omega);
    const omega_dot = vec3.vec3MulMat3(
        vec3.mat3Inv(d.moment_of_inertia),
        vec3.vec3Sub(moment, vec3.vec3Cross(s.omega, I_omega)),
    );

    s.omega = vec3.vec3Add(s.omega, vec3.vec3Scale(omega_dot, dt));
    const half_omega_quat: vec3.QuatF64 = .init(
        s.omega.x() * 0.5,
        s.omega.y() * 0.5,
        s.omega.z() * 0.5,
        0,
    );

    const qdot = vec3.quatMul(half_omega_quat, s.q);
    s.q = vec3.quatAdd(s.q, vec3.quatScale(qdot, dt)).normalized();
    return;
}

test "Validate timesteps" {
    const d = drone.Drone{
        .drag_coeff = 0.0,
        .mass_kg = 0.5,
        .moment_of_inertia = .{
            .row1 = .init(2, 0, 0),
            .row2 = .init(0, 0.5, 0),
            .row3 = .init(0, 0, 2),
        },
    };

    var state: drone.State = .{
        .x_cg = .init(0, 0, 0),
        .omega = .init(0, 0, 0),
        .q = vec3.yawPitchRollToQuat(0, 0, 0),
        .v = .init(0, 0, 0),
    };

    const prop: drone.Propeller = .{
        .angular_velocity = 100,
        .cg_to_prop = .init(1, 0, 0),
        .moment_coefficient = 0.1,
        .thrust_coefficient = 0.5,
        .direction = constants.PropSpinDirection.cw,
    };
    const props: [1]drone.Propeller = .{prop};

    timestep(0.005, d, &props, &state);

    try std.testing.expectApproxEqAbs(0, state.v.x(), 1e-12);
    try std.testing.expectApproxEqAbs(0, state.v.y(), 1e-12);
    try std.testing.expectApproxEqAbs(-49.95096675, state.v.z(), 1e-8);

    try std.testing.expectApproxEqAbs(0, state.x_cg.x(), 1e-12);
    try std.testing.expectApproxEqAbs(0, state.x_cg.y(), 1e-12);
    try std.testing.expectApproxEqAbs(-0.24975483375, state.x_cg.z(), 1e-11);

    try std.testing.expectApproxEqAbs(0, state.omega.x(), 1e-12);
    try std.testing.expectApproxEqAbs(50, state.omega.y(), 1e-12);
    try std.testing.expectApproxEqAbs(2.5, state.omega.z(), 1e-11);

    try std.testing.expectApproxEqAbs(0, state.q.x(), 1e-12);
    try std.testing.expectApproxEqAbs(0.12403234937465502, state.q.y(), 1e-12);
    try std.testing.expectApproxEqAbs(0.006201617468732751, state.q.z(), 1e-11);
    try std.testing.expectApproxEqAbs(0.9922587949972401, state.q.w(), 1e-11);

    timestep(0.005, d, &props, &state);

    try std.testing.expectApproxEqAbs(-1.23072190e+01, state.v.x(), 1e-7);
    try std.testing.expectApproxEqAbs(-7.69201185e-02, state.v.y(), 1e-8);
    try std.testing.expectApproxEqAbs(-9.83635311e+01, state.v.z(), 1e-7);

    try std.testing.expectApproxEqAbs(-6.15360950e-02, state.x_cg.x(), 1e-9);
    try std.testing.expectApproxEqAbs(-3.84600592e-04, state.x_cg.y(), 1e-9);
    try std.testing.expectApproxEqAbs(-7.41572489e-01, state.x_cg.z(), 1e-9);

    try std.testing.expectApproxEqAbs(-0.46875, state.omega.x(), 1e-12);
    try std.testing.expectApproxEqAbs(100, state.omega.y(), 1e-12);
    try std.testing.expectApproxEqAbs(5, state.omega.z(), 1e-11);

    try std.testing.expectApproxEqAbs(-0.0011280012096230429, state.q.x(), 1e-12);
    try std.testing.expectApproxEqAbs(0.36096743708693396, state.q.y(), 1e-12);
    try std.testing.expectApproxEqAbs(0.01790701920276581, state.q.z(), 1e-12);
    try std.testing.expectApproxEqAbs(0.9324057998744074, state.q.w(), 1e-12);
}