const std = @import("std");

pub const Vec3F64 = struct {
    data: @Vector(4, f64),

    pub fn init(x_value: f64, y_value: f64, z_value: f64) Vec3F64 {
        return .{ .data = .{ x_value, y_value, z_value, 0 } };
    }

    pub fn x(self: Vec3F64) f64 {
        return self.data[0];
    }

    pub fn y(self: Vec3F64) f64 {
        return self.data[1];
    }

    pub fn z(self: Vec3F64) f64 {
        return self.data[2];
    }
};

pub fn vec3Add(a: Vec3F64, b: Vec3F64) Vec3F64 {
    return .{ .data = a.data + b.data };
}

pub fn vec3Sub(a: Vec3F64, b: Vec3F64) Vec3F64 {
    return .{ .data = a.data - b.data };
}

pub fn vec3Dot(a: Vec3F64, b: Vec3F64) f64 {
    return @reduce(.Add, a.data * b.data);
}

pub fn vec3Cross(a: Vec3F64, b: Vec3F64) Vec3F64 {
    const x_component = a.y() * b.z() - a.z() * b.y();
    const y_component = a.z() * b.x() - a.x() * b.z();
    const z_component = a.x() * b.y() - a.y() * b.x();
    return .init(x_component, y_component, z_component);
}

pub const Mat3F64 = struct {
    row1: Vec3F64,
    row2: Vec3F64,
    row3: Vec3F64,
};

pub fn mat3Add(a: Mat3F64, b: Mat3F64) Mat3F64 {
    return .{
        .row1 = vec3Add(a.row1, b.row1),
        .row2 = vec3Add(a.row2, b.row2),
        .row3 = vec3Add(a.row3, b.row3),
    };
}

pub fn mat3Sub(a: Mat3F64, b: Mat3F64) Mat3F64 {
    return .{
        .row1 = vec3Sub(a.row1, b.row1),
        .row2 = vec3Sub(a.row2, b.row2),
        .row3 = vec3Sub(a.row3, b.row3),
    };
}

pub fn mat3Mul(a: Mat3F64, b: Mat3F64) Mat3F64 {
    const c_row1: Vec3F64 = .init(
        a.row1.x() * b.row1.x() + a.row1.y() * b.row2.x() + a.row1.z() * b.row3.x(),
        a.row1.x() * b.row1.y() + a.row1.y() * b.row2.y() + a.row1.z() * b.row3.y(),
        a.row1.x() * b.row1.z() + a.row1.y() * b.row2.z() + a.row1.z() * b.row3.z(),
    );

    const c_row2: Vec3F64 = .init(
        a.row2.x() * b.row1.x() + a.row2.y() * b.row2.x() + a.row2.z() * b.row3.x(),
        a.row2.x() * b.row1.y() + a.row2.y() * b.row2.y() + a.row2.z() * b.row3.y(),
        a.row2.x() * b.row1.z() + a.row2.y() * b.row2.z() + a.row2.z() * b.row3.z(),
    );

    const c_row3: Vec3F64 = .init(
        a.row3.x() * b.row1.x() + a.row3.y() * b.row2.x() + a.row3.z() * b.row3.x(),
        a.row3.x() * b.row1.y() + a.row3.y() * b.row2.y() + a.row3.z() * b.row3.y(),
        a.row3.x() * b.row1.z() + a.row3.y() * b.row2.z() + a.row3.z() * b.row3.z(),
    );

    return .{ .row1 = c_row1, .row2 = c_row2, .row3 = c_row3 };
}

pub fn vec3MulMat3(a: Mat3F64, b: Vec3F64) Vec3F64 {
    return .init(
        vec3Dot(a.row1, b),
        vec3Dot(a.row2, b),
        vec3Dot(a.row3, b),
    );
}

// Uses the Quaternion definition of q = w + xi + yj + zk
// and with i^2 = j^2 = k^2 = ijk = -1
// I grabbed a lot of this math from https://imadrahmoune.com/rotations-with-quaternions/
pub const QuatF64 = struct {
    data: @Vector(4, f64),
    pub fn init(x_value: f64, y_value: f64, z_value: f64, w_value: f64) QuatF64 {
        return .{ .data = .{ x_value, y_value, z_value, w_value } };
    }
    pub fn x(self: QuatF64) f64 {
        return self.data[0];
    }

    pub fn y(self: QuatF64) f64 {
        return self.data[1];
    }

    pub fn z(self: QuatF64) f64 {
        return self.data[2];
    }
    pub fn w(self: QuatF64) f64 {
        return self.data[3];
    }
    pub fn mag(q: QuatF64) f64 {
        return std.math.sqrt(@reduce(.Add, q.data * q.data));
    }

    pub fn normalized(q: QuatF64) QuatF64 {
        const magnitude = q.mag();
        return .{ .data = q.data / @as(@Vector(4, f64), @splat(magnitude)) };
    }

    pub fn conjugate(q: QuatF64) QuatF64 {
        return .{ .data = q.data * @as(@Vector(4, f64), .{ -1, -1, -1, 1 }) };
    }
};

