;;; Lisplab, level2-matrix-dge.lisp
;;; Optimizations for real matrices.

;;; 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)

(defmethod mfill ((a vector-d) value)
  (let ((x (coerce value 'double-float))
	(store (vector-store a)))
    (declare (type type-blas-store store))
    (fill store x)))

(defmethod copy ((a vector-d))
  (let ((store (vector-store a)))
    (declare (type type-blas-store store))
    (make-instance (class-of a)
		   :store (copy-seq store)
		   :dim (dim a))))

(defmethod mreverse ((a vector-d))
  (let ((store (vector-store a)))
    (declare (type type-blas-store store))
    (make-instance (class-name (class-of a))
		   :store (reverse store)
		   :dim (dim a))))

(defmethod copy-contents ((a vector-d) (b vector-d) &optional (converter nil))
  (let ((store-a (vector-store a))
	(store-b (vector-store b)))
    (declare (type type-blas-store store-a store-b))
    (if converter 
	(map-into store-b converter store-a)
	(copy-matrix-stores store-a store-b)))    
    b)

(defmethod mmap-into ((out vector-d) f (a vector-d) &rest args)
  (apply #'map-into 
	 (vector-store out)
	 (lambda (&rest args)
	   (coerce (apply f args) 'double-float))
	 (vector-store a) (mapcar #'vector-store args))
    out)

(defmethod mmap ((type (eql nil)) f (a vector-d) &rest args)
  "No output. Called for side effects."
  (if args
      (apply #'map
	     nil
	     f 
	     (vector-store a) (mapcar #'vector-store args))
      (let ((store (vector-store a)))
	(declare (type type-blas-store store))
	(map nil f store)))
  nil)

(defmethod msum ((m vector-d))
  (let ((sum 0d0))
    (with-elements-df-1 (vector-store m) x
      (incf sum x))
    sum))

(defmethod mmax ((m vector-d))
  (let* ((store (vector-store m))
	 (max (aref store 0))
	 (idx 0))
    (declare (type type-blas-store store)
	     (type double-float max)
	     (type type-blas-idx idx))
    (dotimes (i (length store))
      (when (> (aref store i) max)
	(setf max (aref store i)
	      idx i)))
    (values max idx)))

(defmethod mmin ((m vector-d))
  (let* ((store (vector-store m))
	 (min (aref store 0))
	 (idx 0))
    (declare (type type-blas-store store)
	     (type double-float min)
	     (type type-blas-idx idx))
    (dotimes (i (length store))
      (when (< (aref store i) min)
	(setf min (aref store i)
	      idx i)))
    (values min idx)))

(defmethod mabsmax ((m vector-d))
  (let* ((store (vector-store m))
	 (max (aref store 0))
	 (idx 0))
    (declare (type type-blas-store store)
	     (type double-float max)
	     (type type-blas-idx idx))
    (dotimes (i (length store))
      (when (> (abs (aref store i)) (abs max))
	(setf max (aref store i)
	      idx i)))
    (values max idx)))

(defmethod mabsmin ((m vector-d))
  (let* ((store (vector-store m))
	 (min (aref store 0))
	 (idx 0))
    (declare (type type-blas-store store)
	     (type double-float min)
	     (type type-blas-idx idx))
    (dotimes (i (length store))
      (when (< (abs (aref store i)) (abs min))
	(setf min (aref store i)
	      idx i)))
    (values min idx)))

(defmethod mminmax ((m vector-d))
  (let* ((store (vector-store m))
	 (max (aref store 0))
	 (min (aref store 0)))
    (declare (type type-blas-store store)
	     (type double-float max min))
    (dotimes (i (length store))
      (when (> (aref store i) max)
	(setf max (aref store i)))
      (when (< (aref store i) min)
	(setf min (aref store i))))      
    (list min max)))

(defmethod .some (pred (a vector-d) &rest args)
  (let ((stores (mapcar #'vector-store (cons a args))))
    (apply #'some pred stores)))

(defmethod .every (pred (a vector-d) &rest args)
  (let ((stores (mapcar #'vector-store (cons a args))))
    (apply #'every pred stores)))


;;; Matrix and real

(define-constant +generic-function-dfa-df-map+ 
    '((.add . +_dfa-df)
      (.sub . -_dfa-df)
      (.mul . *_dfa-df)
      (.div . /_dfa-df)
      ;; (.expt . ^_dfa-df) taken out specially
      (.max . max_dfa-df) 
      (.min . min_dfa-df)
      (.= . =_dfa-df)
      (./= . /=_dfa-df)
      (.> . >_dfa-df)
      (.>= . >=_dfa-df)
      (.< . <_dfa-df)
      (.<= . <=_dfa-df)
      ))

(defmacro defmethod-dfa-df (name underlying-function)  
  (let ((a (gensym "a"))
	(b (gensym "b"))
	(c (gensym "c")))
  `(defmethod ,name ((,a vector-d) (,b real))
     (let ((,c (mcreate ,a)))
       (,underlying-function (vector-store ,a) (coerce ,b 'double-float) (vector-store ,c))
       ,c))))

(defmacro expand-generic-function-dfa-df-map ()
  (cons 'progn
      (mapcar (lambda (x)
		`(defmethod-dfa-df ,(car x) ,(cdr x)))
	      +generic-function-dfa-df-map+)))

(expand-generic-function-dfa-df-map)

;;; Real and vector 

(define-constant +generic-function-df-dfa-map+ 
    '((.add . +_df-dfa)
      (.sub . -_df-dfa)
      (.mul . *_df-dfa)
      (.div . /_df-dfa)
      ;; (.expt . ^_df-dfa) taken out specially  
      (.max . max_df-dfa) 
      (.min . min_df-dfa)
      (.= . =_df-dfa)
      (./= . /=_df-dfa)
      (.> . >_df-dfa)
      (.>= . >=_df-dfa)
      (.< . <_df-dfa)
      (.<= . <=_df-dfa)
      ))

(defmacro defmethod-df-dfa (name underlying-function)  
  (let ((a (gensym "a"))
	(b (gensym "b"))
	(c (gensym "c")))
  `(defmethod ,name ((,a real) (,b vector-d))
     (let ((,c (mcreate ,b)))
       (,underlying-function (coerce ,a 'double-float) (vector-store ,b) (vector-store ,c))
       ,c))))

(defmacro expand-generic-function-df-dfa-map ()
  (cons 'progn
      (mapcar (lambda (x)
		`(defmethod-df-dfa ,(car x) ,(cdr x)))
	      +generic-function-df-dfa-map+)))

(expand-generic-function-df-dfa-map)

;;;; Vector and vector

(define-constant +generic-function-dfa-dfa-map+ 
    '((.add . +_dfa-dfa)
      (.sub . -_dfa-dfa)
      (.mul . *_dfa-dfa)
      (.div . /_dfa-dfa)
      ;; (.expt . ^_dfa-dfa) taken out specially
      (.max . max_dfa-dfa) 
      (.min . min_dfa-dfa)
      (.= . =_dfa-dfa)
      (./= . /=_dfa-dfa)
      (.> . >_dfa-dfa)
      (.>= . >=_dfa-dfa)
      (.< . <_dfa-dfa)
      (.<= . <=_dfa-dfa)
      ))

(defmacro defmethod-dfa-dfa (name underlying-function)  
  (let ((a (gensym "a"))
	(b (gensym "b"))
	(c (gensym "c")))
  `(defmethod ,name ((,a vector-d) (,b vector-d))
     (let ((,c (mcreate ,a)))
       (,underlying-function (vector-store ,a) (vector-store ,b) (vector-store ,c))
       ,c))))

(defmacro expand-generic-function-dfa-dfa-map ()
  (cons 'progn
      (mapcar (lambda (x)
		`(defmethod-dfa-dfa ,(car x) ,(cdr x)))
	      +generic-function-dfa-dfa-map+)))

(expand-generic-function-dfa-dfa-map)

;;; The expt is an exception, since negative input can give complex exponent
;;; and since the general case is slow.

(defun all-integer-elements-p (a)
  "Helper function for .expt"
  (declare (type-blas-store a))
  (dotimes (i (length a))
    (multiple-value-bind (div mod) (ftruncate (aref a i))
      (declare (ignore div))
      (unless (zerop mod)
	(return-from all-integer-elements-p nil))))
  t)

(defmethod .expt ((a vector-d) (b vector-d))
  (cond ((>= (mmin a) 0d0)
	 (let ((c (mcreate a)))
	   (^_dfa>=0-dfa (vector-store a) (vector-store b) (vector-store c))
	   c))
	((all-integer-elements-p (vector-store b))
	 (let ((c (mcreate a)))
	   (^_dfa-dfa (vector-store a) (vector-store b) (vector-store c))
	   c))
	(t (let ((c (mcreate* a :element-type :z))
		 (a2 (mcreate* a :element-type :z))
		 (b2 (mcreate* a :element-type :z)))
	     ;; There should be some better way to upgrade 
	     ;; the element type of a vector while keeping the contents.
	     (copy-contents a a2)
	     (copy-contents b b2)
	     ;; This could in theory be a litle more clever 
	     ;; since if all exponents are even, the output 
	     ;; is still real. However, recognizing even integers 
	     ;; from float input is risky.
	     (^_cdfa-cdfa (vector-store a2) (vector-store b2) (vector-store c))
	     c))))
	  
(defmethod .expt ((a vector-d) (b real))
  "There is a lot of fuzz going on in here. The reason is because
the important special cases of exponents -3,-2,-1,0,1,2,3 are 
a factor 10 faster than the general case on SBCL. 
Furthermore, output can be complex for non-integer exponent."
  (multiple-value-bind (div mod) (truncate b)
    (if (= 0 mod)
	(let ((c (mcreate a)))
	  (case div 
	    (-3 (^_dfa--3 (vector-store a) (vector-store c)))
	    (-2 (^_dfa--2 (vector-store a) (vector-store c)))
	    (-1 (^_dfa--1 (vector-store a) (vector-store c)))
	    (0  (mfill c 1))
	    (1  (copy-contents a c))
	    (2  (^_dfa-+2 (vector-store a) (vector-store c)))
	    (3  (^_dfa-+3 (vector-store a) (vector-store c)))
	    (t  (^_dfa-df (vector-store a) 
			  (coerce div 'double-float) 
			  (vector-store c))))
	  c)
	(let ((min (mmin a)))
	  (if (>= min 0)
	      (let ((c (mcreate a)))
		(^_dfa>=0-df (vector-store a) 
			     (coerce b 'double-float) 
			     (vector-store c))
		c)
	      (let ((c (mcreate* a :element-type :z))
		    (a2 (mcreate* a :element-type :z)))
		(copy-contents a a2)
		(^_cdfa-cdf (vector-store a2) 
			    (coerce b '(complex double-float)) 
			    (vector-store c))
		c))))))
		
(defmethod .expt ((a real) (b vector-d) )
  (if (>= a 0)
      (let ((c (mcreate b)))
	(^_df>=0-dfa (coerce a 'double-float) 
		     (vector-store b) 
		     (vector-store c))
	c)
      (if (all-integer-elements-p (vector-store b))
	  (let ((c (mcreate b)))	      
	    (^_df-dfa (coerce a 'double-float)  
		      (vector-store b) 
		      (vector-store c))
	    c)
	  (let ((c (mcreate* b :element-type :z))
		(b2 (mcreate* b :element-type :z)))
	    (copy-contents b b2)
	    (^_cdf-cdfa (coerce a '(complex double-float)) 
			(vector-store b2) 
			(vector-store c))
	    c))))


(define-constant +generic-function-dfa-to-dfa-map+ ;really bad name    
    ;; Ordinary functions that map real to real
    ;; Some functions like sqrt are not part of the list 
    ;; since they possibly spit out complex numbers
    '((.sin . sin_dfa-to-dfa) 
      (.cos . cos_dfa-to-dfa) 
      (.tan . tan_dfa-to-dfa) 
      (.atan . atan_dfa-to-dfa)
      (.sinh . sinh_dfa-to-dfa) 
      (.cosh . cosh_dfa-to-dfa) 
      (.tanh . tanh_dfa-to-dfa) 
      (.asinh . asinh_dfa-to-dfa) 
      #+nil (.acosh . acosh_dfa-to-dfa) 
      (.exp . exp_dfa-to-dfa) 
      #+nil (.ln . log_dfa) 
      #+nil (.sqrt . sqrt_dfat) 
      #+nil (.conj . conjugate_dfa)
      #+nil (.re . realpart_dfa) 
      #+nil (.im . imagpart_dfa) 
      (.abs . abs_dfa-to-dfa)
      (.sgn . signum_dfa-to-dfa)))

(defmacro defmethod-dfa-to-dfa (name underlying-function)  
  (let ((a (gensym "a"))
	(b (gensym "b")))
  `(defmethod ,name ((,a vector-d))
     (let ((,b (mcreate ,a)))
       (,underlying-function (vector-store ,a) (vector-store ,b) )
       ,b))))

(defmacro expand-generic-function-dfa-to-dfa-map ()
  (cons 'progn
      (mapcar (lambda (x)
		`(defmethod-dfa-to-dfa ,(car x) ,(cdr x)))
	      +generic-function-dfa-to-dfa-map+)))

(expand-generic-function-dfa-to-dfa-map)

;;; Some exceptions, where output can be either real or complex.

(defmethod .re ((a vector-d))
  (copy a))

(defmethod .im ((a vector-d))
  (mcreate a))

(defmethod .conj ((a vector-d))
  (copy a))

(defmethod .sqrt ((a vector-d))
  (if (>= (mmin a) 0d0)
      (let ((out (mcreate a)))
	(sqrt_dfa-to-dfa (vector-store a) (vector-store out))
	out)
      (let ((out (mcreate* a :element-type :z)))
	(sqrt_dfa-to-cdfa (vector-store a) (vector-store out))
	out)))
	
(defmethod .ln ((a vector-d))
  (if (> (mmin a) 0d0)
      (let ((out (mcreate a)))
	(log_dfa-to-dfa (vector-store a) (vector-store out))
	out)
      (let ((out (mcreate* a :element-type :z)))
	(log_dfa-to-cdfa (vector-store a) (vector-store out))
	out)))
	    
(defmethod .asin ((a vector-d))
  (destructuring-bind (min max) 
      (mminmax a)    
    (if (and (>= min -1d0) 
	     (<= max 1d0))
	(let ((out (mcreate a)))
	  (asin_dfa-to-dfa (vector-store a) (vector-store out))
	  out)
	(let ((out (mcreate* a :element-type :z)))
	  (asin_dfa-to-cdfa (vector-store a) (vector-store out))
	  out))))	

(defmethod .acos ((a vector-d))
  (destructuring-bind (min max) 
      (mminmax a)    
    (if (and (>= min -1d0) 
	     (<= max 1d0))
	(let ((out (mcreate a)))
	  (acos_dfa-to-dfa (vector-store a) (vector-store out))
	  out)
	(let ((out (mcreate* a :element-type :z)))
	  (acos_dfa-to-cdfa (vector-store a) (vector-store out))
	  out))))

(defmethod .atanh ((a vector-d))
  (destructuring-bind (min max) 
      (mminmax a)    
    (if (and (> min -1d0) 
	     (< max 1d0))
	(let ((out (mcreate a)))
	  (atanh_dfa-to-dfa (vector-store a) (vector-store out))
	  out)
	(let ((out (mcreate* a :element-type :z)))
	  (atanh_dfa-to-cdfa (vector-store a) (vector-store out))
	  out))))

(defmethod .acosh ((a vector-d))
  (if (>= (mmin a) 1d0)
      (let ((out (mcreate a)))
	(acosh_dfa-to-dfa (vector-store a) (vector-store out))
	out)
      (let ((out (mcreate* a :element-type :z)))
	(acosh_dfa-to-cdfa (vector-store a) (vector-store out))
	out)))

;;; Some special functions with with complex output for real input

(defmethod .besh1 (n (x vector-d))
  (let ((out (mcreate* x :element-type :z)))
    (mmap-into out (lambda (x) (.besh1 n x)) x)))

(defmethod .besh2 (n (x vector-d))
  (let ((out (mcreate* x :element-type :z)))
    (mmap-into out (lambda (x) (.besh2 n x)) x)))




#|
;;; Now these sad cases where output might be complex for real input

(define-constant +generic-function-dfa-to-cdfa-map+ ;really bad name    
    '((.sqrt . sqrt_dfa)
      (.ln . log_dfa)
      (.asin . asin_dfa)
      (.acos . acos_dfa)
      (.atanh . atanh_dfa)
      ;; Some more?
      )) 

(defmacro defmethod-dfa-to-cdfa (name underlying-function)  
  (let ((a (gensym "a"))
	(b (gensym "b"))
	(spec (gensym "spec")))
  `(defmethod ,name ((,a vector-d))
     (let* ((,spec (cons :z (cdr (type-spec ,a))))
	    (,b (make-matrix-instance ,spec (dim ,a) 0)))
       (,underlying-function (vector-store ,a) (vector-store ,b) )
       ,b))))

(defmacro expand-generic-function-dfa-to-cdfa-map ()
  (cons 'progn
      (mapcar (lambda (x)
		`(defmethod-dfa-to-cdfa ,(car x) ,(cdr x)))
	      +generic-function-dfa-to-cdfa-map+)))

(expand-generic-function-dfa-to-cdfa-map)
|#