;;; linear-algebra.scm  --  Jens Axel Søgaard

; Note: Vektors are represented as lists of numbers.
;       Matrices are represented as lists of vektors.

; Note: The routines in the Schemathics library
;       represent matrices as vectors of vectors.
;       Such a matrix is termed a smatrix here
;       (for "Schemathics matrix")

; Note: All "non-trivial" operations are handled
;       by converting to and from smatrices.

(module linear-algebra mzscheme
  (provide (all-defined))
  (require mzlib/match
           mzlib/list
           (only scheme/base vector-copy!)
           )

;;; SUPPORT FOR SMATRIX

(require (prefix sm: "matrix.ss"))

(define (matrix->smatrix m)
  (list->vector (map list->vector m)))

(define (smatrix->matrix sm)
  (map vector->list (vector->list sm)))

(define vektor? pair?)

(define matrix?
  (match-lambda 
    [()            #t]
    [((x))         (number? x)]
    [((x ...) ...) #t]
    [_             #f]))


;;; VEKTOR OPERATIONS

(define (inner-product v w)
  (do ([sum 0 (+ sum (* (car v) (car w)))]
       [v   v (cdr v)]
       [w   w (cdr w)])
    [(null? v) sum]))

(define dot inner-product)

(define (norm-squared v)
  (inner-product v v))

(define (norm v)
  (sqrt (norm-squared v)))

(define (scale-vector s v)
  (map (lambda (vi) (* s vi)) v))

(define (vektor+vektor v w)
  (map + v w))

(define (vektor-vektor v w)
  (map - v w))

(define (matrix*matrix v w)
  (map (lambda (vi)
         (map (lambda (wj)
                (inner-product vi wj))
              (transpose w)))
       v))


;;; BASIC MATRIX OPERATIONS

(define (transpose x)
  (cond
    [(matrix? x) (apply map list x)]
    [(vektor? x) (transpose (list x))]
    [else        (error "transpose: matrix or vektor expected; given: " x)]))

;;; CONSTRUCTING MATRICES

(define (matrix-map f m)
  (map (lambda (row)
         (map f row))
       m))

; build-matrix : (N N -> alpha) N -> alpha
;   build an mxn matrix with entries 
;     [[(f 1 1) (f 1 2) ... (f 1 n)] ... [(f m 1) (f m 2) ... (f m n)]]
(define (build-matrix f m n)
  (define (build-row f m n)
    (do ([j   n   (- j 1)]
         [row '() (cons (f m j) row)])
      [(= j 0) row]))
  (do ([i    m   (- i 1)]
       [rows '() (cons (build-row f i n) rows)])
    [(= i 0) rows]))


; (kronecker i j) = 1 if i=j, otherwise 0
(define (kronecker i j)
  (if (= i j) 1 0))

; identity-matrix : N -> matrix 
;   return the identity matix of dimension nxn
(define (identity-matrix n)
  (build-matrix kronecker n n))

; diagonal-matrix : (list alpha) -> (matrix alpha)
;   return a diagonal matrix with diagonal elements from the given list
(define (diagonal-matrix diag-elements)
  (let ([m         (length diag-elements)]
        [diagonal  (list->vector diag-elements)])
    (build-matrix (lambda (i j) (* (kronecker i j)
                                   (vector-ref diagonal (- i 1))))
                  m m)))


;;; GETTING AND SETTING PIECES

; get : matrix N      -> vektor  (row)
; get : matrix N N    -> number  (entry)
; get : matrix 'all N -> vektor  (column)
(define get
  (case-lambda
    [(m i)   (list-ref m (- i 1))]                             ; i'th row
    [(m i j) (if (eq? i 'all)                                  ; j'th column
                 (map (lambda (row) (list-ref row (- j 1))) m)
                 (list-ref (list-ref m (- i 1)) (- j 1)))]))   ; (i,j)th element

; trace : matrix -> R
;  (trace m)   = trace of m = sum of diagonal elements
; trace : (matrix alpha) ((list alpha) -> beta) -> beta
;  (trace m f) = (f list-of-diagonal-elements)
(define trace
  (case-lambda
    [(m)   (trace m (lambda (l) (do ([l l (cdr l)] [sum 0 (+ (car l) sum)]) [(null? l) sum])))]
    [(m f) (let ([n (length (list-ref m 0))])
             (do ([i    n   (- i 1)]
                  [diag '() (cons (get m i i) diag)])
               [(= i 0) (f diag)]))]))

; dimensions : matrix -> (list N N)
;   return a list with the row and col dimensions
(define (dimensions m)
  (list (length m)
        (if (null? m)
            0
            (length (list-ref m 0)))))

;;; MATRIX INVERSION

; inverse : matrix -> matrix
;  return the inverse of a square matrix
(define (inverse m)
  (smatrix->matrix
   (sm:invert-matrix 
    (matrix->smatrix m))))

; lu-smatrix->list-of-lu : smatrix -> (list matrix matrix)
;   convert a smatrix containing a lu-decompistion
;   to a list of two matricees
(define (lu-smatrix->list-of-lu lu)
  (let* ([m (vector-length lu)]
         [n (vector-length (vector-ref lu 0))]
         [l   (build-matrix (lambda (i j)
                              (cond
                                [(> j i) 0]
                                [(= j i) 1]
                                [else    (vector-ref (vector-ref lu (- i 1))
                                                     (- j 1))]))
                            m n)]
         [u   (build-matrix (lambda (i j)
                              (cond
                                [(< j i) 0]
                                [else    (vector-ref (vector-ref lu (- i 1))
                                                     (- j 1))]))
                            m n)])
    (list l u)))


; lu : matrix -> (list matrix matrix)
;   return a list with the matrices in the lu-decompostion of m
(define (lu m)
  (let ([sm  (matrix->smatrix m)])
    (sm:lu-factor/pivoting! sm)
    (lu-smatrix->list-of-lu sm)))

; make-solver : matrix -> (vektor -> vektor)
;  given a matrix M return a procedure f, such
;  that (f v) will return a solution to Mx=v.
(define (make-solver m)
  (let ([solver (sm:make-solver/pivoting (matrix->smatrix m))])
    (lambda (v)
      (vector->list (solver (list->vector v))))))

; linear-solve : matrix vektor ... -> (list vektor)
;  return a list of solutions to the equations Mx=b1, Mx=b2, ...
(define (linear-solve m b . bs)
  (map (make-solver m) (cons b bs)))
  
(define (count-cycles p)
  (let* ([l (vector-length p)]
         [v (make-vector l)])
    (vector-copy! v 0 p)
    (let next-cycle ([n 0] [i 0])
      (cond 
        [(= i l) n]
        [(not (vector-ref v i))
         (next-cycle n (add1 i))]
        [else
         (let follow-cycle ([j i])
           (let ([k (vector-ref v j)])
             (vector-set! v j #f)
             (if (vector-ref v k)
                 (follow-cycle k)
                 (next-cycle (add1 n) (add1 i)))))]))))
  
(define (permutation-sign p)
  (let ([n (vector-length p)])
    (if (even? (- n (count-cycles p))) 1 -1)))

; determinant : matrix -> number
;  return the determinant of m
#;(define (det m)
    (define (product l) (apply * l))
    (let ([u (matrix->smatrix m)])
      (sm:lu-factor! u)
      (trace (smatrix->matrix u) product)))
  
(define (det m)
  (define (product l) (apply * l))
  (let* ([u (matrix->smatrix m)]
         [p (sm:lu-factor/pivoting! u)])
    (* (permutation-sign p)
       (trace (smatrix->matrix u) product))))

; (rank m) = dimension of row space 
;          = dimension of column space
;          = number of pivot columns in the row echelon form of m
(define (rank m)
  (length (sm:row-echelon-form! (matrix->smatrix m))))

(define (no-columns m)
  (second (dimensions m)))

(define (no-rows m)
  (length m))

; nullity is the dimension of 
(define (nullity m)
  (- (no-columns m) (rank m)))


; basis : (list vektor) -> (list vektor)
;   remove vektors from the list such that the
;   resulting list is a basis
(define (basis list-of-vektors)
  (let ([pivot-cols (sm:row-echelon-form! 
                     (matrix->smatrix (transpose list-of-vektors)))]
        [vektors (list->vector list-of-vektors)])
    (map (lambda (i) (vector-ref vektors i))
         pivot-cols)))

; dimension : (list vektor) -> N
;   return the dimension of the space generated 
;   by the given vektors
(define (dimension list-of-vektors)
  (length (basis list-of-vektors)))

; zero-vector? : vector -> boolean
;   is all entries of v zero?
(define (zero-vector? v)
  (let ([len (vector-length v)])
    (let loop ([i 0])
      (cond
        [(= i len)                #t]
        [(zero? (vector-ref v i)) (loop (+ i 1))]
        [else                     #f]))))


; row-space : matrix -> (list vektor)
;  return list of generators for the row space
(define (row-space m)
  (let ([sm (matrix->smatrix m)])
    (sm:row-echelon-form! sm)
    (map vector->list
         (filter (lambda (v) (not (zero-vector? v)))
                 (vector->list sm)))))

(define (column-space m)
  (let ([pivot-cols (sm:row-echelon-form! (matrix->smatrix m))]
        [columns    (list->vector (transpose m))])
    (map (lambda (i) (vector-ref columns i))
         pivot-cols)))


;(define (null-space m)
;  ...)



;;; ORTHOGONALIZATION

;http://documents.wolfram.com/v5/Add-onsLinks/StandardPackages/LinearAlgebra/Orthogonalization.html

; the projection of v on w         
(define (projection v w)
  (scale-vector (/ (inner-product v w)
                   (inner-product w w))
                w))

(define (projection-on-orthonormal-basis v ortho-basis)
  (let ([terms       (map (lambda (b) (projection v b))
                          ortho-basis)]
        [zero-vektor (map (lambda (x) 0) v)])
    (do ([proj  zero-vektor (vektor+vektor proj (car terms))]
         [terms terms (cdr terms)])
      [(null? terms) proj])))

(define (normalize v)
  (scale-vector (/ (norm v)) v))

(define (gram-schmidt list-of-vektors normalize?)
  (let ([basis (basis list-of-vektors)])
    (do ([basis       (cdr basis)        (cdr basis)]
         [ortho-basis (list (car basis)) (let ([v (car basis)])
                                           (cons (vektor-vektor v
                                                                (projection-on-orthonormal-basis 
                                                                 v ortho-basis))
                                                 ortho-basis))])
      [(null? basis) (reverse (if normalize?
                                  (map normalize ortho-basis)
                                  ortho-basis))])))

(define (orthogonal-basis list-of-vektors)
  (gram-schmidt list-of-vektors #f))

(define (orthonormal-basis list-of-vektors)
  (gram-schmidt list-of-vektors #t))

; qr : matrix -> (list matrix matrix)
;  Find Q and R such that M=QR, where Q has
;  has orthonormal column vectors and R is
;  an upper-triangular invertible matrix
(define (qr m)
  (let ([q (orthonormal-basis (transpose m))])
    (list (transpose q) (matrix*matrix (inverse q) m))))


(define (symmetric? m)
  (equal? m (transpose m)))

(define (complex-conjugate m)
  (if (number? m)
      (make-rectangular (real-part m) (- (imag-part m)))
      (matrix-map (lambda (z) (make-rectangular (real-part z) (- (imag-part z))))
                  m)))

(define (hermitian? m)
  (equal? m (complex-conjugate (transpose m))))

(define (commutes? m1 m2)
  (equal? (matrix*matrix m1 m2)
          (matrix*matrix m2 m1)))

(define (normal? m)
  (commutes? m (complex-conjugate (transpose m))))



;;; EIGENVALUES AND EIGENVECTORS

(define (jacobi a n d v nrot)
  (define-syntax vr      (syntax-rules () [(_ v i)     (vector-ref  v i)]))
  (define-syntax vs!     (syntax-rules () [(_ v i x)   (vector-set! v i x)]))
  (define-syntax mr      (syntax-rules () [(_ m i j)   (vr (vr m i) j)]))
  (define-syntax ms!     (syntax-rules () [(_ m i j x) (vs! (vr m i) j x)]))
  (define-syntax define* (syntax-rules ()
                           [(_ () v)           (begin)]
                           [(_ (id ids ...) v) (begin (define id v) (define* (ids ...) v))]))
  (define-syntax :=      (syntax-rules () [(_ id v) (set! id v)]))
  (define-syntax for     (syntax-rules () [(_ id from to body ...) (do ([id from (+ id 1)]) [(> id to)] body ...)]))
  
  (define* (i j iq ip) 0)
  (define* (tresh theta tau t sm s h g c) 0.0)
  (define-syntax rotate  (syntax-rules () [(_ a i j k l) (begin (:= g (mr a i j))
                                                                (:= h (mr a k l))
                                                                (ms! a i j (- g (* s (+ h (* g tau)))))
                                                                (ms! a k l (+ h (* s (- g (* h tau))))))]))
  (define b (make-vector (+ n 1)))
  (define z (make-vector (+ n 1)))
  
  (for ip 1 n
       (for iq 1 n
            (ms! v ip iq 0.0))
       (ms! v ip ip 1.0))
  
  (for ip 1 n
       (let ([t (mr a ip ip)])
         (vs! d ip t)
         (vs! b ip t)
         (vs! z ip 0.0)))
  (:= nrot 0)
  (let/ec return
    (for i 1 50
         (:= sm 0.0)
         (for ip 1 (- n 1)
              (for iq (+ ip 1) n
                   (:= sm (+ sm (abs (mr a ip iq))))))
         (if (= sm 0.0)
             (return a d v))
         (if (< i 4)
             (:= tresh (/ (* 0.2 sm) (* n n)))
             (:= tresh 0.0))
         (for ip 1 (- n 1)
              (for iq (+ ip 1) n
                   (:= g (* 100.0 (abs (mr a ip iq))))
                   (if (and (> i 4) 
                            (= (+ (abs (vr d ip)) g) (abs (vr d ip)))
                            (= (+ (abs (vr d iq)) g) (abs (vr d iq))))
                       (ms! a ip iq 0.0)
                       (when (> (abs (mr a ip iq)) tresh)
                         (:= h (- (vr d iq) (vr d ip)))
                         (if (= (+ (abs h) g) (abs h))
                             (:= t (/ (mr a ip iq) h))
                             (begin
                               (:= theta (/ (* 0.5 h) (mr a ip iq)))
                               (:= t (/ 1.0 (+ (abs theta) (sqrt (+ 1.0 (* theta theta))))))
                               (if (< theta 0.0) (:= t (- t)))))
                         (:= c (/ 1.0 (sqrt (+ 1.0 (* t t)))))
                         (:= s (* t c))
                         (:= tau (/ s (+ 1.0 c)))
                         (:= h (* t (mr a ip iq)))
                         (vs! z ip (- (vr z ip) h))
                         (vs! z iq (+ (vr z iq) h))
                         (vs! d ip (- (vr d ip) h))
                         (vs! d iq (+ (vr d iq) h))
                         (ms! a ip iq 0.0)
                         (for j 1 (- ip 1)
                              (rotate a j ip j iq))
                         (for j (+ ip 1) (- iq 1)
                              (rotate a ip j j iq))
                         (for j (+ iq 1) n
                              (rotate a ip j iq j))
                         (for j 1 n
                              (rotate v j ip j iq))
                         (:= nrot (+ nrot 1)))))
              (for ip 1 n
                   (vs! b ip (+ (vr b ip) (vr z ip)))
                   (vs! d ip (vr b ip))
                   (vs! z ip 0.0)))))
  (error "jacobi: too many iterations"))


  ;(define (eigenvalues m)
;  ...)
;
;(define (eigenvectors m)
;  ...)
;
;(define (eigensystem m)
;  ...)
;
;(define (characteristic-polynomial m)
;  ...)

  
)