\import Algebra.Linear.Matrix
\import Algebra.Module
\import Algebra.Module.LinearMap
\import Algebra.Monoid
\import Algebra.Ring
\import Data.Array
\import Function.Meta
\import Logic
\import Meta
\import Paths
\import Paths.Meta
\import Set

\func Trace {R : Ring} {n : Nat} : LinearMap (MatrixModule R n n) (RingLModule R) \cowith
  | func M => BigSumHom (diag M)
  | func-+ => rewrite (diag.func-+, inv BigSumHom.func-+) idp
  | func-*c => rewrite (diag.func-*c, inv BigSumHom.func-*c) idp
  \where {
    \func diag : LinearMap (MatrixModule R n n) (ArrayLModule n (RingLModule R)) \cowith
      | func M => \lam i => M i i
      | func-+ => exts \lam _ => idp
      | func-*c => exts \lam _ => idp

    \func BigSumHom : LinearMap (ArrayLModule n (RingLModule R)) (RingLModule R) \cowith
      | func e => R.BigSum e
      | func-+ => R.BigSum_+
      | func-*c => inv LModule.*c_BigSum-ldistr
  }

\lemma Trace-transpose {R : Ring}{n : Nat} (M : MatrixModule R n n) : Trace M = Trace (transpose M)
  => idp

\open MatrixRing

\func Trace-prod {R : CRing}{n m : Nat}(A : MatrixModule R n m)(B : MatrixModule R m n)
  : Trace (A product B) = Trace (B product A)
  => rewrite (Trace-prod-unfold A B, R.BigSum-transpose, Trace-prod-unfold B A, help) idp
\where \protected {
  \lemma help
    : R.BigSum (\lam j => R.BigSum (\lam k => A k j R.* B j k)) =
      R.BigSum (\lam i => R.BigSum (\lam j => B i j R.* A j i))
    => R.BigSum-ext \lam j => R.BigSum-ext \lam i => R.*-comm

  \lemma Trace-prod-unfold {R : Ring} {n m : Nat} (X : MatrixModule R n m) (Y : MatrixModule R m n)
    : Trace (X product Y) = R.BigSum (\lam i => R.BigSum (\lam j => X i j R.* Y j i))
    => idp
}

\lemma Trace-conjugation {R : CRing} {n : Nat} {A B : Matrix R n n} {T : Monoid.Inv B} : Trace (T.inv product A product B) = Trace A
  => rewrite (Trace-prod (T.inv product A) B, inv $ product-assoc B T.inv A, T.inv-right, Ring.ide-left {MatrixRing R n}) idp

\lemma Trace-diagonal {R : Ring} {l : Array R} : Trace (diagonal l) = R.BigSum l
  => unfold (rewrite diag-diagonal idp)
\where {
  \lemma diag-diagonal {n : Nat} (l : Array R n) : Trace.diag (diagonal l) = l
    => exts \lam _ => mcases \with {
      | yes p => idp
      | no n1 => absurd (n1 idp)
    }
}

\lemma Trace-ide {R : Ring}{n : Nat} : Trace (Ring.ide {MatrixRing R n}) = R.natCoef n
  => rewrite (Trace-diagonal {R} {replicate n R.ide}, R.BigSum_replicate1) idp