const std = @import("std");
const math = std.math;
const testing = std.testing;
const Vector = @import("vector.zig").Vector;

pub fn Matrix(comptime N: usize) type {
  return packed struct {
    const Self = @This();
    pub const Scalar = f32;

    values: [N][N]Scalar,

    /// Initialies a matrix.
    pub fn init(values: [N][N]Scalar) Self {
      return .{ .values = values };
    }

    /// Creates a matrix filled with the given value.
    pub fn filled(n: Scalar) Self {
      var result: Self = undefined;

      comptime var i = 0;
      inline while (i < N) : (i += 1) {
        comptime var j = 0;
        inline while (j < N) : (j += 1) {
          result.values[i][j] = n;
        }
      }

      return result;
    }

    /// A zero initialized matrix.
    pub const ZERO: Self = comptime Self.filled(0);

    /// An identity initialized matrix.
    pub const IDENTITY: Self = comptime brk: {
      var result: Self = undefined;

      comptime var i = 0;
      while (i < N) : (i += 1) {
        comptime var j = 0;
        while (j < N) : (j += 1) {
          result.values[i][j] = if (i == j) 1 else 0;
        }
      }

      break :brk result;
    };

    /// Transposes this matrix
    pub fn transpose(self: Self) Self {
      var result: Self = undefined;

      comptime var i = 0;
      inline while (i < N) : (i += 1) {
        comptime var j = 0;
        inline while (j < N) : (j += 1) {
          result.values[j][i] = self.values[i][j];
        }
      }

      return result;
    }

    pub fn invert(src: Self) ?Self {
        // https://github.com/stackgl/gl-mat4/blob/master/invert.js
        const a = @bitCast([16]f32, src.values);

        const a00 = a[0];
        const a01 = a[1];
        const a02 = a[2];
        const a03 = a[3];
        const a10 = a[4];
        const a11 = a[5];
        const a12 = a[6];
        const a13 = a[7];
        const a20 = a[8];
        const a21 = a[9];
        const a22 = a[10];
        const a23 = a[11];
        const a30 = a[12];
        const a31 = a[13];
        const a32 = a[14];
        const a33 = a[15];

        const b00 = a00 * a11 - a01 * a10;
        const b01 = a00 * a12 - a02 * a10;
        const b02 = a00 * a13 - a03 * a10;
        const b03 = a01 * a12 - a02 * a11;
        const b04 = a01 * a13 - a03 * a11;
        const b05 = a02 * a13 - a03 * a12;
        const b06 = a20 * a31 - a21 * a30;
        const b07 = a20 * a32 - a22 * a30;
        const b08 = a20 * a33 - a23 * a30;
        const b09 = a21 * a32 - a22 * a31;
        const b10 = a21 * a33 - a23 * a31;
        const b11 = a22 * a33 - a23 * a32;

        // Calculate the determinant
        var det = b00 * b11 - b01 * b10 + b02 * b09 + b03 * b08 - b04 * b07 + b05 * b06;

        if (std.math.approxEqAbs(f32, det, 0, 1e-8)) {
            return null;
        }
        det = 1.0 / det;

        const out = [16]f32{
            (a11 * b11 - a12 * b10 + a13 * b09) * det, // 0
            (a02 * b10 - a01 * b11 - a03 * b09) * det, // 1
            (a31 * b05 - a32 * b04 + a33 * b03) * det, // 2
            (a22 * b04 - a21 * b05 - a23 * b03) * det, // 3
            (a12 * b08 - a10 * b11 - a13 * b07) * det, // 4
            (a00 * b11 - a02 * b08 + a03 * b07) * det, // 5
            (a32 * b02 - a30 * b05 - a33 * b01) * det, // 6
            (a20 * b05 - a22 * b02 + a23 * b01) * det, // 7
            (a10 * b10 - a11 * b08 + a13 * b06) * det, // 8
            (a01 * b08 - a00 * b10 - a03 * b06) * det, // 9
            (a30 * b04 - a31 * b02 + a33 * b00) * det, // 10
            (a21 * b02 - a20 * b04 - a23 * b00) * det, // 11
            (a11 * b07 - a10 * b09 - a12 * b06) * det, // 12
            (a00 * b09 - a01 * b07 + a02 * b06) * det, // 13
            (a31 * b01 - a30 * b03 - a32 * b00) * det, // 14
            (a20 * b03 - a21 * b01 + a22 * b00) * det, // 15
        };
        return Self{
            .values = @bitCast([4][4]f32, out),
        };
    }

    /// Multiplies 2 matrices together.
    pub fn mul(self: Self, other: Self) Self {
      var result: Self = undefined;

      const a = self.transpose();
      const b = other;

      comptime var i = 0;
      inline while (i < N) : (i += 1) {
        comptime var j = 0;
        inline while (j < N) : (j += 1) {
          const row: @Vector(N, Scalar) = a.values[j];
          const column: @Vector(N, Scalar) = b.values[i];
          const products: [N]Scalar = row * column;

          var sum = @floatCast(f32, 0);
          comptime var k = 0;
          inline while (k < N) : (k += 1) {
            sum += products[k];
          }

          result.values[i][j] = sum;
        }
      }

      return result;
    }

    /// Multiplies 2 matrices together.
    pub fn mulAssign(self: *Self, other: Self) void {
      const a = self.transpose();
      const b = other;

      comptime var i = 0;
      inline while (i < N) : (i += 1) {
        comptime var j = 0;
        inline while (j < N) : (j += 1) {
          const row: @Vector(N, Scalar) = a.values[j];
          const column: @Vector(N, Scalar) = b.values[i];
          const products: [N]Scalar = row * column;

          var sum = @floatCast(f32, 0);
          comptime var k = 0;
          inline while (k < N) : (k += 1) {
            sum += products[k];
          }

          self.values[i][j] = sum;
        }
      }
    }

    const VecN = Vector(N);
    pub fn apply(self: Self, vec: VecN) VecN {
      var result: VecN = undefined;

      comptime var i = 0;
      inline while (i < N) : (i += 1) {
        result.values[i] = 0;

        comptime var j = 0;
        inline while (j < N) : (j += 1) {
          result.values[i] += self.values[i][j] * vec.values[j];
        }
      }

      return result;
    }
  };
}

