import Base

# Mlib_linalg.bend — Mathlib.LinearAlgebra finite fragment: 2x2 Nat matrices.
#
# SCOPE CAP (honest): this is the Nat 2x2 computational core only.
# - Computational: add / mul / det / trace / eq over closed values.
# - det uses Nat.sub truncation: exact only when a*d >= b*c (all laws below
#   satisfy this; callers must ensure it).
# - General semiring laws (assoc/comm/distr over symbolic matrices) are NOT
#   claimed here: they belong to nat.bend delegation (add_comm/add_assoc) and
#   are left as future work. All laws below are closed instances that hold
#   by checker normalization ({==}).
# - F32 matrices are compute-only (F32 is axiomatic, unprovable by design).

type Mat2 is Data:
  MMat{a: Nat, b: Nat, c: Nat, d: Nat}

# --- constructors ---

def mmat_mk(a: Nat, b: Nat, c: Nat, d: Nat) -> Mat2:
  MMat{a, b, c, d}

def mmat_zero() -> Mat2:
  MMat{0n, 0n, 0n, 0n}

def mmat_id() -> Mat2:
  MMat{1n, 0n, 0n, 1n}

def mmat_example() -> Mat2:
  MMat{4n, 2n, 1n, 3n}

# --- extractors (match on parameters only) ---

def mmat_a(m: Mat2) -> Nat:
  match m:
    case MMat{a, b, c, d}:
      a

def mmat_b(m: Mat2) -> Nat:
  match m:
    case MMat{a, b, c, d}:
      b

def mmat_c(m: Mat2) -> Nat:
  match m:
    case MMat{a, b, c, d}:
      c

def mmat_d(m: Mat2) -> Nat:
  match m:
    case MMat{a, b, c, d}:
      d

# --- operations (fields marked + where reused) ---

def mmat_add(a: Mat2, b: Mat2) -> Mat2:
  match a b:
    case MMat{a1, b1, c1, d1} MMat{a2, b2, c2, d2}:
      MMat{Nat.add(a1, a2), Nat.add(b1, b2), Nat.add(c1, c2), Nat.add(d1, d2)}

def mmat_mul(+a: Mat2, +b: Mat2) -> Mat2:
  match a b:
    case MMat{+a1, +b1, +c1, +d1} MMat{+e1, +f1, +g1, +h1}:
      MMat{Nat.add(Nat.mul(a1, e1), Nat.mul(b1, g1)),
        Nat.add(Nat.mul(a1, f1), Nat.mul(b1, h1)),
        Nat.add(Nat.mul(c1, e1), Nat.mul(d1, g1)),
        Nat.add(Nat.mul(c1, f1), Nat.mul(d1, h1))}

def mmat_det(m: Mat2) -> Nat:
  match m:
    case MMat{a, b, c, d}:
      Nat.sub(Nat.mul(a, d), Nat.mul(b, c))

def mmat_trace(m: Mat2) -> Nat:
  match m:
    case MMat{a, b, c, d}:
      Nat.add(a, d)

def mmat_eq(a: Mat2, b: Mat2) -> Bool:
  match a b:
    case MMat{a1, b1, c1, d1} MMat{a2, b2, c2, d2}:
      Bool.and(Bool.and(Nat.is_eq(a1, a2), Nat.is_eq(b1, b2)),
        Bool.and(Nat.is_eq(c1, c2), Nat.is_eq(d1, d2)))

# --- closed laws (all definitional) ---

law mmat_add_1234_5678:
  {mmat_add(mmat_mk(1n, 2n, 3n, 4n), mmat_mk(5n, 6n, 7n, 8n))
    == mmat_mk(6n, 8n, 10n, 12n) : Mat2}

def mmat_add_1234_5678():
  {==}

law mmat_mul_12_34:
  {mmat_mul(mmat_mk(1n, 2n, 3n, 4n), mmat_mk(5n, 6n, 7n, 8n))
    == mmat_mk(19n, 22n, 43n, 50n) : Mat2}

def mmat_mul_12_34():
  {==}

law mmat_det_example:
  {mmat_det(mmat_example()) == 10n : Nat}

def mmat_det_example():
  {==}

law mmat_trace_example:
  {mmat_trace(mmat_example()) == 7n : Nat}

def mmat_trace_example():
  {==}

law mmat_id_mul:
  {mmat_mul(mmat_id(), mmat_example()) == mmat_example() : Mat2}

def mmat_id_mul():
  {==}

law mmat_add_zero:
  {mmat_add(mmat_example(), mmat_zero()) == mmat_example() : Mat2}

def mmat_add_zero():
  {==}

law mmat_eq_refl:
  {mmat_eq(mmat_example(), mmat_example()) == True{} : Bool}

def mmat_eq_refl():
  {==}