package require math::linearalgebra

namespace eval ::math::linearalgebra {
    namespace export joinMatrix cancelrow cancelcol
    namespace export cofactor cofactorMatrix
    namespace export adjointMatrix
    namespace export determinant invert rank
    namespace export mkRandom
    namespace export round
}


# joinMatrix --
#     Join two matrices along the specified dimension
# Arguments:
#     dim        Dimension to join [1,2]
#     mv1,mv2    Matrices to be treated
#
# Result:
#     Matrix result of the join
#
proc ::math::linearalgebra::joinMatrix {dim mv1 mv2} {
    set result $mv1
    switch -exact -- $dim {
        1 {
            foreach r $mv2 {lappend $result $r}
        }
        2 {
            set l [llength $mv1]
            for {set i 0} {$i < $l} {incr i} {
                lset result $i [concat [lindex $mv1 $i] [lindex $mv2 $i]]
            }
        }
        default {
            return -code error "do we work on ${dim}D matrices?"
        }
    }
    return $result
}

# cancelrow --
#     Return a matrix minus the specified row
# Arguments:
#     matrix     Matrix to work on
#     row        to delete
#
# Result:
#     Matrix result of the operation
#
proc ::math::linearalgebra::cancelrow {matrix row} {
    return [concat [lrange $matrix 0 $row-1] [lrange $matrix $row+1 end]]
}

# cancelrow --
#     Return a matrix minus the specified column
# Arguments:
#     matrix     Matrix to work on
#     col        to delete
#
# Result:
#     Matrix result of the operation
#
proc ::math::linearalgebra::cancelcol {matrix col} {
    set result {}
    foreach r $matrix {
        lappend result [concat [lrange $r 0 $col-1] [lrange $r $col+1 end]]
    }
    return $result
}

# cofactor --
#     Compute the cofactor of the specified element
# Arguments:
#     matrix     Matrix to work on
#     row        
#     col        
#
# Result:
#     The cofactor of element at a{row,col}
#
proc ::math::linearalgebra::cofactor {matrix row col} {
    set submatrix [cancelcol [cancelrow $matrix $row] $col]
    return [determinant $submatrix]
}

# cofactorMatrix --
#     Compute the cofactor of the specified matrix
# Arguments:
#     matrix     Matrix to work on
#
# Result:
#     The cofactor of specified matrix
#
proc ::math::linearalgebra::cofactorMatrix {matrix} {
    set result {}
    lassign [shape $matrix] rows cols
    for {set i 0} {$i < $rows} {incr i} {
        set newrow {}
        for {set j 0} {$j < $cols} {incr j} {
            lappend newrow [expr ((($i+$j)%2)==0?1:-1)*[cofactor $matrix $i $j]]
        }
        lappend result $newrow
    }
    return $result
}

# adjointMatrix --
#     The adjoint matrix is the transpose of the cofactor matrix
# Arguments:
#     matrix
#
# Result:
#     The adjoint matrix
#
proc ::math::linearalgebra::adjointMatrix matrix {
    return [transpose [cofactorMatrix $matrix]]
}

# determinant --
#     Compute the cofactor of the specified element
# Arguments:
#     matrix     Matrix to work on
#     row        
#     col        
#
# Result:
#     The cofactor of element at a{row,col}
#
proc ::math::linearalgebra::determinant matrix {
    set shape [shape $matrix]
    if {[lindex $shape 0] != [lindex $shape 1]} {
        return -code error "determinant only defined for a square matrix"
    }

    switch -exact -- [join $shape x] {
        1x1 {
            return [lindex [lindex $matrix 0] 0]
        }
        2x2 {
            lassign [getrow $matrix 0] a b
            lassign [getrow $matrix 1] c d
            return [expr {($a*$d)-($c*$b)}]
        }
        default {
            set det 0
            set sign 0
            set row_no 0
            foreach row $matrix {
                set sign [expr {($sign==1)?(-1):(1)}]
                set det [expr {
                    $det+$sign*[lindex $row 0]*[cofactor $matrix $row_no 0]
                }]
                incr row_no
            }
            return $det
        }
    }
}

# invert --
#     Perform the matrix inversion using the adjoint method
#     Note: probably the Gauss-Jordan elimination over an
#           augmented matrix is *much* faster.
# Arguments:
#     matrix     Matrix to work on
#
# Result:
#     The inverted matrix
#
proc ::math::linearalgebra::invert matrix {
    set shape [shape $matrix]

    set det [determinant $matrix]
    if {$det == 0} {
        return -code error "cannot invert a singular matrix"
    }

    switch -exact -- [join $shape x] {
        2x2 {
           set temp [mkMatrix 2 2]
           setelem temp 0 0 [getelem $matrix 1 1]
           setelem temp 1 1 [getelem $matrix 0 0]
           setelem temp 0 1 [expr {-1*[getelem $matrix 0 1]}]
           setelem temp 1 0 [expr {-1*[getelem $matrix 1 0]}]
        }
        default {
           set temp [adjointMatrix $matrix]
        }
    }

    return [scale_mat [expr 1./$det] $temp]
}

proc ::math::linearalgebra::mkRandom { rows {cols {}} args} {
    if {$cols eq {}} {set cols $rows}
    set min 0
    set max 1
    if {[llength $args] == 1} {
        set max $args
    } elseif {[llength $args] == 2} {
        lassign $args min max
    }
    set result {}
    for {set i 0} {$i < $rows} {incr i} {
       set row {}
       for {set j 0} {$j < $cols} {incr j} {
           lappend row [expr {$min+$max*rand()}]
       }
       lappend result $row
    }
    return $result
}

proc ::math::linearalgebra::round matrix {
    set result {}
    foreach row $matrix {
        set newrow {}
        foreach el $row {
            lappend newrow [expr {round($el)}]
        }
        lappend result $newrow
    }
    return $result
}

proc ::math::combinations {n k} {
    set l {}
    for {set i 0} {$i < $n} {incr i} {lappend l $i}
    return [lcombinations $l $k]
}

proc ::math::lcombinations {list size} {
    if {$size == 0} {return {{}}}
    set retval {}
    for {set i 0} {($i + $size) <= [llength $list]} {incr i} {
        set firstElement [lindex $list $i]
        set remainingElements [lrange $list [expr {$i+1}] end]
        foreach subset [lcombinations $remainingElements [expr {$size-1}]] {
            lappend retval [linsert $subset 0 $firstElement]
        }
    }
    return $retval
}

proc ::math::linearalgebra::rank {matrix {checkR {}}} {
    lassign [shape $matrix] rows cols
    if {$checkR eq {}} {set checkR [expr ($cols>$rows)?$rows:$cols]}
    if {$checkR == 1} {return 1}
    set col_idx [combinations $cols $checkR]
    set row_idx [combinations $rows $checkR]
    foreach coll $col_idx {
        foreach rowl $row_idx {
            set m {}
            for {set i 0} {$i < [llength $rowl]} {incr i} {
                set row {}
                for {set j 0} {$j < [llength $coll]} {incr j} {
                    lappend row [getelem $matrix [lindex $rowl $i] [lindex $coll $j]]
                }
                lappend m $row
            }
            set det [determinant $m]
            if {$det != 0} {return $checkR}
        }
    }
    return [rank $matrix [incr checkR -1]]
}