pub fn yawPitchRollToQuat(yaw_radians: f64, pitch_radians: f64, roll_radians: f64) QuatF64 {
    const cos_roll = @cos(roll_radians / 2);
    const sin_roll = @sin(roll_radians / 2);
    const cos_pitch = @cos(pitch_radians / 2);
    const sin_pitch = @sin(pitch_radians / 2);
    const cos_yaw = @cos(yaw_radians / 2);
    const sin_yaw = @sin(yaw_radians / 2);

    const quat: QuatF64 = .init(
        sin_roll * cos_pitch * cos_yaw - cos_roll * sin_pitch * sin_yaw,
        cos_roll * sin_pitch * cos_yaw + sin_roll * cos_pitch * sin_yaw,
        cos_roll * cos_pitch * sin_yaw - sin_roll * sin_pitch * cos_yaw,
        cos_roll * cos_pitch * cos_yaw + sin_roll * sin_pitch * sin_yaw,
    );
    return quat.normalized();
}

pub fn quatMul(a: QuatF64, b: QuatF64) QuatF64 {
    return .init(
        a.w() * b.x() + a.x() * b.w() + a.y() * b.z() - a.z() * b.y(),
        a.w() * b.y() - a.x() * b.z() + a.y() * b.w() + a.z() * b.x(),
        a.w() * b.z() + a.x() * b.y() - a.y() * b.x() + a.z() * b.w(),
        a.w() * b.w() - a.x() * b.x() - a.y() * b.y() - a.z() * b.z(),
    );
}

pub fn quatApply(a: QuatF64, b: Vec3F64) Vec3F64 {
    const rotated_quat = quatMul(quatMul(a, .init(b.x(), b.y(), b.z(), 0)), a.conjugate());
    return .init(rotated_quat.data[0], rotated_quat.data[1], rotated_quat.data[2]);
}

test "vec3Add adds properly" {
    const a = Vec3F64.init(1, 2, 3);
    const b = Vec3F64.init(3, 1, 0);
    const c = vec3Add(a, b);

    try std.testing.expect(std.math.approxEqAbs(f64, c.x(), 4, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.y(), 3, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.z(), 3, 1e-12));
}
test "vec3Sub subs properly" {
    const a = Vec3F64.init(1, 2, 3);
    const b = Vec3F64.init(3, 1, 0);
    const c = vec3Sub(a, b);

    try std.testing.expect(std.math.approxEqAbs(f64, c.x(), -2, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.y(), 1, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.z(), 3, 1e-12));
}
test "vec3Dot dots properly" {
    const a = Vec3F64.init(1, 2, 3);
    const b = Vec3F64.init(3, 1, 0);
    try std.testing.expect(std.math.approxEqAbs(f64, vec3Dot(a, b), 5, 1e-12));
}
test "vec3Cross crosses properly" {
    const a = Vec3F64.init(1, 2, 3);
    const b = Vec3F64.init(3, 1, 0);

    const c = vec3Cross(a, b);

    try std.testing.expect(c.x() == -3);
    try std.testing.expect(c.y() == 9);
    try std.testing.expect(c.z() == -5);
}

test "Mat3 add works" {
    const a: Mat3F64 = .{
        .row1 = Vec3F64.init(1, 2, 3),
        .row2 = Vec3F64.init(0, 1, 5),
        .row3 = Vec3F64.init(0, 0, 9),
    };
    const b: Mat3F64 = .{
        .row1 = Vec3F64.init(0, 1, 0),
        .row2 = Vec3F64.init(0, 1, 1),
        .row3 = Vec3F64.init(0, 0, 1),
    };

    const c = mat3Add(a, b);

    try std.testing.expect(c.row1.x() == 1);
    try std.testing.expect(c.row1.y() == 3);
    try std.testing.expect(c.row1.z() == 3);
    try std.testing.expect(c.row2.x() == 0);
    try std.testing.expect(c.row2.y() == 2);
    try std.testing.expect(c.row2.z() == 6);
    try std.testing.expect(c.row3.x() == 0);
    try std.testing.expect(c.row3.y() == 0);
    try std.testing.expect(c.row3.z() == 10);
}

test "Mat3 sub works" {
    const a: Mat3F64 = .{
        .row1 = Vec3F64.init(1, 2, 3),
        .row2 = Vec3F64.init(0, 1, 5),
        .row3 = Vec3F64.init(0, 0, 9),
    };
    const b: Mat3F64 = .{
        .row1 = Vec3F64.init(0, 1, 0),
        .row2 = Vec3F64.init(0, 1, 1),
        .row3 = Vec3F64.init(0, 0, 1),
    };

    const c = mat3Sub(a, b);

    try std.testing.expect(c.row1.x() == 1);
    try std.testing.expect(c.row1.y() == 1);
    try std.testing.expect(c.row1.z() == 3);
    try std.testing.expect(c.row2.x() == 0);
    try std.testing.expect(c.row2.y() == 0);
    try std.testing.expect(c.row2.z() == 4);
    try std.testing.expect(c.row3.x() == 0);
    try std.testing.expect(c.row3.y() == 0);
    try std.testing.expect(c.row3.z() == 8);
}

