;;; Lisplab, level3-fft-zge.lisp
;;; Methods for fast Fourier transforms specialized on blas 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.

;;; TODO should use the normal ref-blas-complex-store ?

;;; TODO make shift routines for complex matrices also.

(in-package :lisplab)

;; Optimized shift methods for real matrix. 
(defmethod fft-shift ((m matrix-base-dge))
  (let* ((rows (rows m))
         (fr (floor rows 2))
         (cr (ceiling rows 2))
         (cols (cols m))
         (fc (floor cols 2))
         (cc (ceiling cols 2))
         (store-m (vector-store m))
         (m2 (mcreate m))
         (store-m2 (vector-store m2)))
         
    (declare (type type-blas-store store-m store-m2)
             (type type-blas-idx fr cr fc cc rows cols))
    (dotimes (i rows)
      (dotimes (j cols)
        (setf (aref store-m2 (column-major-idx i j rows))
              (aref store-m (column-major-idx  (if (< i fr) (+ i cr) (- i fr))
                                               (if (< j fc) (+ j cc) (- j fc))
                                               rows)))))
    m2))

(defmethod ifft-shift ((m matrix-base-dge))
  (let* ((rows (rows m))
         (fr (floor rows 2))
         (cr (ceiling rows 2))
         (cols (cols m))
         (fc (floor cols 2))
         (cc (ceiling cols 2))
                 (store-m (vector-store m))
         (m2 (mcreate m))
         (store-m2 (vector-store m2)))
         
    (declare (type type-blas-store store-m store-m2)
             (type type-blas-idx fr cr fc cc rows cols))
    (dotimes (i rows)
      (dotimes (j cols)
        (setf (aref store-m2 (column-major-idx i j rows))
              (aref store-m (column-major-idx  (if (< i cr) (+ i fr) (- i cr))
                                               (if (< j cc) (+ j fc) (- j cc))
                                               rows)))))
    m2))

;;;; The implementing methods

(defmethod fft1! ((x matrix-lisp-zge) &key)
  (assert (= 1 (logcount (rows x))))
  (dotimes (i (cols x))
    (fft-radix-2-blas-complex-store! :f (vector-store x) (rows x) (* (rows x) i) 1))
  x)

(defmethod ifft1! ((x matrix-lisp-zge) &key)
  (assert (= 1 (logcount (rows x))))
  (dotimes (i (cols x))
    (fft-radix-2-blas-complex-store! :r (vector-store x) (rows x) (* (rows x) i) 1))
  x)

(defmethod fft2! ((x matrix-lisp-zge) &key)
  (assert (and (= 1 (logcount (rows x)))
	       (= 1 (logcount (cols x)))))
  (fft1! x) 
  (dotimes (i (rows x))
    (fft-radix-2-blas-complex-store! :f (vector-store x) (cols x) i (rows x)))
  x)

(defmethod ifft2! ((x matrix-lisp-zge) &key)
  (assert (and (= 1 (logcount (rows x)))
	       (= 1 (logcount (cols x)))))

  (ifft1! x) 
  (dotimes (i (rows x))
    (fft-radix-2-blas-complex-store! :r (vector-store x) (cols x) i (rows x)))
  x)

(defun fft-radix-2-blas-complex-store! (direction x n start step)
  "Destrutive, radix 2 fast fourier transform. Direction is either :f for 
forward or :r for reverse transform. Input must be a
vector of complex double float"
  (let* ((ftx x)
	 (sign (ecase direction (:f -1d0) (:r 1d0))))
    (declare (type-blas-idx n start step)
	     (double-float sign)
	     (type-blas-store ftx))
    (bit-reverse-blas-complex-store! ftx n start step)
    ;; apply fft recursion
    (dotimes (bit (floor (log n 2)))
      (let* ((dual (expt 2 bit))
	     (W #C(1d0 0d0))
	     (tmp (- (exp (/ (* sign %i pi) dual)) 1d0 )))
	(declare (type type-blas-idx dual)
		 (type (integer 0 30) bit)
		 (type (complex double-float) W tmp))
	(loop for b from 0 below n by (* 2 dual) do 
	      (let* ((wd (ref-blas-complex-store2 ftx 
						  (truly-the type-blas-idx (+ b dual)) 
						  start 
						  step)))
		(declare (type-blas-idx b)
			 ((complex double-float) Wd))
		(setf (ref-blas-complex-store2 ftx (truly-the type-blas-idx (+ b dual)) 
					       start 
					       step) 
		      (- (ref-blas-complex-store2 ftx b start step) wd))
		(incf (ref-blas-complex-store2 ftx b start step) 
		      wd)))   
	(loop for a from 1 to (- dual 1) do 
	      ;; trignometric recurrence for w-> exp(i theta) w 
	      (incf w (* w tmp))
	      (loop for b from 0 to (- n 1) by (* 2 dual) do
		    (let* ((i (truly-the type-blas-idx  (+ b a)))
			   (j (truly-the type-blas-idx (+ i dual)))
			   (wd (* w (ref-blas-complex-store2 ftx j start step))))
		      (declare ( type-blas-idx b i j)
			       ((complex double-float) Wd))
		      (setf (ref-blas-complex-store2 ftx j start step) 
			    (- (ref-blas-complex-store2 ftx i start step) wd))
		      (incf (ref-blas-complex-store2 ftx i start step) 
			    wd))))))
    ftx))

(defun bit-reverse-blas-complex-store! (vec n start step)
  ;; This is the Goldrader bit-reversal algorithm
  (let ((j 0))
    (declare (type-blas-store vec)
	     (type-blas-idx j n start step))
    (dotimes (i (- n 1))
      (let ((k (floor n 2)))
	(declare (type-blas-idx i k))
	(when (< i j)
	  (let ((tmp (ref-blas-complex-store2 vec i start step)))
	    (setf (ref-blas-complex-store2 vec i start step) 
		  (ref-blas-complex-store2 vec j start step)
		  (ref-blas-complex-store2 vec j start step) 
		  tmp)))
	(do () ((> k j))
	  (setf j (the type-blas-idx (- j k))
		k (floor k 2)))
	(incf j k))))
  vec)