from sequtils import newSeqWith, map
from utils import `...`, `-->`
from vector import Vector, newVector, createVector
from options import Option, none
import future
from random import randomize, random

type
  Matrix*[T] = object
    rows, columns : int 
    data : seq[T]
  SMatrix[W, H: static[int], T] =
    array[1..W, array[1..H, T]]
  Pivot* = ref object
    data* : Vector[int]
    permutations* : int
  SingularMatrixError* = object of ValueError
  SizeError* = object of ValueError
  MatrixRef[T] = ref Matrix[T]

proc `[]`(p : Pivot, index : int) : int = p.data[index]
proc `[]=`(p : var Pivot, index : int, value : int) = p.data[index] = value
proc len(p : Pivot) : int = p.data.len

proc newPivot[T](size : int) : Pivot =
    result = new (Pivot) 
    result = Pivot(data: newVector[int](size), permutations:0)
    for i in 0...size:
        result[i] = i

type AbtractMAtrix = Matrix or SMAtrix

iterator iter_walter[T](m: Matrix[T]) : (int,int, T) {.closure.} =
  for i in 0...m.rows:
      for j in 0...m.columns:
        yield (i, j, m[i,j])

iterator items*[T](m: Matrix[T]): (int,int, T) =
  for i in 0...m.rows:
      for j in 0...m.columns:
        yield (i, j, m[i,j])

proc size*[T](m : Matrix[T]) : (int, int) = (m.rows, m.columns)

proc newMatrix*[T](rows, columns : int, init : T = 0) : Matrix[T] = 
    result = Matrix[T](rows: rows, columns:columns, data: newSeq[T](rows * columns))
    for i,j,_ in items(result):
        result[i,j] = init

proc newMatrix*[T](rows, columns : int, values : openarray[T]) : Matrix[T] = 
    result = Matrix[T](rows: rows, columns:columns, data: newSeq[T](rows * columns))
    for i,j,_ in items(result):
        result[i,j] = values[i * columns + j] 

proc identity*[T](sz : int) : Matrix[T] =
    result = newMatrix[T](sz,sz,0)
    for i in 0...sz:
        result[i,i] = 1


proc `[]`*[T](m : Matrix[T], r,c :int) : T =
    m.data[r*m.columns + c] 

proc `[]`*[T](m : var Matrix[T], r,c :int) : var T =
    m.data[r*m.columns + c] 

proc `[]=`*[T](m : var Matrix[T], r,c :int, newValue : T) =
    m.data[r*m.columns + c] = newValue

# proc `[]=`*[T](m : var Matrix[T], r,c :int, newValue : T) =
#     m.data[r*m.columns + c] = newValue

proc `$`*[T](m : Matrix[T]) : string = 
    result = "["
    for i,j,v in items(m):
        if j == 0:
            if i > 0: result &= " " 
            result &= "["
        result &= $v
        if j == (m.columns - 1): 
            result &= "]"
            if i != (m.rows - 1):
                result &= ",\n"
            else:
                result &= "]\n"
        else: result &= ", "

proc `*`*[T](m1 : Matrix[T], m2 : Matrix[T]) : Matrix[T] =
    result = newMatrix[T](m1.rows, m2.columns, 0)
    for i in 0...result.rows:
        for j in 0...result.columns:
            for k in 0...m1.columns:
                result[i, j] = result[i, j] + m1[i, k] * m2[k, j]

proc `*`*[T](m1 : Matrix[T], v2 : Vector[T]) : Vector[T] =
    result = newVector[T](m1.rows, 0)
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            result[i] = result[i] + m1[i, j] * v2[j]

proc `+`*[T](m1 : Matrix[T], m2 : Matrix[T]) : Matrix[T] =
    result = newMatrix[T](m1.rows, m1.columns)
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            result[i,j] = m1[i,j] + m2[i,j]

proc `+`*[T](m1 : Matrix[T], v : T) : Matrix[T] =
    result = newMatrix[T](m1.rows, m1.columns)
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            result[i,j] = m1[i,j] + v


