;;; -*- Mode: LISP; Syntax: Common-Lisp -*-
;;; Interface to FFTW for Allegro Common Lisp  DOUBLE VERSION

;;; THIS WORKS!!! DOUBLE FLOAT DCT and FFT, any length
;;; However, because arrays in lisp can move at GC, it is imperative
;;; to compute the plans for the fft immediately before calling the fft.
;;; This means that "optimal" plans are not computed, just estimates.
;;; for power of 2 size, this may not matter.

;;; it is possible to find interfaces to FFTW from lisp using CFFI
;;; But they (a) include ALL the interfaces and (b) seem to erect a huge
;;; infrastructure to get things to work that we don't need.

(load  "c:/fftw/libfftw3-3.dll")  

(ff:def-foreign-call 
 (fftw_plan "fftw_plan_dft_1d") ;; 1d complex FFT
 ((n :int) 
  (in  (* :double) (array double-float(*)))
  (out (* :double) (array double-float(*)))
  (sign :int) 
  (flags :unsigned-int))
 :returning :foreign-address)

(ff:def-foreign-call 
    (fftw_plan_r2r "fftw_plan_r2r_1d") ;; real-to-real 1d transform, e.g. discrete cosine
    ((n :int)
     (in  (* :double) (array double-float(*)))
     (out (* :double) (array double-float(*)))
     (kind :int)
     (flags :unsigned-int))
 :returning :foreign-address)

(ff:def-foreign-call (fftw_execute_dft "fftw_execute_dft") ;; to use new array or arrays
    ((n :foreign-address)
     (in  (* :double) (array double-float(*))) ;; these will be complexes stored as 
     (out (* :double) (array double-float(*)))) ;; successive entries
  :returning :void)

(ff:def-foreign-call (fftw_execute_r2r "fftw_execute_r2r") ;; to use new array or arrays
    ((n :foreign-address)
     (in  (* :double) (array double-float(*)))
     (out (* :double) (array double-float(*))))
  :returning :void)

(ff:def-foreign-call (fftw_destroy_plan "fftw_destroy_plan")
    ((n :foreign-address))
  :returning :void)

;; we are not using these ...
;;(ff:def-foreign-call (fftw_malloc "fftw_malloc") ((count :int)) :returning :foreign-address)

;;(ff:def-foreign-call (fftw_free "fftw_free") ((location :foreign-address)) :returning :void)

;;#+ignore 
;;(defmacro mref(A i) ;; memory reference in fftw world
;;  `(sys::memref-int ,A 0 (ash ,i 3) :double-float))

(defmacro mref(A i)`(aref ,A ,i))

;;;;;;;;;;;; end of the interfacing routines that we need for FFTW. 

