/* Special matrices */

package scientific.linear.special

import std.unittest.*
import std.unittest.testmacro.*
import std.math.*

import scientific.numbers.CNumber
import scientific.numbers.Float
import scientific.numbers.Real
import scientific.numbers.power
import scientific.numbers.Complex64
import scientific.numbers.approxEqualC
import scientific.linear.Vector
import scientific.linear.vector
import scientific.linear.zeros
import scientific.linear.Matrix
import scientific.linear.matrix
import scientific.linear.empty
import scientific.linear.approxEqual
import scientific.linear.arange
import scientific.linear.diag
import scientific.linear.toComplex64
import scientific.linear.expm
import scientific.linear.hstack
import scientific.linear.vstack
import scientific.linear.*
import scientific.stats.utils.comb
import scientific.utils.AssertionException


public func vander<T>(x: Vector<T>, N: Int64, increasing!:Bool = false): Matrix<T> where T <: Float<T> {
    let sz = x.size()
    let res = empty<T>(sz, N)
    for (i in 0..sz) {
        for (j in 0..N) {
            if (increasing) {
                res[i,j] = power(x[i], T.fromInt(j))
            } else {
                res[i,j] = power(x[i], T.fromInt(N-1-j))
            }
        }
    }
    return res
}

public func circulant<T>(x: Vector<T>): Matrix<T> where T <: CNumber<T> {
    let sz = x.size()
    let res = empty<T>(sz, sz)
    for (i in 0..sz) {
        for (j in 0..sz) {
            if (i + j < sz) {
                res[i,j] = x[i+j]
            } else {
                res[i,j] = x[i+j-sz]
            }
        }
    }
    return res
}

public func companion<T>(x: Vector<T>): Matrix<T> where T <: Float<T> {
    let sz = x.size() - 1
    let res = zeros<T>(sz, sz)
    for (i in 0..sz) {
        res[0,i] = T.fromInt(-1) * x[i+1] / x[0]
    }
    for (i in 1..sz) {
        res[i,i-1] = T.fromInt(1)
    }
    return res
}

public func fiedler<T>(x: Vector<T>): Matrix<T> where T <: Real<T> {
    let sz = x.size()
    let res = zeros<T>(sz, sz)
    for (i in 0..sz) {
        for (j in 0..sz) {
            res[i,j] = T.abs(x[i] - x[j])
        }
    }
    return res
}

public func hankel<T>(x: Vector<T>): Matrix<T> where T <: CNumber<T> {
    let sz = x.size()
    let res = zeros<T>(sz, sz)
    for (i in 0..sz) {
        for (j in 0..i+1) {
            res[j,i-j] = x[i]
        }
    }
    return res
}

public func hankel<T>(x: Vector<T>, y: Vector<T>): Matrix<T> where T <: CNumber<T> {
    let sx = x.size()
    let sy = y.size()
    let res = empty<T>(sx, sy)
    for (i in 0..sx) {
        for (j in 0..i+1) {
            res[j,i-j] = x[i]
        }
    }
    for (i in 0..sx) {
        for (j in sx-i..sy) {
            res[i,j] = y[j+i-sx+1]
        }
    }
    return res
}

public func hilbert<T>(N: Int64): Matrix<T> where T <: Float<T> {
    let res = empty<T>(N, N)
    for (i in 0..N) {
        for (j in 0..N) {
            res[i,j] = T.fromInt(1) / T.fromInt(i + j + 1)
        }
    }
    return res
}

public func leslie<T>(f: Vector<T>, s: Vector<T>): Matrix<T> where T <: CNumber<T> {
    let sz = f.size()
    if (s.size() != sz - 1) {
        throw AssertionException("leslie: size of vectors does not match")
    }

    let res = zeros<T>(sz, sz)
    for (i in 0..sz) {
        res[0,i] = f[i]
    }
    for (i in 0..sz-1) {
        res[i+1,i] = s[i]
    }
    return res
}

public func pascal(n: Int64, kind!:String="symmetric"): Matrix<Int64> {
    let res = zeros<Int64>(n, n)
    for (i in 0..n) {
        for (j in 0..n) {
            if (kind == "symmetric") {
                res[i,j] = comb(i+j, i)
            } else if (kind == "lower") {
                res[i,j] = comb(i, j)
            } else if (kind == "upper") {
                res[i,j] = comb(j, i)
            } else {
                throw AssertionException("pascal: unrecognized kind")
            }
        }
    }
    return res
}

