;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; LUSH Lisp Universal Shell
;;;   Copyright (C) 2002 Leon Bottou, Yann Le Cun, AT&T Corp, NECI.
;;; Includes parts of TL3:
;;;   Copyright (C) 1987-1999 Leon Bottou and Neuristique.
;;; Includes selected parts of SN3.2:
;;;   Copyright (C) 1991-2001 AT&T Corp.
;;;
;;; 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., 59 Temple Place, Suite 330, Boston, MA 02111, USA
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; $$
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


(libload "dh-compile")
(libload "libc/constants")

(libload "fftw-config")

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
#? *** fftw: an interface to the FFTW3 Fourier Transform library
;; This is an interface to the FFTW3 Fourier Transform
;; library. To use this interface, you must first install
;; the FFTW3 library and development package.
;; fftw3 is distributed with most Linux distros
;; as fftw3 or libfftw3 and libfftw3-devel.

#? * High Level Fourier Transform functions

#? (spectrogram <signal> <window> <stride>)
;; compute and return the spectrogram of 
;; signal in vector <signal>. The Fourier transform
;; is performed on Hanning windows of width <window>,
;; that are stepped by <stride> samples.
;; example:
;; {<code> (spectrogram m 512 256) </code>}
(de spectrogram (m win stride)
  (let* ((n (idx-dim m 0))
	 (nn (div (- n win)  stride))
	 (mm (narrow m 0 (+ (* nn stride) win) 0))
	 (z (unfold mm 0 win stride))
	 (spectre (double-matrix (idx-dim z 0) (idx-dim z 1)))
	 (w (double-matrix (idx-dim z 1)))
	 (hw (make-hanning-window win)))
    (idx-bloop ((in z) (out spectre))
      (idx-mul in hw w)
      (idx1-power-spectrum w out))
    ;; chop off the top half of the frequency spectrum, which is zero
    (narrow spectre 1 (div (idx-dim spectre 1) 2) 0)))


#? (idx1-power-spectrum <in> <out>)
;; compute the power spectrum of a waveform <in>
;; and put the result in <out>. <in> and <out>
;; must be idx1 of double with identical sizes.
;; only the bottom half of out is filled with useful data
(de idx1-power-spectrum (in out)
  ((-idx1- (-double-)) in)
  ((-idx1- (-double-)) out)
  (when (<> (idx-dim in 0) (idx-dim out 0)) 
    (error "in and out must have the same size"))
  (let ((ft (matrix (idx-dim in 0) 2)))
    (fftw-dft-r2c-1d in ft t)
    (idx1-complex2norm ft out))
  ())

