;;; Lisplab, level3-linalg-generic.lisp
;;; Non-spcialized matrix methods.

;;; Copyright (C) 2009 Joern Inge Vestgaarden
;;;
;;; This program is free software; you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published by
;;; the Free Software Foundation; either version 2 of the License, or
;;; (at your option) any later version.
;;;
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License along
;;; with this program; if not, write to the Free Software Foundation, Inc.,
;;; 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.


;;; TODO clean up. Move read and write out

(in-package :lisplab)

(defmethod mtr ((matrix matrix-base))
  (let ((ans 0))
    (dotimes (i (rows matrix))
      (setf ans (.+ ans (mref matrix i i))))
    ans))

(defmethod mtp ((a matrix-base))
  (let ((b (mcreate a 0 (list (cols a) (rows a)))))
    (dotimes (i (rows b))
      (dotimes (j (cols b))
	(setf (mref b i j) (mref a j i))))
    b))

(defmethod mct ((a matrix-base))
  (.conj (mtp a)))

#+old (defmethod m* ((a matrix-base) (b matrix-base))
  (let ((c (mcreate a 0 (list (rows a) (cols b)))))
    (dotimes (i (rows c))
      (dotimes (j (cols c))
	(dotimes (k (cols a))
	  (setf (mref c i j) (.+ (mref c i j)
				 (.* (mref a i k) (mref b k j)))))))
    c))

(defmethod m* ((a matrix-base) (b matrix-base))
  (let ((c (mcreate a 0 (list (rows a) (cols b))))
	(a (mtp a)))
    (dotimes (i (rows c))
      (dotimes (j (cols c))
	(setf (mref c i j) 
	      (col-col-mul-sum A i b j))))
    c))

(defmethod m/ ((a matrix-base) (b matrix-base))
  (m* a (minv b)))

(defmethod minv ((a matrix-base))
  (minv! (copy a)))

(defmethod minv! ((a matrix-base))
  "Matrix inversion based on LU-factorization."
  (let ((LU (copy A)))
    (destructuring-bind (LU p det) 
	(LU-factor! LU  (make-permutation-vector (rows A))) 
      (mfill A 0) ; Use A for the results
      (dotimes (i (rows A))
	(let ((col (view-col A (vref p i))))
	  (setf (vref col i) 1)
	  (LU-solve! LU col))))
    A))

(defmethod LU-factor! ((A matrix-base) p)
  ;; Translation from GSL. 
  ;; Destructive LU factorization. The outout is PA=LU,
  ;; stored in one matrix, where the diagonal elements belong
  ;; to U and L has implicite ones at diagonal.

  ;; TODO: handle permutations better!

  ;; TODO: Change unatural i and j indexing.
  (let ((N (rows A))
	(sign 1))
    (dotimes (i N)
      (setf (vref p i) i))
    (dotimes (j (1- N))
      (let ((i-pivot j)
	    (max (.abs (mref A j j))))
	(loop for i from (1+ j) below N do
	     (let ((Aij (.abs (mref  A i j))))
	       (when (.> Aij max)
		 (setf max Aij
		       i-pivot i))))
	(unless (= i-pivot j)
	  (dotimes (i N)
	    (psetf (mref A j i) (mref A i-pivot i) 
		   (mref A i-pivot i) (mref A j i)))
	  (psetf (vref p j) (vref p i-pivot) 
		 (vref p i-pivot) (vref p j))
	  (setf sign (- sign))))
      (unless (zerop (mref A j j))
	(loop for i from (1+ j) below N do
		(let ((Aij (./ (mref A i j) (mref A j j))))
		  (setf (mref A i j) Aij)
		  (loop for k from (1+ j) below N do		       
		       (setf (mref A i k) 
			     (.- (mref A i k) 
				 (.* Aij (mref A j k)))))))))
    (list A p sign)))
	  
(defmethod LU-factor ((A matrix-base))
  (destructuring-bind (A p sign)
      (LU-factor! (copy A)
		  (make-permutation-vector (rows A)))
    (let ((L A)
	  (U (copy A))
	  (Pmat (mcreate A 0)))
      (w/mat L (x i j) (cond ((> i j) x) ((= i j) 1) (t 0)))
      (w/mat U (x i j) (cond ((<= i j) x) (t 0)))
      (dotimes (i (rows P))
	(setf (mref Pmat (vref p i) i) 1))
      (list L U Pmat))))

(defun L-solve! (L x)
  ;; Solves Lx=b
  (loop for i from 1 below (size x) do
       (let ((sum (vref x i)))
	 (loop for j from 0 below i do
	      (setf sum (.- sum (.* (mref L i j) (vref x j)))))
	 (setf (vref x i) sum))) 
  x)

(defun U-solve! (U x)
  ;; Solves Ux=b
  (let* ((N (size x))
	 (N-1 (1- N)))
    (setf (vref x N-1) (./ (vref x N-1) 
			   (mref U N-1 N-1)))
    (loop for i from (- N-1 1) downto 0 do
	 (let ((sum (vref x i)))
	   (loop for j from (1+ i) below N do
		(setf sum (.- sum (.* (mref U i j) (vref x j)))))
	   (setf (vref x i) (./ sum 
				(mref U i i)))))
      x))

(defmethod LU-solve! ((LU matrix-base) (x matrix-base))
  (L-solve! LU x)
  (U-solve! LU x)
  x)

(defmethod lin-solve ((A matrix-base) (b matrix-base))
  (destructuring-bind (LU pvec det) 
      (LU-factor! (copy A) (make-permutation-vector (rows A)))    
    (let ((b2 (copy b)))
    (dotimes (i (size b))
      (setf (vref b2 i) (vref b (vref pvec i)))) 
    (LU-solve! LU b2))))
  
(defmethod mdet ((A matrix-base))
  (destructuring-bind (LU pvec det) 
      (LU-factor! (copy A) (make-permutation-vector (rows A)))    
    (dotimes (i (rows A))
      (setf det (.* det (mref LU i i))))
    det))

    

;;; Alternative code
      

#+nil (defmethod minv! ((a matrix-base))
  ;; Flawed. Does not work on when pivoting is needed
  "Brute force O(n^3) implementation of matrix inverse.
Think I'll keep this for the general case since it works also 
when the elements cannot be ordered, unlike the LU-based version"
  (let* ((size (rows a))
	 (temp 0))
    (dotimes (i size a)
      (setf temp (mref a i i))
      (dotimes (j size)
	(setf (mref a i j) (if (= i j)
			       (./ (mref a i j))
			       (./ (mref a i j) temp))))
      (dotimes (j size)
	(unless (= i j)
	  (setf temp (mref a j i)
		(mref a j i) 0)
	  (dotimes (k size)
	    (setf (mref a j k)
		  (.- (mref a j k) (.* temp (mref a i k))))))))))