proc `-`*[T](m1 : Matrix[T], m2 : Matrix[T]) : Matrix[T] =
    result = newMatrix[T](m1.rows, m2.columns)
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            result[i,j] = m1[i,j] - m2[i,j]

proc `-`*[T](m1 : Matrix[T], v : T) : Matrix[T] =
    result = newMatrix[T](m1.rows, m1.columns)
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            result[i,j] = m1[i,j] - v

proc `+=`*[T](m1 : var Matrix[T], m2 : Matrix[T]) =
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            m1[i,j] += m2[i,j]

proc `+=`*[T](m1 : var Matrix[T], v : T) : Matrix[T] =
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            m1[i,j] += v

proc `-=`*[T](m1 : var Matrix[T], v : T) : Matrix[T] =
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            m1[i,j] -= v

proc `-=`*[T](m1 : var Matrix[T], m2 : Matrix[T]) =
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            m1[i,j] -= m2[i,j]

proc `-`*[T](m : Matrix[T]) : Matrix[T] =
    result = newMatrix[T](m.rows, m.columns)
    for i,j,v in m:
        result[i,j] = -v

proc `==`*[T](m1 : Matrix[T], m2 : Matrix[T]) : bool =
    if m1.size() != m2.size():
        return false
    for i in 0...m1.rows:
        for j in 0...m1.columns:
            if m1[i,j] != m2[i,j]:
                return false
    return true

proc clone*[T](m : Matrix[T]) : Matrix[T] = newMatrix[T](m.rows, m.columns, m.data)

proc transpose*[T](m : Matrix[T]) : Matrix[T] =
    result = newMatrix[T](m.rows, m.columns)
    for i, j, v in m:
        result[j, i] = v

proc det*[T](m : Matrix[T]) : T =
    var clone = m.clone()
    clone.gauss_jordan_low()
    result = 1
    for i in 0...clone.rows:
        result *= clone[i, i]

proc swap_rows[T](m : var Matrix[T], id1 : int, id2 : int, pivot : Pivot=nil, other : MatrixRef[T]=nil) =
    for i in 0...m.columns:
        let tmp = m[id1, i]
        m[id1, i] = m[id2, i]
        m[id2, i] = tmp
        if other != nil:
            other[].swap_rows(id1, id2)
    if pivot != nil:
        var pv = pivot
        let tmp = pv[id1]
        pv[id1] = pv[id2]
        pv[id2] = tmp
        pv.permutations =  pv.permutations + 1

proc add_row[T](m : var Matrix[T], sourceIndex : int, destIndex : int, factor : T, other : MatrixRef[T]=nil) =
    for i in 0...m.columns:
        m[destIndex, i] = m[destIndex, i] + m[sourceIndex, i] * factor
    if other != nil:
        other[].add_row(sourceIndex, destIndex, factor)

proc gauss_jordan_low*[T](m : var Matrix[T], other : MatrixRef[T]=nil) =
    var pivot = newPivot[T](m.rows)
    for i in 0...m.rows:
        if m[i, i] == 0:
            for j in (i + 1)...m.columns:
                if m[j, i] != 0:
                    m.swap_rows(i, j, pivot, other)
                    break
        for j in (i + 1)...m.rows:
            if m[i, i] != 0:
                let factor = -m[j, i] / m[i, i]
                m.add_row(i, j, factor, other)

proc gauss_jordan_high*[T](m : var Matrix[T], other : MatrixRef[T]=nil) = 
    var pivot = newPivot[T](m.rows)
    for i in m.rows-->0:
        if m[i, i] == 0:
            for j in i-->0:
                if m[j, i] != 0:
                    m.swap_rows(i, j, pivot, other)
                    break
        for j in i-->0:
            if m[i, i] != 0:
                let factor = -m[j, i] / m[i, i]
                m.add_row(i, j, factor, other)

