--
-- Matrices
--

M.* %s %t := withSymbols [i, j, k] s~i~j . t_j
M.*' %s %t := withSymbols [i, j, k] s~i~j .' t_j

M.power %t n := foldl M.* t (take (n - 1) (repeat1 t))
--M.power %m n := repeatedSquaring M.* m n

M.comm %m1 %m2 := withSymbols [i, j, k] m1~i~j . m2_j_k - m2~i~j . m1_j_k

M.inverse %m :=
  let d := M.det m
   in generateTensor
        2#(match m as matrix with
          | cons #%2 #%1 _ $A $B $C $D ->
            if isEven (%1 + %2)
              then M.det (M.join A B C D) / d
              else - (M.det (M.join A B C D) / d))
        (tensorShape m)

trace %t := withSymbols [i] sum (contract t~i_i)

matrix :=
  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))]
          | _ -> []
    | cons #$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]

M.join %A %B %C %D :=
  let ashape := tensorShape A
      a1 := nth 1 ashape
      a2 := nth 2 ashape
      bshape := tensorShape B
      b1 := nth 1 bshape
      b2 := nth 2 bshape
      cshape := tensorShape C
      c1 := nth 1 cshape
      c2 := nth 2 cshape
      dshape := tensorShape D
      d1 := nth 1 dshape
      d2 := nth 2 dshape
      m1 := max a1 b1
      m2 := max a2 c2
      n1 := max c1 d1
      n2 := max b2 d2
   in generateTensor
        2#(match (%1, %2) as (integer, integer) with
          | (?(<= a1), ?(<= a2)) -> A_%1_%2
          | (?(<= m1), _) -> B_%1_(%2 - a2)
          | (_, ?(<= m2)) -> C_(%1 - a1)_%2
          | (_, _) -> D_(%1 - m1)_(%2 - m2))
        [m1 + n1, m2 + n2]

--
-- Determinant
--
evenAndOddPermutations n :=
  let (es, os) := evenAndOddPermutations' n
   in (map 1#(\i -> nth i %1) es, map 1#(\i -> nth i %1) os)

evenAndOddPermutations0 n :=
  let (es, os) := evenAndOddPermutations' n
   in ( map 1#(\i -> nth (i + 1) (map 1#(%1 - 1) %1)) es
      , map 1#(\i -> nth (i + 1) (map 1#(%1 - 1) %1)) os )

evenAndOddPermutations' n :=
  match n as integer with
    | #1 -> ([[1]], [])
    | #2 -> ([[1, 2]], [[2, 1]])
    | _ ->
      let (es, os) := evenAndOddPermutations' (n - 1)
          es' := map (++ [n]) es
          os' := map (++ [n]) os
       in ( es' ++ concat
                     (map
                        (\i -> map 1#(permutate i n %1) os')
                        (between 1 (n - 1)))
          , os' ++ concat
                     (map
                        (\i -> map 1#(permutate i n %1) es')
                        (between 1 (n - 1))) )

permutate x y xs :=
  match xs as list eq with
    | $hs ++ #x :: $ms ++ #y :: $ts -> hs ++ y :: ms ++ x :: ts
    | $hs ++ #y :: $ms ++ #x :: $ts -> hs ++ x :: ms ++ y :: ts

M.determinant %m :=
  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

M.det := M.determinant

--
-- Eigenvalues and eigenvectors
--
M.eigenvalues %m :=
  match tensorShape m as list integer with
    | #2 :: #2 :: [] ->
      let (e1, e2) := qF (M.det (T.- m (scalarToTensor x [2, 2]))) x
       in [e1, e2]
    | _ -> undefined

M.eigenvectors %m :=
  match tensorShape m as list integer with
    | #2 :: #2 :: [] ->
      let (e1, e2) := qF (M.det (T.- m (scalarToTensor x [2, 2]))) x
       in [ (e1, clearIndex (T.- m (scalarToTensor e1 [2, 2]))_i_1)
          , (e2, clearIndex (T.- m (scalarToTensor e2 [2, 2]))_i_1) ]
    | _ -> undefined

--
-- LU decomposition
--
M.LU %x :=
  match tensorShape x as list integer with
    | #2 :: #2 :: [] ->
      let L := generateTensor
                 2#(match compare %1 %2 as ordering with
                   | less -> 0
                   | equal -> 1
                   | greater -> b_%1_%2)
                 [2, 2]
          U := generateTensor
                 2#(match compare %1 %2 as ordering with
                   | greater -> 0
                   | _ -> c_%1_%2)
                 [2, 2]
          m := M.* L U
          ret := solve
                   [ (m_1_1, x_1_1, c_1_1)
                   , (m_1_2, x_1_2, c_1_2)
                   , (m_2_1, x_2_1, b_2_1)
                   , (m_2_2, x_2_2, c_2_2) ]
       in (substitute ret L, substitute ret U)
    | #3 :: #3 :: [] ->
      let L := generateTensor
                 2#(match compare %1 %2 as ordering with
                   | less -> 0
                   | equal -> 1
                   | greater -> b_%1_%2)
                 [3, 3]
          U := generateTensor
                 2#(match compare %1 %2 as ordering with
                   | greater -> 0
                   | _ -> c_%1_%2)
                 [3, 3]
          m := M.* L U
          ret := solve
                   [ (m_1_1, x_1_1, c_1_1)
                   , (m_1_2, x_1_2, c_1_2)
                   , (m_1_3, x_1_3, c_1_3)
                   , (m_2_1, x_2_1, b_2_1)
                   , (m_2_2, x_2_2, c_2_2)
                   , (m_2_3, x_2_3, c_2_3)
                   , (m_3_1, x_3_1, b_3_1)
                   , (m_3_2, x_3_2, b_3_2)
                   , (m_3_3, x_3_3, c_3_3) ]
       in (substitute ret L, substitute ret U)
    | _ -> undefined

--
-- Utility
--
generateMatrixFromQuadraticExpr f xs :=
  generateTensor
    2#(coefficient2 f (nth %1 xs) (nth %2 xs))
    [length xs, length xs]