from random import initRand, rand
from nwo/utils import `...`
from mmath/rational import newRational, Rational, one, zero, abs, `>`, `<`, `==`, `*`, `+`, `-`, `/`, `/=`, `+=`, `sqrt`
from mmath/hmatrix import det, lu, from_pivot, newHMatrix, HMatrix, lu_det, invert, identity,
    clone, tril, triu, lu_solve, `*`, `-`, `+`, `+=`, `-=`, `==`, norm2, transpose, lup, squared_norm2
from mmath/hvector import buildHVector, newHVector, `-`, abs, norm
import unittest

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

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

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

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

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

    test "*":
        block:
            let mtx  = newHMatrix[int8](3, 3, [1i8,2,3,4,5,6,8,7,9])
            let res = mtx * buildHVector[int8](1i8, 2, 3)
            check(res == buildHVector[int8](14i8, 32, 49))
        block:
            let mtx = newHMatrix[int](2, 5, [1,4,8,2,5,7,3,6,9,0])
            let res = mtx * mtx.transpose()
            check(res == newHMatrix[int](2, 2, [110,85,85,175]))
        block:
            let mtx = newHMatrix[int](2, 5, [1,4,8,2,5,7,3,6,9,0])
            let res = mtx.transpose() * mtx
            check(res == newHMatrix[int](5, 5, [
                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[float32](3)
        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 = newHMatrix[Rational[int64]](5, 5, 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  = newHMatrix[float64](100, 100, arr)
        var lu = mtx.clone()
        let pivot = lu.lup()
        var b = newHVector[float64](100, 0.0f)
        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  = newHMatrix[int8](3, 3, [1i8,2,3,4,5,6,7,8,9])
        block:
            let upper = mtx.triu()
            check(upper == newHMatrix[int8](3, 3, [1i8,2,3,0,5,6,0,0,9]))
        block:
            let upper = mtx.triu(1)
            check(upper == newHMatrix[int8](3, 3, [1i8,2,3,0,1,6,0,0,1]))

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

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