#? (idx1-complex2norm <in> <out>)
;; Compute the vector of norms  
;; of a vector of complex numbers.
;; <in> must be an Nx2 idx2 of doubles
;; where element (i,0) contains the real part
;; of the i-th complex and (i,1) its imaginary
;; part. On output, the idx1 of double <out>
;; (which must be of dimension N) will contain
;; the norms.
(de idx1-complex2norm (in out)
  ((-idx2- (-double-)) in)
  ((-idx1- (-double-)) out)
  (when (<> 2 (idx-dim in 1)) (error "not a complex vector"))
  (when (<> 1 (idx-modulo in 1)) (error "real and imaginary parts not contiguous"))
  (cidx-bloop ("i" ("in" in) ("out" out)) #{ *out = sqrt( in[0]*in[0] + in[1]*in[1] ); #}) ())


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

#? * masking window utilities

#? (make-hanning-window-tanh <n> <p>)
;; make <n>-sample Hanning window with a piece 
;; of hyperbolic tangent function at each end.
(de make-hanning-window-tanh (n p)
  ((-int-) n p)
  (let* ((out (double-matrix n))
	 (o (* 0.5 p))
	 (k (/ 5 (* 0.5 p))))
    (for (i 0 p)
	 (let ((x (* 0.5 (1+ (tanh (* k (- i o)))))))
	   (out i x)
	   (out (- (- p 1) i -1) x)))
    out))

#? (make-hanning-window-tanh <n>)
;; make an <n>-sample Hanning window with 
;; a piece of cosine at each end.
(de make-hanning-window (n)
  (let* ((r (double-matrix n))
	 (tr (int (* n 0.1))))
    (idx-d1fill r 1.0)
    (for (i 0 tr) 
	 (r i (* 0.5 (- 1 (cos (* (/ 3.1415927 tr) i)))))
	 (r (- (- n 1) i) (r i)))
    r))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
(progn
  (defconstant "FFTW_FORWARD" -1)
  (defconstant "FFTW_BACKWARD" +1)
  ;; /* documented flags */
  (defconstant "FFTW_MEASURE" 0)
  (defconstant "FFTW_DESTROY_INPUT" 1)
  (defconstant "FFTW_UNALIGNED" 2)
  (defconstant "FFTW_CONSERVE_MEMORY" 4)
  (defconstant "FFTW_EXHAUSTIVE" 8)
  (defconstant "FFTW_PRESERVE_INPUT" 16)
  (defconstant "FFTW_PATIENT" 32)
  (defconstant "FFTW_ESTIMATE" 64)

  (defconstant "FFTW_NO_SIMD" (bitshl 1 17))
  )

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
#? * Low-Level Fast Fourier Transform funtions (FFT)
;; These functions perform 1D and 2D FFTs
;; from complex to complex, and real to complex.


#? (fftw-dft-c2c-1d <in> <out> <forward> <norm>)
;; compute the complex-to-complex 1D fourier 
;; transform of complex vector <in>, and put
;; result in <out>. <forward> is an integer which
;; should be set to 1 for a forward transform and
;; to -1 for an inverse transform. <norm> is a 
;; boolean that should be set to nil for no
;; normalization and t for normalization.
(de fftw-dft-c2c-1d (in out forward norm)
  ((-idx2- (-double-)) in)
  ((-idx2- (-double-)) out)
  ((-int-) forward)
  ((-bool-) norm)
  (when (<> 2 (idx-dim in 1)) (error "second dimension of input should be 2"))
  (when (<> 2 (idx-dim out 1)) (error "second dimension of output should be 2"))
  (when (<> (idx-dim in 0) (idx-dim out 0)) (error "input and output have different sizes"))
  (when (or (not (contiguousp in)) (not (contiguousp out))) (error "vector not contiguous"))
  (let ((n (idx-dim in 0)))
    ((-int-) n)
    #{{ fftw_plan p;
        p = fftw_plan_dft_1d($n, IDX_PTR($in,fftw_complex), IDX_PTR($out,fftw_complex), 
                             (($forward>0)?FFTW_FORWARD:FFTW_BACKWARD), FFTW_ESTIMATE);
        fftw_execute(p);
        fftw_destroy_plan(p);
     }#}
    (when norm
      (let ((s (/ 1 (to-double n))))
	((-double-) s)
	(cidx-bloop ("i" ("out" out)) #{ out[0] *= $s; out[1] *= $s; #} ))) ()))

#? (fftw-dft-r2c-1d <in> <out> <norm>)
;; compute the real-to-complex 1D fourier 
;; transform of real vector <in>, and put
;; result in the complex vector <out>. 
;; <norm> is a boolean that should be set to nil 
;; for no normalization and t for normalization.
(de fftw-dft-r2c-1d (in out norm)
  ((-idx1- (-double-)) in)
  ((-idx2- (-double-)) out)
  ((-bool-) norm)
  (when (<> 2 (idx-dim out 1)) (error "second dimension of output should be 2"))
  (when (<> (idx-dim in 0) (idx-dim out 0)) (error "input and output have different sizes"))
  (when (or (not (contiguousp in)) (not (contiguousp out))) (error "vector not contiguous"))
  (let ((n (idx-dim in 0)))
    ((-int-) n)
    #{{ fftw_plan p;
        p = fftw_plan_dft_r2c_1d($n, IDX_PTR($in,double), IDX_PTR($out,fftw_complex), FFTW_ESTIMATE);
        fftw_execute(p);
        fftw_destroy_plan(p);
     }#}
    (when norm
      (let ((s (/ 1 (to-double n))))
	((-double-) s)
	(cidx-bloop ("i" ("out" out)) #{ out[0] *= $s; out[1] *= $s; #} ))) ()))


#? (fftw-dft-c2c-2d <in> <out> <forward> <norm>)
;; compute the complex-to-complex 2D fourier 
;; transform of complex matrix <in>, and put
;; result in <out>. <forward> is an integer which
;; should be set to 1 for a forward transform and
;; to -1 for an inverse transform. <norm> is a 
;; boolean that should be set to nil for no
;; normalization and t for normalization.
(de fftw-dft-c2c-2d (in out forward norm)
  ((-idx3- (-double-)) in)
  ((-idx3- (-double-)) out)
  ((-int-) forward)
  ((-bool-) norm)
  (when (<> 2 (idx-dim in 2)) (error "second dimension of input should be 2"))
  (when (<> 2 (idx-dim out 2)) (error "second dimension of output should be 2"))
  (when (or (<> (idx-dim in 0) (idx-dim out 0)) (<> (idx-dim in 1) (idx-dim out 1))) (error "input and output have different sizes"))
  (when (or (not (contiguousp in)) (not (contiguousp out))) (error "vector not contiguous"))
  (let ((n0 (idx-dim in 0))
	(n1 (idx-dim in 1)))
    ((-int-) n0 n1)
    #{{ fftw_plan p;
        p = fftw_plan_dft_2d($n1, $n0, IDX_PTR($in,fftw_complex), IDX_PTR($out,fftw_complex), 
                             (($forward>0)?FFTW_FORWARD:FFTW_BACKWARD), FFTW_ESTIMATE);
        fftw_execute(p);
        fftw_destroy_plan(p);
     }#}
    (when norm
      (let ((s (/ 1 (to-double (* n0 n1)))))
	((-double-) s)
	(cidx-bloop ("i" "j" ("out" out)) #{ out[0] *= $s; out[1] *= $s; #} ))) ()))


#? (fftw-dft-r2c-2d <in> <out> <norm>)
;; compute the real-to-complex 2D fourier 
;; transform of real matrix <in>, and put
;; result in the complex vector <out>. 
;; <norm> is a boolean that should be set to nil 
;; for no normalization and t for normalization.
(de fftw-dft-r2c-2d (in out norm)
  ((-idx2- (-double-)) in)
  ((-idx3- (-double-)) out)
  ((-bool-) norm)
  (when (<> 2 (idx-dim out 2)) (error "second dimension of output should be 2"))
  (when (or (<> (idx-dim in 0) (idx-dim out 0)) (<> (idx-dim in 1) (idx-dim out 1))) (error "input and output have different sizes"))
  (when (or (not (contiguousp in)) (not (contiguousp out))) (error "vector not contiguous"))
  (let ((n0 (idx-dim in 0))
	(n1 (idx-dim in 1)))
    ((-int-) n0 n1)
    #{{ fftw_plan p;
        p = fftw_plan_dft_r2c_2d($n1, $n0, IDX_PTR($in,double), IDX_PTR($out,fftw_complex), FFTW_ESTIMATE);
        fftw_execute(p);
        fftw_destroy_plan(p);
     }#}
    (when norm
      (let ((s (/ 1 (to-double (* n0 n1)))))
	((-double-) s)
	(cidx-bloop ("i" "j" ("out" out)) #{ out[0] *= $s; out[1] *= $s; #} ))) ()))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; this function is not finished
(de fft (in &optional (forward 1) (norm ()))
  (error "function not implemented")
  (let ((m ()))
    (if (= (classof (idx-storage in)) (classof (idx-storage (double-matrix-nc))))
	(setq m in)
      (if (and (> forward 0) (<> 2 (idx-dim in (1- (idx-ndim in)))))
	  (progn (setq m (apply matrix (idx-dim in))) (idx-copy in m))
	(********)))
    (selectq (idx-ndim m)
      (1
       (if (> forward 0)
	   (let ((r (matrix (idx-dim m 0) 2))) (fftw-dft-r2c-1d m r) r)
	 (let ((s (matrix (idx-dim m 0) 2)) (r (matrix (idx-dim m 0) 2))) 
	   (idx-copy s (fftw-dft-r2c-1d m r) r))))
      (2 
       (if (= 2 (idx-dim m 1))
	   (let ((r (matrix (idx-dim m 0) 2))) (fftw-dft-c2c-1d m r) r)
	 (let ((r (matrix (idx-dim m 0) 2))) (fftw-dft-r2c-1d m r) r))))))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; low level interface

;; fftw_execute(const X(plan) p)
;; fftw_plan_dft(int rank, const int *n,				\
;; fftw_plan_dft_1d(int n, C *in, C *out, int sign,			\
;; fftw_plan_dft_2d(int nx, int ny,					\
;; fftw_plan_dft_3d(int nx, int ny, int nz,				\
;; fftw_plan_many_dft(int rank, const int *n,			\
;; fftw_plan_guru_dft(int rank, const X(iodim) *dims,		\
;; fftw_plan_guru_split_dft(int rank, const X(iodim) *dims,		\
;; fftw_execute_dft(const X(plan) p, C *in, C *out)
;; fftw_execute_split_dft(const X(plan) p, R *ri, R *ii, R *ro, R *io)
;; fftw_plan_many_dft_r2c(int rank, const int *n,			\
;; fftw_plan_dft_r2c(int rank, const int *n,				\
;; fftw_plan_dft_r2c_1d(int n,R *in,C *out,unsigned flags)
;; fftw_plan_dft_r2c_2d(int nx, int ny,				\
;; fftw_plan_dft_r2c_3d(int nx, int ny,				\
;; fftw_plan_many_dft_c2r(int rank, const int *n,			\
;; fftw_plan_dft_c2r(int rank, const int *n,				\
;; fftw_plan_dft_c2r_1d(int n,C *in,R *out,unsigned flags)
;; fftw_plan_dft_c2r_2d(int nx, int ny,				\
;; fftw_plan_dft_c2r_3d(int nx, int ny,				\
;; fftw_plan_guru_dft_r2c(int rank, const X(iodim) *dims,		\
;; fftw_plan_guru_dft_c2r(int rank, const X(iodim) *dims,		\
;; fftw_plan_guru_split_dft_r2c(int rank, const X(iodim) *dims,	\
;; fftw_plan_guru_split_dft_c2r(int rank, const X(iodim) *dims,	\
;; fftw_execute_dft_r2c(const X(plan) p, R *in, C *out)
;; fftw_execute_dft_c2r(const X(plan) p, C *in, R *out)
;; fftw_execute_split_dft_r2c(const X(plan) p, R *in, R *ro, R *io)
;; fftw_execute_split_dft_c2r(const X(plan) p, R *ri, R *ii, R *out)
;; fftw_plan_many_r2r(int rank, const int *n,			\
;; fftw_plan_r2r(int rank, const int *n, R *in, R *out,		\
;; fftw_plan_r2r_1d(int n, R *in, R *out,				\
;; fftw_plan_r2r_2d(int nx, int ny, R *in, R *out,			\
;; fftw_plan_r2r_3d(int nx, int ny, int nz,				\
;; fftw_plan_guru_r2r(int rank, const X(iodim) *dims,		\
;; fftw_execute_r2r(const X(plan) p, R *in, R *out)
;; fftw_destroy_plan(X(plan) p)
;; fftw_forget_wisdom(void)
;; fftw_cleanup(void)
;; fftw_plan_with_nthreads(int nthreads)
;; fftw_init_threads(void)
;; fftw_cleanup_threads(void)
;; fftw_export_wisdom_to_file(FILE *output_file)
;; fftw_export_wisdom_to_string(void)
;; fftw_export_wisdom(void (*write_char)(char c, void *), void *data)
;; fftw_import_system_wisdom(void)
;; fftw_import_wisdom_from_file(FILE *input_file)
;; fftw_import_wisdom_from_string(const char *input_string)
;; fftw_import_wisdom(int (*read_char)(void *), void *data)
;; fftw_fprint_plan(const X(plan) p, FILE *output_file)
;; fftw_print_plan(const X(plan) p)
;; fftw_malloc(size_t n)
;; fftw_free(void *p)
;; fftw_flops(const X(plan) p, double *add, double *mul, double *fma)


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(let ((dhc-make-lushflags (concat dhc-make-lushflags fftw-cflags)))
  (dhc-make-with-libs 
   () 
   (list fftw-libfftw3)
   ;; Headers
   #{ #include <fftw3.h> #}
   ;; functions
   fftw-dft-c2c-1d
   fftw-dft-r2c-1d
   fftw-dft-c2c-2d
   fftw-dft-r2c-2d
   idx1-complex2norm
   idx1-power-spectrum
   ))
