--
-- Matrices
--

inductive pattern Matrix a:=
  | quadCons (Matrix a) (Matrix a) (Matrix a) (Matrix a)
  | matCons Integer Integer a (Matrix a) (Matrix a) (Matrix a) (Matrix a)

def matrix : Matcher (Matrix MathExpr) :=
  matcher
    | quadCons $ $ $ $ as (mathExpr, matrix, matrix, matrix) with
      | $tgt ->
        match tensorShape tgt as list integer with
          | $m :: $n :: _ ->
            [(tgt_1_1, tgt_1_(2, n), tgt_(2, m)_1, tgt_(2, m)_(2, n))]
          | _ -> []
    | matCons #$i #$j $ $ $ $ $ as (mathExpr, matrix, matrix, matrix, matrix) with
      | $tgt ->
        let ns := tensorShape tgt
            m := nth 1 ns
            n := nth 2 ns
         in [ ( tgt_i_j
            , tgt_(1, i - 1)_(1, j - 1)
            , tgt_(1, i - 1)_(j + 1, n)
            , tgt_(i + 1, m)_(1, j - 1)
            , tgt_(i + 1, m)_(j + 1, n) ) ]
    | #$val as () with
      | $tgt -> if val = tgt then [()] else []
    | $ as (something) with
      | $tgt -> [tgt]

def M.inverse (m: Matrix MathExpr) : Matrix MathExpr :=
  let d := M.det m
   in generateTensor
        (\[i, j] ->
          match m as matrix with
          | matCons #j #i _ $A $B $C $D ->
            if isEven (i + j)
              then M.det (M.join A B C D) / d
              else - (M.det (M.join A B C D) / d))
        (tensorShape m)

def M.* (s: Matrix MathExpr) (t: Matrix MathExpr) : Matrix MathExpr := 
  withSymbols [i, j, k] (s~i~j . t_j_k)

def M.*' (s: Matrix MathExpr) (t: Matrix MathExpr) : Matrix MathExpr := 
  withSymbols [i, j, k] (s~i~j .' t_j_k)

def M.power (t: Matrix MathExpr) (k: Integer) : Matrix MathExpr := 
  foldl M.* t (take (k - 1) (repeat1 t))

def M.comm (m1: Matrix MathExpr) (m2: Matrix MathExpr) : Matrix MathExpr := 
  withSymbols [i, j, k] m1~i~j . m2_j_k - m2~i~j . m1_j_k

def M.join (A: Matrix MathExpr) (B: Matrix MathExpr) (C: Matrix MathExpr) (D: Matrix MathExpr)
  : Matrix MathExpr :=
  let ashape := tensorShape A
      bshape := tensorShape B
      cshape := tensorShape C
      dshape := tensorShape D
  in let a1 := nth 1 ashape
         a2 := nth 2 ashape
         b1 := nth 1 bshape
         b2 := nth 2 bshape
         c1 := nth 1 cshape
         c2 := nth 2 cshape
         d1 := nth 1 dshape
         d2 := nth 2 dshape
     in let m1 := max a1 b1
            m2 := max a2 c2
            n1 := max c1 d1
            n2 := max b2 d2
        in generateTensor
             (\match as list integer with
               | [$i & ?(<= a1), $j & ?(<= a2)] -> A_i_j
               | [$i & ?(<= m1), $j]            -> B_i_(j - a2)
               | [$i,            $j & ?(<= m2)] -> C_(i - a1)_j
               | [$i,            $j]            -> D_(i - m1)_(j - m2))
             [m1 + n1, m2 + n2]

--
-- Determinant
--
def M.determinant (m: Matrix MathExpr) : MathExpr :=
  match tensorShape m as list integer with
    | [#0, #0] -> 1
    | [$n, #n] ->
      let (es, os) := evenAndOddPermutations' n
       in sum (map (\e -> product (map2 (\i j -> m_i_j) (between 1 n) e)) es) -
            sum (map (\o -> product (map2 (\i j -> m_i_j) (between 1 n) o)) os)
    | _ -> undefined

def M.det (m: Matrix MathExpr) : MathExpr := M.determinant m