const std = @import("std");

const mat4f = mat(4, 4, f32);
const vec4f = mat(4, 1, f32);

const mat4i = mat(4, 4, i32);
const vec4i = mat(4, 1, i32);

const mat4u = mat(4, 4, u32);
const vec4u = mat(4, 1, u32);

pub fn mat(comptime R_: usize, comptime C_: usize, comptime T_: type) type {
    return struct {
        pub const Rows = R_;
        pub const Cols = C_;
        pub const T = T_;

        data: [Cols][Rows]T,

        pub fn mul(l: @This(), r: anytype) MatMulReturnType(@This(), @TypeOf(r)) {
            return matmul(l, r);
        }
    };
}

fn MatMulReturnType(comptime L: type, comptime R: type) type {
    if (L.Cols != R.Rows) @compileError("invalid dimensions");

    const x: L.T = std.mem.zeroes(L.T);
    const y: R.T = std.mem.zeroes(R.T);
    const T = @TypeOf(x + y);

    return mat(L.Rows, R.Cols, T);
}

pub fn matmul(lhs: anytype, rhs: anytype) MatMulReturnType(@TypeOf(lhs), @TypeOf(rhs)) {
    @setFloatMode(.optimized);

    const L = @TypeOf(lhs);
    const R = @TypeOf(rhs);
    const Ret = MatMulReturnType(L, R);

    var res = std.mem.zeroes(Ret);

    if (L.Cols != R.Rows) @compileError("invalid dimensions");

    inline for (0..R.Cols) |col| {
        inline for (0..L.Rows) |row| {
            inline for (0..L.Cols) |k| {
                res.data[col][row] += lhs.data[k][row] * rhs.data[col][k];
            }
        }
    }

    return res;
}

export fn c_matmul_f_4x4_1x4(lhs: *const anyopaque, rhs: *const anyopaque, out: *anyopaque) void {
    const l: *const mat4f = @alignCast(@ptrCast(lhs));
    const r: *const vec4f = @alignCast(@ptrCast(rhs));
    const o: *vec4f = @alignCast(@ptrCast(out));

    o.* = matmul(l.*, r.*);
}

export fn c_matmul_i_4x4_1x4(lhs: *const anyopaque, rhs: *const anyopaque, out: *anyopaque) void {
    const l: *const mat4i = @alignCast(@ptrCast(lhs));
    const r: *const vec4i = @alignCast(@ptrCast(rhs));
    const o: *vec4i = @alignCast(@ptrCast(out));

    o.* = matmul(l.*, r.*);
}

export fn c_matmul_u_4x4_1x4(lhs: *const anyopaque, rhs: *const anyopaque, out: *anyopaque) void {
    const l: *const mat4u = @alignCast(@ptrCast(lhs));
    const r: *const vec4u = @alignCast(@ptrCast(rhs));
    const o: *vec4u = @alignCast(@ptrCast(out));

    o.* = matmul(l.*, r.*);
}

test "matmul" {
    // note, column major; it's transposed.
    const m: mat4u = .{ .data = .{
        .{ 85, 84, 87, 37 },
        .{ 33, 54, 49, 83 },
        .{ 96, 97, 3, 13 },
        .{ 69, 12, 45, 77 },
    } };
    const u: vec4u = .{ .data = .{.{ 37, 69, 94, 87 }} };

    const actual: vec4u = matmul(m, u);

    const expect: vec4u = .{ .data = .{.{ 20449, 16996, 10797, 15017 }} };

    try std.testing.expectEqualDeep(expect, actual);
}