;;; Lisplab, level3-linalg-dge.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.

(in-package :lisplab)

#+todo (defmethod mtr (matrix)
  (let ((ans 0))
    (dotimes (i (rows matrix))
      (setf ans (.+ ans (mref matrix i i))))
    ans))

(defmethod mtp ((a matrix-base-dge))
  (let* ((N (rows a))
	 (M (cols a))
	 (c (mcreate a 0 (list M N)))
	 (a2 (vector-store a))
	 (c2 (vector-store c)))
    (declare (type-blas-store a2 c2)
             (type-blas-idx M N))
    (macrolet ((refa (i j) `(ref-blas-real-store A2 ,i ,j N))
	       (refc (i j) `(ref-blas-real-store C2 ,i ,j M)))
      (dotimes (i M)
	(dotimes (j N)
	  (setf (refc i j) (refa j i)))))
    c))

(defmethod mct ((a matrix-base-dge))
  (mtp a))

(defmethod m* ((a matrix-base-dge) (b matrix-base-dge))
  (let* ((N (rows a))
	 (M (cols b))
	 (S (rows b))
	 (c (mcreate a 0 (list N M)))
	 (a2 (vector-store (mtp a))) ; optimization. consideres the transpose
	 (b2 (vector-store b))
	 (c2 (vector-store c)))
    (declare (type-blas-store a2 b2 c2)
             (type-blas-idx N M S))
    (macrolet ((refa (i j) `(ref-blas-real-store A2 ,j ,i N))
	       (refb (i j) `(ref-blas-real-store B2 ,i ,j S))
	       (refc (i j) `(ref-blas-real-store C2 ,i ,j N)))
      (dotimes (i N)
	(dotimes (j M)
	  (let ((cij 0d0))
	    (declare (double-float cij))
	    (dotimes (k S)
	      (incf cij (* (refa i k ) (refb k j))))
	    (setf (refc i j) cij)))) 
      c)))

(defmethod LU-factor! ((A matrix-base-dge) 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.
  ;; Assumes the permutation vector to be initilized
  (let ((N (rows A))
	(sign 1)
	(A2 (vector-store A)))
    (declare (type-blas-idx N) 
	     (fixnum sign)
	     (type-blas-store a2)
	     (type-permutation-vector p))
    (macrolet ((xxx (i j) `(ref-blas-real-store A2 ,i ,j N)))
      (dotimes (j (1- N))
	(let ((i-pivot j)
	      (max (abs (xxx j j))))
	  (loop for i from (1+ j) below N do
	       (let ((Aij (abs (xxx i j))))
		 (when (> Aij max)		   
		   (setf max Aij
			 i-pivot i))))
	(unless (= i-pivot j)
	  (dotimes (i N)
	    (let ((tmp (xxx j i)))
	      (setf (xxx j i) (xxx i-pivot i) 
		    (xxx i-pivot i) tmp)))
	  (let ((tmp (vref p j))) 
	    (setf (vref p j) (vref p i-pivot) 
		  (vref p i-pivot) tmp)
	    (setf sign (- sign)))))
      (unless (zerop (xxx j j))
	(loop for i from (1+ j) below N do
		(let ((Aij (/ (xxx i j) (xxx j j))))
		  (setf (xxx i j) Aij)
		  (loop for k from (1+ j) below N do		       
		       (setf (xxx i k) 
			     (- (xxx i k) 
				 (* Aij (xxx j k)))))))))
    (list A p sign))))
	  
(defun L-solve!-blas-real (L x col)
  ;; Solve Lx=b
  (declare (matrix-base-dge L x))
  (let ((L2 (vector-store L))
	(x2 (vector-store x))
	(N (rows x)))
    (declare (type-blas-store L2 x2)
	     (type-blas-idx N col))
    (macrolet ((refL (i j) `(ref-blas-real-store L2 ,i ,j N))
	       (refx (i)   `(ref-blas-real-store x2 ,i col N)))
      (loop for i from 1 below N do
	   (let ((sum (refx i)))
	     (declare (double-float sum))
	     (loop for j from 0 below i do
		  (decf sum (* (refL i j) (refx j))))
	     (setf (refx i) sum)))))
  x)

(defun U-solve!-blas-real (U x col)
  (declare (matrix-base-dge U x))
  (let* ((U2 (vector-store U))
	 (x2 (vector-store x))
	 (N (rows x))
	 (N-1 (1- N))
	 (N-2 (1- N-1)))
    (declare (type-blas-store U2 x2)
	     (type-blas-idx N-1 N-2 N col))
    (macrolet ((refU (i j) `(ref-blas-real-store U2 ,i ,j N))
	       (refx (i)   `(ref-blas-real-store x2 ,i col N)))
      (setf (refx N-1) (/ (refx N-1) (refU N-1 N-1)))
      (loop for i from N-2 downto 0 do
	   (let ((sum (refx i)))
	     (declare (double-float sum))
	     (loop for j from (1+ i) below N do
		  (decf sum (* (refU i j) (refx j))))
	     (setf (refx i) (/ sum (refU i i)))))))
      x)

(defun LU-solve!-blas-real (LU x col)
  (L-solve!-blas-real LU x col)
  (U-solve!-blas-real LU x col)
  x)

(defmethod lin-solve ((A matrix-base-dge) (b matrix-base-dge))
  (destructuring-bind (LU pvec sign) 
      (LU-factor! (copy A) (make-permutation-vector (size b)))
    (let ((b2 (copy b)))
      (dotimes (i (rows A))
	(setf (vref b2 (vref pvec i)) (vref b i)))
      (LU-solve!-blas-real LU b2 0))))

(defun minv!-blas-real (A)
  (let ((N (rows A)))
    (if (= N 1)
	(setf (mref A 0 0) (/ 1 (mref A 0 0)))
	(let ((LU (copy A)))
	  (destructuring-bind (LU p det)
	      (LU-factor! LU (make-permutation-vector N))
	    (declare (ignore det))
	    (mfill A 0)
	    (dotimes (i N)
	      (setf (mref A i (vref p i)) 1)
	      (LU-solve!-blas-real LU A  (vref p i)))))))
  A)

(defmethod minv! ((A matrix-base-dge))
  (minv!-blas-real A))

(defmethod minv ((A matrix-base-dge))
  (minv! (copy A)))