;;; -*- Mode: LISP; Syntax: Common-Lisp -*-
;;; Interface to FFTW for Allegro Common Lisp  DOUBLE VERSION
;;; THIS WORKS!!! DOUBLE FLOAT DCT and FFT, any length
;;; 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 (destroy_plan "fftw_destroy_plan")
    ((n :foreign-address))
  :returning :void)


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

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




(defvar *dct2plans* (make-hash-table))
(defvar *dct3plans* (make-hash-table))
(defvar *dctbufsize* 8)
(defvar *dctbuf* (make-array *dctbufsize* :element-type 'double-float))


(defun dct2(input)  ;; DCT-II
  
  ;; 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))))
	(r (if (<= len *dctbufsize*) *dctbuf*
	     (setf *dctbuf* (make-array (setf *dctbufsize* len) 
					 :element-type 'double-float))))
	(s (make-array len :element-type 'double-float)) ;answer will go here
	(dct2plan (gethash len *dct2plans*)))
    (declare (type (array double-float (*)) r s)(optimize (speed 3)(safety 0)))
    ;; if no plan for transforms of this length, make one.
    (cond ((null dct2plan)
	   (setf (gethash len *dct2plans*) 
	     (setf dct2plan (fftw_plan_r2r len r s (if (evenp len) 5 9) 0)))))
    ;; copy and pre-normalize
    ;; clear r? should not be necessary
    (loop for i fixnum from 0 to (1- *dctbufsize*) do
	  (setf (aref r i) 0.0d0))
	(loop for i fixnum from 0 to (1- len) do 
	      (setf (aref r i)(* ni (float (aref input i) 1.0d0))))
	;;execute the plan	      
	(fftw_execute_r2r dct2plan r 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))
;;hm.  (dct2 #(1 2 3 4 5 6 7 8)) answers..
;;#(12.727922279186966d0 -4.555410373911547d0 3.2380243244781607d-16 ...)
;;..12.72792206135786  is better.

(defun dct3(input)  ;; DCT-III. See comments for dct2
  (let*((len (length input))
	(ni (sqrt (/ 1.0d0 len)))
	(r (if (<= len *dctbufsize*) *dctbuf*
	     (setf *dctbuf* (make-array (setf *dctbufsize* len) 
					 :element-type 'double-float))))
	(s (make-array len :element-type 'double-float))
	(dct3plan (gethash len *dct3plans*)))
    (declare (type (array double-float (*)) r s)(optimize (speed 3)(safety 0)))
    (cond ((null dct3plan)
	   (setf (gethash len *dct3plans*) 
	     (setf dct3plan (fftw_plan_r2r len r s (if (evenp len) 4 8) 0)))))
    ;; copy and pre-normalize
	(loop for i fixnum from 0 to (1- len) do 
	      (setf (aref r i)(* ni (float (aref input i) 1.0d0))))
	;;execute the plan	      
	(fftw_execute_r2r dct3plan r s)
	s))

;(defun rtdct(x)(dct3(dct2 x))) ;; round-trip test should be identity.

(defvar *fftplans* (make-hash-table))
(defvar *fftiplans* (make-hash-table))
(defvar *fftbufsize* 8)
(defvar *fftbuf* (make-array *fftbufsize* :element-type 'double-float))

(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 (if (<= (ash len 1) *fftbufsize*) *fftbuf*
	     (setf *fftbuf* (make-array (setf *fftbufsize* (ash len 1)) 
					:element-type 'double-float))))	 
	(fftplan (gethash len *fftplans*))
	(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
    (cond ((null fftplan)
	   (setf (gethash len *fftplans*) 
	    ; (setf fftplan (fftw_plan len r r -1 0)) ;; optimal
	     (setf fftplan (fftw_plan len r r -1 64)) ;; estimate
	     )))
    ;; copy real/imag parts of input into adjacent array locations and pre-normalize
    (loop for i fixnum from 0 to (1- len)  do 
	  (setf (aref r (+ i i))(* ni (realpart(aref input i)))) ;;0, 2, 4
	  (setf (aref r (1+ (+ i i)))(* ni (imagpart(aref input i)))) ) ;; 1, 3,
    (fftw_execute_dft fftplan 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 (aref r (+ i i))
				      (aref r (1+(+ i i))))))
    ans))

(defun unplan()
  (maphash #'(lambda(i j)(destroy_plan j)) *fftplans*)
  (maphash #'(lambda(i j)(destroy_plan j)) *fftiplans*)
  (maphash #'(lambda(i j)(destroy_plan j)) *dct2plans*)
  (maphash #'(lambda(i j)(destroy_plan j)) *dct3plans*)
  
  (setf *fftplans* (make-hash-table))
  (setf *fftiplans* (make-hash-table))
  (setf *fftbufsize* 8)
  (setf *fftbuf* (make-array *fftbufsize* :element-type 'double-float))
  (setf *dct2plans* (make-hash-table))
  (setf *dct3plans* (make-hash-table))
  (setf *dctbufsize* 8)
  (setf *dctbuf* (make-array *dctbufsize* :element-type 'double-float))
  'plans-gone
  )

(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 (if (<= (ash len 1) *fftbufsize*) *fftbuf*
	     (setf *fftbuf* (make-array (setf *fftbufsize* (ash len 1)) 
					 :element-type 'double-float))))	 
	(fftiplan (gethash len *fftiplans*))
	(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
    (cond ((null fftiplan)
	   (setf (gethash len *fftiplans*) 
	     (setf fftiplan (fftw_plan len r r 1 0))))) ;; 1 means inverse.
    ;; copy real/imag parts of input into adjacent array locations and pre-normalize
    (loop for i fixnum from 0 to (1- len)  do 
	  (setf (aref r (+ i i))(float (realpart(aref input i)) 1.0d0)) ;;0, 2, 4
	  (setf (aref r (1+(+ i i)))(float(imagpart(aref input i)) 1.0d0)) ) ;; 1, 3,
    (fftw_execute_dft fftiplan 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 (aref r (+ i i))
				      (aref r (1+(+ i i))))))
    ans))

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