fn expectMatrixEqual(comptime N: usize, expected: [N][N]f32, actual: Matrix(N)) void {
  comptime var i = 0;
  inline while (i < N) : (i += 1) {
    comptime var j = 0;
    inline while (j < N) : (j += 1) {
      testing.expectEqual(expected[i][j], actual.values[i][j]);
    }
  }
}

test "matrix initialization" {
  const A = Matrix(2).init(.{ .{ 1, 2 }, .{ 3, 4 } });
  expectMatrixEqual(2, [2][2]f32{ .{ 1, 2, }, .{ 3, 4 } }, A);
}

test "matrix zero" {
  expectMatrixEqual(3, Matrix(3).filled(0).values, Matrix(3).ZERO);
}

test "matrix identity" {
  expectMatrixEqual(3, [3][3]f32{ .{ 1, 0, 0 }, .{ 0, 1 , 0 }, .{ 0, 0, 1 } }, Matrix(3).IDENTITY);
}

test "matrix transpose" {
  const actual = Matrix(3).init(.{ .{ 1, 2, 3 }, .{ 4, 5, 6 }, .{ 7, 8, 9 } });
  const expected = [3][3]f32{ .{ 1, 4, 7 }, .{  2, 5, 8 }, .{ 3, 6, 9 } };
  expectMatrixEqual(3, expected, actual.transpose());
}

test "matrix multiplication" {
  const A = Matrix(2).init(.{ .{ 1, 2 }, .{ 3, 4 } });
  const B = Matrix(2).init(.{ .{ 6, 7 }, .{ 8, 9 } });
  const expected = [2][2]f32{ .{ 27, 40 }, .{ 35, 52 } };
  expectMatrixEqual(2, expected, A.mul(B));
}