(defun dct2(input)  ;; DCT-II dct
  
  ;; input array of REAL numbers, (double float here, probably) 
  ;; output array of REAL numbers
  ;;  for even N, FFTW_REDFT10, or DCT-II, is "the" DCT.
  ;;  for  odd N, FFTW_RODFT10, or DCT-II, is "the" DCT.
  ;; Their inverses are RED01 and ROD01, respectively.
  (let*((len (length input))
	(ni (sqrt (/ 1.0d0  (ash len 2))))
	(s (make-array len :element-type 'double-float));answer will go here
	;;(s (make-array len :element-type 'double-float :allocation :static-reclaimable )) ;maybe better?
	)
    (declare (type (array double-float (*))  s)
	     (fixnum len)(double-float ni)
	     (optimize (speed 3)(safety 0)))
            ;; copy and pre-normalize
    	(loop for i fixnum from 0 to (1- len) do 
	      (setf (mref s i)(* ni (aref input i))))
	;;execute the plan, which is a cheap (estimate-only) version.
	(fftw_execute_r2r (fftw_plan_r2r len s s (if (evenp len) 5 9) 64) s s)
	;; return the array that was allocated for this.
	;; Note that we are NOT over-writing the input, which some
	;; people seem to like to do in FORTRAN world.
	s))

(defun dct3(input)  ;; DCT-III. See comments for dct2 
  (let*((len (length input))
	(ni (sqrt (/ 1.0d0 len)))
	(s (make-array len :element-type 'double-float))
	;;(s (make-array len :element-type 'double-float :allocation :static-reclaimable))
	)
    (declare (type (array double-float (*)) s) (fixnum len)(double-float ni)
	     (optimize (speed 3)(safety 0)))
     ;; copy, convert to float if necessary,  and pre-normalize
	(loop for i fixnum from 0 to (1- len) do 
	      (setf (mref s i)(* ni (aref input i))))
	;;execute the plan	      
	(fftw_execute_r2r (fftw_plan_r2r len s s (if (evenp len) 4 8) 64) s s)
	s))
					;
(defun rtdct(x)(dct3(dct2 x))) ;; round-trip test should be identity.

(defun fft(input) 
  ;; input is array of possibly complex numbers, 
  ;; returns array of complex double-float.
  ;; computes a plan, if there is not one.
  (let*((len (length input)) (ni (/ 1.0d0 len))
	(r  (make-array (+ len len) :element-type 'double-float)) ;; workspace
	(ans (make-array len))) ; complex numbers. I suppose we could put them in "input"
    (declare (fixnum len)
	     (type (array double-float (*)) r) 
	     (optimize (speed 3)(safety 0)))
    ;; array is twice the length of the input to make room for real & imag parts
    ;; copy real/imag parts of input into adjacent array locations and pre-normalize
;    (print 'copy)
    (loop for i fixnum from 0 to (1- len)  do 
	  (setf (mref r (+ i i))(* ni (realpart(aref input i)))) ;;0, 2, 4
	  (setf (mref r (+  i i 1)) (* ni (imagpart(aref input i)))) ) ;; 1, 3,
    (fftw_execute_dft (fftw_plan len r r -1 64) r r)	      ;; do the FFT
    (loop for i fixnum from 0 to (1- len)  do ;; copy parts into complex numbers
	  (setf (aref ans i) (complex (mref r (+ i i))
				      (mref r (1+(+ i i))))))
    ans))

(defun ffti(input) 
  ;; input is array of possibly complex numbers, 
  ;; returns array of complex double-float.
  ;; computes a plan, if there is not one.
  (let* ((len (length input)) 
	 (r  (make-array (ash len 1) :element-type 'double-float))
	 (ans(make-array len)))		; complex numbers. I suppose we could put them in "input"
    (declare (fixnum len)
	     (type (array double-float (*)) r) (optimize (speed 3)(safety 0)))
    ;; array is twice the length of the input to make room for real & imag parts
    ;; copy real/imag parts of input into adjacent array locations and pre-normalize
    (loop for i fixnum from 0 to (1- len)  do 
	  (setf (mref r (+ i i))(float (realpart(aref input i)) 1.0d0)) ;;0, 2, 4
	  (setf (mref r (1+(+ i i)))(float(imagpart(aref input i)) 1.0d0)) ) ;; 1, 3,
    (fftw_execute_dft(fftw_plan len r r 1 64) r r) ;; do the FFT
    (loop for i fixnum from 0 to (1- len)  do  ;; copy parts into complex numbers
	  (setf (aref ans i) (complex (mref r (+ i i))
				      (mref r (+ i i 1)))))
    ans))

(defun rtl(ar)  (ffti(fft ar))) ;; round trip test of fft

#|
(setf tt (coerce (loop for i from 0 to 1999 collect 1.0d0) '(array double-float (*))))
(setf ss (coerce (loop for i from 0 to 500 collect 1.0d0) 'array))


|#

;;; here's a way of computing an optimal plan (cf fftw documents) and reusing it.

;;; say you have many arrays of the same length that you are going to compute fft.
;;; compute a workspace and a plan this way.

;;;  (setf myplan (fftplan (length input) 0)) ;; allocates buffer and computes optitmal plan.
;;;  for each array A, use this..
;;;  (fftp-do A 0 myplan)


(defun fftplan(len flag); flag is 0 for optimal plan, 64 for estimate
  (let ((r (make-array (+ len len) :element-type 'double-float :allocation :static-reclaimable)))
    (list len 
	  r
	  (fftw_plan len r r -1 flag))))
	  

(defun fftp-do(input plan)   (apply #'fft-with-plan input plan))

(defun fftk(input)(fftp-do input (fftplan (length input) 64))) ;; just use estimate

(defun fft-with-plan (input len r plan)  ; fft with pre-computed plan
  ;; input is array of possibly complex numbers, 
  ;; returns array of complex double-float.
  ;; computes a plan, if there is not one.
  (let  ((ni (/ 1.0d0 len))
	(ans (make-array len))) ; complex numbers. I suppose we could put them in "input"
    (declare (fixnum len)
	     (type (array double-float (*)) r) 
	     (optimize (speed 3)(safety 0)))
    ;; array is twice the length of the input to make room for real & imag parts
    ;; copy real/imag parts of input into adjacent array locations and pre-normalize
;    (print 'copy)
    (loop for i fixnum from 0 to (1- len)  do 
	  (setf (mref r (+ i i))(* ni (realpart(aref input i)))) ;;0, 2, 4
	  (setf (mref r (+  i i 1)) (* ni (imagpart(aref input i)))) ) ;; 1, 3,
    (fftw_execute_dft plan r r)	      ;; do the FFT
    (loop for i fixnum from 0 to (1- len)  do ;; copy parts into complex numbers
	  (setf (aref ans i) (complex (mref r (+ i i))
				      (mref r (1+(+ i i))))))
    ans))