public func dft(n: Int64): Matrix<Complex64> {
    var x = toComplex64(arange<Float64>(0.0, Float64(n), 1.0).toMatrix())
    let real: Float64 = Float64.getPI() * 1.0 / Float64(n)
    var tempj = Complex64(real, 0.0)
    var temp: Matrix<Complex64> = Complex64(0.0, -2.0) * (tempj * x)
    var omegas = matExp(temp)
    var result = getExp(omegas)
    return result
}

public func blockDiag<T>(a: Array<Matrix<T>>): Matrix<T> where T <: CNumber<T> {
    let size = a.size
    var arr = zeros<Int64>(size, 2)
    var rowCount: Int64 = 0
    var colCount: Int64 = 0
    for (i in 0..size-1) {
        rowCount = rowCount + a[i].getRows()
        colCount = colCount + a[i].getCols()
        arr[i+1,0] = rowCount
        arr[i+1,1] = colCount
    }
    var res = zeros<T>(rowCount+1, colCount+1)
    for (i in 0..size) {
        for (j in 0..a[i].getRows()) {
            for(k in 0..a[i].getCols()) {
                res[arr[i,0]+j, arr[i,1]+k] = a[i][j,k]
            }
        }
    }
    return res
}

public func toeplitz<T>(c: Vector<T>, r: Vector<T>): Matrix<T> where T <: CNumber<T> {
    let row = c.size()
    let col = r.size()
    var res = zeros<T>(row, col)
    for (i in 0..row) {
        res[i,0] = c[i]
    }
    for (i in 0..col) {
        res[0,i] = r[i]
    }
    for (i in 0..col) {
        for (j in 0..row) {
            if (i + j < col) {
                res[j,i+j] = res[0,i]
            }
        }
    }
    for (i in 0..row) {
        for (j in 0..col) {
            if (i + j < row) {
                res[i+j,j] = res[i,0]
            }
        }
    }
    return res
}

public func convolutionMatrix<T>(a: Vector<T>, n: Int64, mode!:String="full"): Matrix<T> where T <: CNumber<T> {
    var m = a.size()
    var row: Int64 = 0
    if (mode == "full") {
        row = m + n - 1
    } else if (mode == "same") {
        row = max(m, n)
    } else if (mode == "valid") {
        row = max(m, n) - min(m, n) + 1
    }
    var col = n
    var res = zeros<T>(row, col)
    if (mode == "full") {
        var cvector = zeros<T>(row)
        var rvector = zeros<T>(col)
        for (i in 0..m) {
            cvector[i] = a[i]
        }
        rvector[0] = a[0]
        return toeplitz(cvector, rvector)
    } else if (mode == "same") {
        var cvector = zeros<T>(row)
        var rvector = zeros<T>(col)
        var mid = m / 2
        var temp = 0
        if (m % 2 == 0) {
            temp = mid - 1
        } else {
            temp = mid
        }
        for (i in 0..temp+1) {
            rvector[i] = a[temp-i]
        }
        for (i in 0..temp+1) {
            cvector[i] = a[temp+i]
        }
        return toeplitz(cvector, rvector)
    } else {
        var cvector = zeros<T>(row)
        var rvector = zeros<T>(col)
        for (i in 0..m) {
            rvector[i] = a[m-i-1]
        }
        cvector[0] = a[m-1]
        return toeplitz(cvector, rvector)
    }
}

public func hadamard(a: Int64): Matrix<Int64> {
    var H = matrix<Int64>([[1]])
    var lg2 = Int64(log2(Float64(a)))
    for (i in 0..lg2) {
        var MH = H * (-1) 
        var temp = hstack(H, H)
        var temp2 = hstack(H, MH)
        H = vstack(temp, temp2)
    }
    return H
}

public func helmert(n: Int64): Matrix<Float64> {
    var temp = zeros<Int64>(n, n)
    for (i in 0..n) {
        for (j in 0..n) {
            temp[i,j] = 1
        }
    }
    var H = tril(temp, k: -1) - diag(arange(0, n, 1))
    var d = arange(0, n, 1) * arange(1, n+1, 1)
    d[0] = n
    for (j in 0..n) {
        H[0,j] = 1
    }
    var res = zeros<Float64>(n, n)
    for (i in 0..n){
        for (j in 0..n) {
            res[i,j] = Float64(H[i,j]) / sqrt(Float64(d[i]))
        }
    }
    return res
}


