import linalg

-- Check that the optimized matmul gives the same answers as the naive one
amat = for i:(Fin 100) j:(Fin 100). n_to_f $ ordinal (i, j)

:p tiled_matmul(amat, amat) ~~ naive_matmul amat amat
> True

-- Check that the inverse of the inverse is identity.
mat = [[11.,9.,24.,2.],[1.,5.,2.,6.],[3.,17.,18.,1.],[2.,5.,7.,1.]]
:p mat ~~ (invert (invert mat))
> True

-- Check that solving gives the inverse.
v = [1., 2., 3., 4.]
:p v ~~ (mat **. (solve mat v))
> True

-- Check that det and exp(logdet) are the same.
(s, logdet) = sign_and_log_determinant mat
:p (determinant mat) ~~ (s * (exp logdet))
> True

-- Matrix integer powers.
:p matrix_power mat 0 ~~ eye
> True
:p matrix_power mat 1 ~~ mat
> True
:p matrix_power mat 2 ~~ (mat ** mat)
> True
:p matrix_power mat 5 ~~ (mat ** mat ** mat ** mat ** mat)
> True

:p trace mat == (11. + 5. + 18. + 1.)
> True

-- Check that we can linearize LU decomposition
-- This is a regression test for Issue #842.
snd(linearize (\x. snd $ sign_and_log_determinant [[x]]) 1.0)(2.0)
> 2.

-- Check that we can differentiate through LU decomposition
-- This is a regression test for Issue #848.
grad (\x. (pivotize [[x]]).sign) 1.0
> 0.

grad (\x. snd $ sign_and_log_determinant [[x]]) 2.0
> 0.5

-- Check forward_substitute solve by comparing
-- against zero-padding and doing the full solve.
def padLowerTriMat(mat:LowerTriMat n v) -> n=>n=>v given (n|Ix, v|Add) =
  for i j.
    if (ordinal j)<=(ordinal i)
      then mat[i,unsafe_project j]
      else zero

lower : LowerTriMat (Fin 4) Float = arb $ new_key 0
lower_padded = padLowerTriMat lower
vec : (Fin 4)=>Float = arb $ new_key 0

forward_substitute lower vec ~~ solve lower_padded vec
> True