test "Mat3 mul works" {
    const a: Mat3F64 = .{
        .row1 = Vec3F64.init(1, 2, 3),
        .row2 = Vec3F64.init(0, 1, 5),
        .row3 = Vec3F64.init(0, 0, 9),
    };
    const b: Mat3F64 = .{
        .row1 = Vec3F64.init(0, 1, 0),
        .row2 = Vec3F64.init(0, 1, 1),
        .row3 = Vec3F64.init(0, 0, 1),
    };

    const c = mat3Mul(a, b);

    try std.testing.expect(c.row1.x() == 0);
    try std.testing.expect(c.row1.y() == 3);
    try std.testing.expect(c.row1.z() == 5);
    try std.testing.expect(c.row2.x() == 0);
    try std.testing.expect(c.row2.y() == 1);
    try std.testing.expect(c.row2.z() == 6);
    try std.testing.expect(c.row3.x() == 0);
    try std.testing.expect(c.row3.y() == 0);
    try std.testing.expect(c.row3.z() == 9);
}

test "vec3 mul mat3 works" {
    const b: Mat3F64 = .{
        .row1 = Vec3F64.init(0, 1, 0),
        .row2 = Vec3F64.init(0, 1, 1),
        .row3 = Vec3F64.init(0, 0, 1),
    };
    const a = Vec3F64.init(1, 2, 3);

    const c = vec3MulMat3(b, a);

    try std.testing.expect(std.math.approxEqAbs(f64, c.x(), 2, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.y(), 5, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.z(), 3, 1e-12));
}

test "quat magnitude" {
    const q: QuatF64 = .init(1.0, 2.0, 3.0, 4.0);
    try std.testing.expect(std.math.approxEqAbs(f64, q.mag(), std.math.sqrt(30.0), 1e-12));
}

test "quat normalized" {
    const q: QuatF64 = .init(1.0, 2.0, 3.0, 4.0);
    const expected: QuatF64 = .init(
        1.0 / std.math.sqrt(30.0),
        2.0 / std.math.sqrt(30.0),
        3.0 / std.math.sqrt(30.0),
        4.0 / std.math.sqrt(30.0),
    );

    const normalized = q.normalized();

    try std.testing.expect(std.math.approxEqAbs(f64, normalized.x(), expected.x(), 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, normalized.y(), expected.y(), 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, normalized.z(), expected.z(), 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, normalized.w(), expected.w(), 1e-12));
}

test "quat mul" {
    const a: QuatF64 = .init(2, 3, 4, 1);
    const b: QuatF64 = .init(0, 1, 0, 1);

    const c = quatMul(a, b);

    try std.testing.expect(std.math.approxEqAbs(f64, c.w(), -2, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.x(), -2, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.y(), 4, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, c.z(), 6, 1e-12));
}

test "quat euler" {
    const quat = yawPitchRollToQuat(0.5, 0.5, 1);
    try std.testing.expect(std.math.approxEqAbs(f64, quat.x(), 0.3963648, 1e-6));
    try std.testing.expect(std.math.approxEqAbs(f64, quat.y(), 0.3252922, 1e-6));
    try std.testing.expect(std.math.approxEqAbs(f64, quat.z(), 0.0954433, 1e-6));
    try std.testing.expect(std.math.approxEqAbs(f64, quat.w(), 0.8532119, 1e-6));
}

test "quat conjugate" {
    const quat: QuatF64 = .init(1, 1, 1, 1);
    const conj = quat.conjugate();
    try std.testing.expect(std.math.approxEqAbs(f64, conj.x(), -1, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, conj.y(), -1, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, conj.z(), -1, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, conj.w(), 1, 1e-12));
}

test "quat apply to vec " {
    const quat: QuatF64 = yawPitchRollToQuat(std.math.pi / 2.0, 0, 0);

    try std.testing.expect(std.math.approxEqAbs(f64, quat.x(), 0, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, quat.y(), 0, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, quat.z(), @sin(std.math.pi / 4.0), 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, quat.w(), @cos(std.math.pi / 4.0), 1e-12));

    const rotated = quatApply(quat, .init(1, 2, 1));

    try std.testing.expect(std.math.approxEqAbs(f64, rotated.x(), -2, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, rotated.y(), 1, 1e-12));
    try std.testing.expect(std.math.approxEqAbs(f64, rotated.z(), 1, 1e-12));

    const forward_body = Vec3F64.init(1, 0, 0);
    const forward_inertial = quatApply(quat, forward_body);

    try std.testing.expectApproxEqAbs(@as(f64, 0), forward_inertial.x(), 1e-12);
    try std.testing.expectApproxEqAbs(@as(f64, 1), forward_inertial.y(), 1e-12);
    try std.testing.expectApproxEqAbs(@as(f64, 0), forward_inertial.z(), 1e-12);
}