from random import initRand, rand
from nwo/utils import `...`
from mmath/smatrix import det, lu, lup, lu_det, lup, det, from_pivot, invert, identity, 
    newSMatrix, SMatrix, clone, tril, triu, lu_solve, `$`, `*`, `-`, `+`, `-=`, `+=`, `==`, `*`, norm2, newSMatrixFromArray,
    gauss_jordan_high, gauss_jordan_low, transpose, squared_norm2
from mmath/svector import buildSVector, `-`, abs, norm, Svector
from mmath/error import SingularMatrixError
from mmath/rational import newRational, Rational, one, zero, abs, `>`, `<`, `==`, `*`, `+`, `-`, `/`, `/=`, `+=`, `sqrt`
import unittest

suite "Nim linear algebra library":
    let test_matrix = newSMatrixFromArray[3, 3, float32]([1.0f32, 2f32, 3f32, 4f32, 5f32, 6f32, 8f32, 7f32, 9f32])

    test "+":
        let mtx1  = newSMatrixFromArray[3, 3, int8]([1i8,2,3,4,5,6,7,8,9])
        let mtx2  = newSMatrixFromArray[3, 3, int8]([-1i8,-2,-3,-4,-5,-6,-7,-8,-9])
        check(mtx1 + mtx2 == SMatrix[3, 3, int8]())

    test "-":
        let mtx  = newSMatrixFromArray[3, 3, int8]([1i8,2,3,4,5,6,7,8,9])
        check(mtx - mtx == SMatrix[3, 3, int8]())

    test "+=":
        var mtx = newSMatrixFromArray[3, 3, int8]([1i8,2,3,4,5,6,7,8,9])
        mtx += mtx
        check(mtx == newSMatrixFromArray[3, 3, int8]([2i8,4,6,8,10,12,14,16,18]))

    test "-=":
        var mtx = newSMatrixFromArray[3, 3, int8]([1i8,2,3,4,5,6,7,8,9])
        mtx -= mtx
        check(mtx == SMatrix[3, 3, int8]())

    test "*":
        block:
            let mtx  = newSMatrixFromArray[3, 3, int8]([1i8,2,3,4,5,6,8,7,9])
            let res = mtx * buildSVector[3, int8](1i8, 2, 3)
            check(res == buildSVector[3, int8](14i8, 32, 49))
        block:
            let mtx = newSMatrixFromArray[2, 5, int]([1,4,8,2,5,7,3,6,9,0])
            let res = mtx * mtx.transpose()
            check(res == newSMatrixFromArray[2, 2, int]([110,85,85,175]))
        block:
            let mtx = newSMatrixFromArray[2, 5, int]([1,4,8,2,5,7,3,6,9,0])
            let res = mtx.transpose() * mtx
            check(res == newSMatrixFromArray[5, 5, int]([
                50, 25, 50, 65, 5,
                25, 25, 50, 35, 20,
                50, 50, 100, 70, 40,
                65, 35, 70, 85, 10,
                5, 20, 40, 10, 25]))
    
        
    test "Determinant":
        check(test_matrix.det() == -9f32)

    test "LU Determinant":
        check(test_matrix.lu_det() == -9f32)

    test "Inverse":
        let err = test_matrix * test_matrix.invert() - identity[3, float32]()
        check(err.norm2() < 1e-5)
    
    test "LU decomposition":
        var rng = initRand(101325)
        var arr : array[0..(25 - 1), Rational[int64]]
        for i in 0..<arr.len:
            arr[i] = newRational(rng.rand(-20..20).int64, 20.int64)
        let mtx = newSMatrixFromArray[5, 5, Rational[int64]](arr)
        var lu = mtx.clone()
        let pivot = lu.lup()
        let l = lu.tril(Rational[int64].one)
        let u = lu.triu()
        let err = pivot * mtx - (l * u)
        check(err.squared_norm2() == Rational[int64].zero)

    test "Linear system solve":
        var rng = initRand(101325)
        var arr : array[0..(100 * 100 - 1), float64]
        for i in 0..<arr.len:
            arr[i] = rng.rand(-100f32..100f32)
        var mtx  = newSMatrixFromArray[100, 100, float64](arr)
        var lu = mtx.clone()
        let pivot = lu.lup()
        var b : SVector[100, float64]
        for i in 0...len(b):
            b[i] = float64(rng.rand(b.len))

        let x = lu.lu_solve(b, pivot)
        let error  = (mtx * x) - b
        check(error.norm() < 1e-5)
    
    test "triu":
        let mtx  = newSMatrixFromArray[3, 3, int8]([1i8,2,3,4,5,6,7,8,9])
        block:
            let upper = mtx.triu()
            check(upper == newSMatrixFromArray[3, 3, int8]([1i8,2,3,0,5,6,0,0,9]))
        block:
            let upper = mtx.triu(1)
            check(upper == newSMatrixFromArray[3, 3, int8]([1i8,2,3,0,1,6,0,0,1]))

    test "tril":
        let mtx  = newSMatrixFromArray[3, 3, uint]([1u,2,3,4,5,6,7,8,9])
        block:
            let lower = mtx.tril()
            check(lower == newSMatrixFromArray[3, 3, uint]([1u,0,0,4,5,0,7,8,9]))
        block:
            let lower = mtx.tril(1)
            check(lower == newSMatrixFromArray[3, 3, uint]([1u,0,0,4,1,0,7,8,1]))

    test "transpose":
        block:
            let mtx = newSMatrixFromArray[3, 3, int]([1, 2, 3, 4, 5, 6, 8, 7, 9])
            let xpose = newSMatrixFromArray[3, 3, int]([1,4,8,2,5,7,3,6,9])
            check(mtx.transpose() == xpose)
        block:
            let mtx = newSMatrixFromArray[2, 5, int]([1,4,8,2,5,7,3,6,9,0])
            let xpose = newSMatrixFromArray[5, 2, int]([1,7,4,3,8,6,2,9,5,0])
            check(mtx.transpose() == xpose)