const std = @import("std");

fn Matrix(comptime T:type) type {
    return struct {
        alloc: ?*std.mem.Allocator,
        data: []T,
        width: usize,
        height: usize,

        const Self = @This();

        pub fn init(allocator: *std.mem.Allocator, w: usize, h: usize) !Self {
            return Self {
                .alloc = allocator,
                .data = try allocator.alloc(T, w * h),
                .width = w,
                .height = h,
            };
        }

        pub fn deinit(self: *Self) void {
            if (self.alloc) |allocator| {
                allocator.free(self.data);
            }
        }

        pub fn fromSlice(slice: []T, w: usize) Self {
            return Self {
                .alloc = null,
                .data = slice,
                .width = w,
                .height = slice.len / w,
            };
        }

        pub fn get(self: *Self, x: usize, y: usize) T {
            return self.data[x*self.width + y];
        }

        pub fn set(self: *Self, x: usize, y: usize, item: T) void {
            self.data[x*self.width + y] = item;
        }
    };
}

fn mul(alloc: *std.mem.Allocator, a: *Matrix(i32), b: *Matrix(i32)) !Matrix(i32) {
    var result = try Matrix(i32).init(alloc, b.width, a.height);

    var i: usize = 0;
    while (i < a.height) : (i += 1) {
        var j: usize = 0;
        while (j < b.width) : (j += 1) {
            result.set(i, j, 0);

            var k: usize = 0;
            while (k < a.width) : (k += 1) {
                result.set(i, j, result.get(i, j) + a.get(i, k) * b.get(k, j));
            }
        }
    }

    return result;
}

fn print(a: *Matrix(i32)) void {
    var i: usize = 0;
    while (i < a.height) : (i += 1) {
        var j: usize = 0;
        while (j < a.width) : (j += 1) {
            std.debug.warn("{} ", a.get(i, j));
        }
        std.debug.warn("\n");
    }
}

pub fn main() !void {
    const arena = &std.heap.ArenaAllocator.init(std.heap.direct_allocator);
    defer arena.deinit();

    const allocator = &arena.allocator;

    var a_data = [_]i32{ 1, 2, 3, 4 };
    var a = Matrix(i32).fromSlice(a_data[0..], 2);

    var b_data = [_]i32{ 2, 5, 6, 7 };
    var b = Matrix(i32).fromSlice(b_data[0..], 2);

    var c = try mul(allocator, &a, &b);

    print(&a);
    print(&b);
    print(&c);
}