@Test
public class TestSpecialMatrix {
    @TestCase
    func testDft2(): Unit {
        let expected_dft = matrix<Complex64>([
            [Complex64(1.0, 0.0), Complex64(1.0, 0.0), Complex64(1.0, 0.0), Complex64(1.0, 0.0), Complex64(1.0, 0.0)],
            [Complex64(1.0, 0.0), Complex64(0.309017, -0.951057), Complex64(-0.809017, -0.587785), Complex64(-0.809017, 0.587785), Complex64(0.309017, 0.951057)],
            [Complex64(1.0, 0.0), Complex64(-0.809017, -0.587785), Complex64(0.309017, 0.951057), Complex64(0.309017, -0.951057), Complex64(-0.809017, 0.587785)],
            [Complex64(1.0, 0.0), Complex64(-0.809017, 0.587785), Complex64(0.309017, -0.951057), Complex64(0.309017, 0.951057), Complex64(-0.809017, -0.587785)],
            [Complex64(1.0, 0.0), Complex64(0.309017, 0.951057), Complex64(-0.809017, 0.587785), Complex64(-0.809017, -0.587785), Complex64(0.309017, -0.951057)]
        ])
        @Assert(approxEqualC(dft(5), expected_dft, atol:1e-6))
    }

    @TestCase
    func testBlockDiag2(): Unit {
        var A = matrix<Int64>(
            [[1, 0],
             [0, 1]]
        )
        var B = matrix<Int64>(
            [[3, 4, 5],
             [6, 7, 8]]
        )
        var C = matrix<Int64>(
            [[7]]
        )
        var arr = [A, B, C]

        let expected_res = matrix<Int64>(
            [[1, 0, 0, 0, 0, 0],
             [0, 1, 0, 0, 0, 0],
             [0, 0, 3, 4, 5, 0],
             [0, 0, 6, 7, 8, 0],
             [0, 0, 0, 0, 0, 7]
        ])
        
        @Assert(blockDiag(arr) == expected_res)
    }

    @TestCase
    func testVander(): Unit {
        let x = vector<Float64>([1.0, 2.0, 3.0, 5.0])
        let expected_res = matrix<Float64>([
            [1.0, 1.0, 1.0],
            [4.0, 2.0, 1.0],
            [9.0, 3.0, 1.0],
            [25.0, 5.0, 1.0]
        ])
        @Assert(approxEqual(vander(x, 3), expected_res))
    }

    @TestCase
    func testVanderIncreasing(): Unit {
        let x = vector<Float64>([1.0, 2.0, 3.0, 5.0])
        let expected_res = matrix<Float64>([
            [1.0, 1.0, 1.0],
            [1.0, 2.0, 4.0],
            [1.0, 3.0, 9.0],
            [1.0, 5.0, 25.0]
        ])
        @Assert(approxEqual(vander(x, 3, increasing:true), expected_res))
    }

    @TestCase
    func testCirculant(): Unit {
        let x = vector<Float64>([1.0, 2.0, 3.0])
        let expected_res = matrix<Float64>([
            [1.0, 2.0, 3.0],
            [2.0, 3.0, 1.0],
            [3.0, 1.0, 2.0]
        ])
        @Assert(approxEqual(circulant(x), expected_res))
    }

    @TestCase
    func testCompanion(): Unit {
        let x = vector<Float64>([1.0, -10.0, 31.0, -30.0])
        let expected_res = matrix<Float64>([
            [10.0, -31.0, 30.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0]
        ])
        @Assert(approxEqual(companion(x), expected_res))
    }

    @TestCase
    func testCompanion2(): Unit {
        let x = vector<Float64>([2.0, -10.0, 31.0, -30.0])
        let expected_res = matrix<Float64>([
            [5.0, -15.5, 15.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0]
        ])
        @Assert(approxEqual(companion(x), expected_res))
    }

    @TestCase
    func testFiedler(): Unit {
        let x = vector<Int64>([1, 4, 12, 45, 77])
        let expected_res = matrix<Int64>([
            [0, 3, 11, 44, 76],
            [3, 0, 8, 41, 73],
            [11, 8, 0, 33, 65],
            [44, 41, 33, 0, 32],
            [76, 73, 65, 32, 0]
        ])
        @Assert(fiedler(x) == expected_res)
    }

    @TestCase
    func testHankel(): Unit {
        let x = vector<Int64>([1, 17, 99])
        let expected_res = matrix<Int64>([
            [1, 17, 99],
            [17, 99, 0],
            [99, 0, 0]
        ])
        @Assert(hankel(x) == expected_res)
    }

    @TestCase
    func testHankel2(): Unit {
        let x = vector<Int64>([1, 2, 3, 4])
        let y = vector<Int64>([4, 7, 7, 8, 9])
        let expected_res = matrix<Int64>([
            [1, 2, 3, 4, 7],
            [2, 3, 4, 7, 7],
            [3, 4, 7, 7, 8],
            [4, 7, 7, 8, 9]
        ])
        @Assert(hankel(x, y) == expected_res)
    }