proc invert*[T](m : Matrix[T]) : Matrix[T] =
    var tmp = m.clone()
    var res = new(MatrixRef[T])
    res[] = identity[T](tmp.rows)
    tmp.gauss_jordan_low(res)
    tmp.gauss_jordan_high(res)
    for i in 0...res.rows:
        let f = tmp[i, i]
        for j in 0...res.columns:
            res[][i, j] /= f
    res[]

proc triu*[T](m : Matrix[T], diag_replace=nil) : Matrix[T] =
    result = Matrix(m.rows, m.columns, 0)
    for i in range(m.rows):
        for j in range(i, m.columns):
            if diag_replace and i == j:
                result[i, j] = diag_replace
            else:
                result[i, j] = m[i, j]

proc tril*[T](m : Matrix[T], diag_replace=nil) : Matrix[T] =
    result = Matrix(m.rows, m.columns, 0)
    for i in range(m.rows):
        for j in range(i + 1):
            if diag_replace and i == j:
                result[i, j] = diag_replace
            else:
                result[i, j] = m[i, j]

proc lu_row[T](m : var Matrix[T], i : int) =
    if m[i, i] == 0:
        raise newException(SingularMatrixError, "Matrix is singular")
    for j in i...m.columns:
        for k in 0...i:
            m[i, j] = m[i, j] - m[i, k] * m[k, j]
    for j in (i + 1)...m.columns:
        for k in 0...i:
            m[j, i] = m[j, i] - m[j, k] * m[k, i]
        m[j, i] = m[j, i] / m[i, i]

proc lu_pivot[T](m : var Matrix[T], i : int, pivot : Pivot) =
    var max = abs(m[i, i])
    var max_index = i
    for j in (i + 1)...m.rows:
        if abs(m[j, i]) > max:
            max = abs(m[i, j])
            max_index = j
    if max_index != i:
        m.swap_rows(i, max_index, pivot)

proc lu*[T](m : var Matrix[T], pivoting=true) : Pivot =
    var pivot = newPivot[T](m.rows)
    if pivoting:
        for i in 0...m.rows:
            m.lu_pivot(i, pivot)
            m.lu_row(i)
    else:
        for i in 0...m.rows:
            m.lu_row(i)
    return pivot

proc lu_solve*[T](m : Matrix[T], b : Vector[T], p : Pivot = nil) : Vector[T] = 
    var pivot = p
    if pivot == nil:
        pivot = new(Pivot)
        pivot = newPivot[T](m.rows)
    var x = newVector[T](m.rows)
    for i in 0...m.rows:
        x[i] = b[pivot[i]]
        for k in 0...i:
            x[i] = x[i] - m[i, k] * x[k]
    for i in m.rows-->0:
        for k in (i + 1)...m.rows:
            x[i] = x[i] - m[i, k] * x[k]
        x[i] = x[i] / m[i, i]
    return x

proc lu_invert*[T](m : Matrix[T], pivot : Pivot=nil) : Matrix[T] =
    if not pivot:
        pivot = newPivot[T](m.rows)
    result = newMatrix[T](m.rows, m.columns)
    for i in 0...m.rows:
        for j in 0...m.rows:
            if pivot[j] == i:
                result[j, i] = 1
            else:
                result[j, i] = 0
            for k in range(j):
                result[j, i] -= m[j, k] * result[k, i]
        for j in m.rows-->0:
            for k in range(j + 1, m.rows):
                result[j, i] -= m[j, k] * result[k, i]
            result[j, i] = result[j, i] / m[j, j]

proc lu_det*[T](m : Matrix[T], pivot : Pivot=nil) : T =
    if not pivot:
        pivot = newPivot[T](m.rows)
    result = 1
    for i in 0...m.rows:
        result *= m[i, i]
    if pivot.permutations mod 2 != 0:
        result *= -1

proc from_pivot[T](pivot : Pivot): Matrix[T] =
    result = Matrix(len(pivot), len(pivot))
    for i in 0...pivot.len:
        result[pivot[i], i] = 1
        # result[j, i] = 1