    @TestCase
    func testHilbert(): Unit {
        let expected_res = matrix<Float64>([
            [1.000000, 0.500000, 0.333333],
            [0.500000, 0.333333, 0.250000],
            [0.333333, 0.250000, 0.200000]
        ])
        @Assert(approxEqual(hilbert<Float64>(3), expected_res, atol:1e-6))
    }

    @TestCase
    func testLeslie(): Unit {
        let f = vector<Float64>([0.1, 2.0, 1.0, 0.1])
        let s = vector<Float64>([0.2, 0.8, 0.7])
        let expected_res = matrix<Float64>([
            [0.1, 2.0, 1.0, 0.1],
            [0.2, 0.0, 0.0, 0.0],
            [0.0, 0.8, 0.0, 0.0],
            [0.0, 0.0, 0.7, 0.0]
        ])
        @Assert(approxEqual(leslie(f, s), expected_res))
    }

    @TestCase
    func testPascal(): Unit {
        let expected_res = matrix<Int64>([
            [1, 1, 1, 1],
            [1, 2, 3, 4],
            [1, 3, 6, 10],
            [1, 4, 10, 20]
        ])
        @Assert(pascal(4) == expected_res)
    }

    @TestCase
    func testPascalLower(): Unit {
        let expected_res = matrix<Int64>([
            [1, 0, 0, 0],
            [1, 1, 0, 0],
            [1, 2, 1, 0],
            [1, 3, 3, 1]
        ])
        @Assert(pascal(4, kind:"lower") == expected_res)
    }

    @TestCase
    func testPascalUpper(): Unit {
        let expected_res = matrix<Int64>([
            [1, 1, 1, 1],
            [0, 1, 2, 3],
            [0, 0, 1, 3],
            [0, 0, 0, 1]
        ])
        @Assert(pascal(4, kind:"upper") == expected_res)
    }

    // @TestCase
    // func testHelmert(): Unit {
    //     var res = helmert(5)
    //     var expected_res = matrix<Float64>([
    //         [0.447214, 0.447214, 0.447214, 0.447214, 0.447214],
    //         [0.707107, -0.707107, 0.000000, 0.000000, 0.000000],
    //         [0.408248, 0.408248, -0.816497, 0.000000, 0.000000],
    //         [0.288675, 0.288675, 0.288675, -0.866025, 0.000000],
    //         [0.223607, 0.223607, 0.223607, 0.223607, -0.894427]
    //     ])
    //     @Assert(approxEqual(res, expected_res, atol:1e-6))
    // }

    @TestCase
    func testHadamard(): Unit {
        var res = hadamard(4)
        let expected_res = matrix<Int64>([
            [1, 1, 1, 1],
            [1, -1, 1, -1],
            [1, 1, -1, -1],
            [1, -1, -1, 1]
        ])
        @Assert(res == expected_res)
    }

    @TestCase
    func testConvolutionMatrix(): Unit {
        let a = vector<Int64>([-1, 4, -2])
        var res = convolutionMatrix(a, 5)
        let expected_res = matrix<Int64>([
            [-1, 0, 0, 0, 0],
            [ 4, -1, 0, 0, 0],
            [-2, 4, -1, 0, 0],
            [ 0, -2, 4, -1, 0],
            [ 0, 0, -2, 4, -1],
            [ 0, 0, 0, -2, 4],
            [ 0, 0, 0, 0, -2]
        ])
        @Assert(res == expected_res)
    }

    @TestCase
    func testConvolutionMatrixValid(): Unit {
        let a = vector<Int64>([-1, 4, -2])
        var res = convolutionMatrix(a, 5, mode:"valid")
        let expected_res = matrix<Int64>([
            [-2, 4, -1, 0, 0],
            [ 0, -2, 4, -1, 0],
            [ 0, 0, -2, 4, -1]
        ])
        @Assert(res == expected_res)
    }

    @TestCase
    func testConvolutionMatrixSame(): Unit {
        let a = vector<Int64>([-1, 4, -2])
        var res = convolutionMatrix(a, 5, mode:"same")
        let expected_res = matrix<Int64>([
            [4, -1, 0, 0, 0],
            [-2, 4, -1, 0, 0],
            [ 0, -2, 4, -1, 0],
            [0, 0, -2, 4, -1],
            [ 0, 0, 0, -2, 4]
        ])
        @Assert(res == expected_res)
    }

    @TestCase
    func testToeplitz(): Unit {
        let a = vector<Int64>([1, 2, 3])
        let b = vector<Int64>([1, 4, 5, 6])
        var res = toeplitz(a, b)
        let expected_res = matrix<Int64>([
            [1, 4, 5, 6],
            [2, 1, 4, 5],
            [3, 2, 1, 4]
        ])
        @Assert(res == expected_res)
    }
}