(*
 *  Copyright (c) 2005, 2006, 2007 Abram Hindle
 *  
 *  This file is part of CaptchaBreaker

 *  CaptchaBreaker 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 3 of the License, or
 *  (at your option) any later version.

 *  Foobar 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, see <http://www.gnu.org/licenses/>.
 *)

let hypot x y = sqrt ( x*.x +. y*.y ) ;;
let pi = 4.0 *. atan 1.0;;
let pi2 = 2.0 *.pi;;
let npi = (-1.0)*.pi ;;


let fstdo () = flush stdout;;
let fftdebug = false;;
let debugf =  if fftdebug then     (fun s -> (  print_string "FFT: "; print_string s;      print_string "\n";      fstdo();)) else (fun s -> () )
;;


class jfft =
  object(self)
	method reverseBits index numBits =
	  let rec loop rev ind = function
	      x when (x=numBits) -> rev
	    | x -> loop ((rev lsl 1) lor (ind land 1)) (ind lsr 1) (x+1)
          in
	  loop 0 index 0
  
	    
	method isPowerOfTwo = function 
	    x when x < 2 -> false
	  | x when (x land (x-1) > 0) -> false
	  | x -> true

	method numberOfBitsNeeded = function
	    x when x < 2 -> 1
	  | x ->
	      let rec loop i x =
		if ((x land (1 lsl i))>0) then
		  i
		else
		  loop (i+1) x
	      in
	      loop 0 x

      method fftrun inverseTransform realIn imagIn realOut imagOut =
	        let numSamples = Array.length realOut in
		let lri = Array.length realIn in
		let lii = Array.length imagIn in
		let lro = Array.length realOut in
		let lio = Array.length imagOut in

		assert(lri = lii && lro = lio && lii = lro );
		assert(self#isPowerOfTwo numSamples);
		Array.fill realOut 0 lro 0.0;
		Array.fill imagOut 0 lio 0.0;
		let angle_numerator = if ( inverseTransform )	then
		       -1.0 *. pi2 
		else pi2
		 in

		let numBits = self#numberOfBitsNeeded numSamples in

		Array.iteri (fun i x -> (
			let j = self#reverseBits  i numBits in
		        realOut.(j) <- x;
	        	imagOut.(j) <- imagIn.(i);
		)) realIn;

		debugf ("numSamples "^(string_of_int numSamples)^" "^(string_of_int numBits));
		let rec blockSizeNumSamplesLoop = function
		    blockSize when blockSize > numSamples -> ()
		  | blockSize ->
		  debugf ("BlockSize: " ^ (string_of_int blockSize));
	          let blockEnd = if (blockSize = 2) then 1 else ((blockSize)/2) in

		  debugf ("BlockEnd: " ^ (string_of_int blockEnd));
		  let delta_angle =  angle_numerator /. (float_of_int blockSize) in
		  let  sm2 = sin ( (-2.0) *. delta_angle ) in
		  let  sm1 = sin ( (-1.0) *. delta_angle ) in
		  let  cm2 = cos ( (-2.0) *. delta_angle ) in
		  let  cm1 = cos ( (-1.0) *. delta_angle ) in
		  let  w = 2.0 *. cm1 in
		  let rec inumSamplesLoop ar0 ai0 = function
		      i when i >= numSamples -> ()
		    | i ->
		    let rec nblockEndloop ar0 ar1 ar2 ai0 ai1 ai2 j = function
			n when n >= blockEnd -> (ar0,ai0)
		      | n ->
			  let ar0' = w*.ar1 -. ar2 in
			  let ar2' = ar1 in
			  let ar1' = ar0' in
			  let ai0' = w*.ai1 -. ai2 in
			  let ai2' = ai1 in
			  let ai1' = ai0' in
			  let k = j + blockEnd in
			  let tr = ar0' *. realOut.(k) -. ai0' *. imagOut.(k) in
			  let ti = ar0' *. imagOut.(k) +. ai0' *. realOut.(k) in
			  realOut.(k) <- realOut.(j) -. tr ;
			  imagOut.(k) <- imagOut.(j) -. ti ;
			  realOut.(j) <- realOut.(j) +. tr;
			  imagOut.(j) <- imagOut.(j) +. ti;
			  nblockEndloop ar0' ar1' ar2' ai0' ai1' ai2' (j+1) (n+1)
		   in

		   let (ar0',ai0') = nblockEndloop ar0 cm1 cm2 ai0 sm1 sm2 i 0 in
		   let i' = i + blockSize in
		   inumSamplesLoop ar0' ai0' i';
		 in
		 inumSamplesLoop 0.0 0.0 0;
		 blockSizeNumSamplesLoop (2*blockSize);

		in
		blockSizeNumSamplesLoop 2;		
		
		  
		if (inverseTransform) then
		  begin
		    let denom = float_of_int numSamples in
		    for i = 0 to (numSamples - 1) do
		      realOut.(i) <- realOut.(i) /. denom;
		      imagOut.(i) <- imagOut.(i) /. denom;
		    done;
		    debugf "InverseTransform";
		  end;

 	(* convert to tail recursion *)

      method fft_f = self#fftrun false

      method fft_i = self#fftrun true


	method magnitude real imag  = sqrt(real*.real+.imag*.imag)
	
	method angle real imag = atan2 real imag
	
        method toReal magnitude angle = magnitude *. (cos angle)
	
	method toImag magnitude angle = magnitude *. (sin angle)
	
end;;



type fftdata = Empty | FFTData of float array * float array
;;
type fftbuffer = FFTBuffer of fftdata array;;
let fftbufferarr (FFTBuffer(a)) = a;;
let fftbufferget (FFTBuffer(a)) b = a.(b);;
let bufferget = fftbufferget;;

let realof (FFTData(a,b)) = a ;;
let imagof (FFTData(a,b)) = b ;;
let fftdatasize (FFTData(a,b)) = Array.length a ;;
let fftbuffersize (FFTBuffer(a)) = Array.length a;;

let createjfft nffts fftsize sr window =
  let fourier = new jfft in
  let create_data () = 
    let real = Array.create fftsize 0.0  in
    let imag = Array.create fftsize 0.0  in
    FFTData(real,imag)
  in
  let temp = create_data () in
  let create_fft_buffer () =
    let l = nffts in
    let buffer = Array.create l Empty in
    Array.iteri (fun i x -> ( buffer.(i) <- create_data() )) buffer;
    FFTBuffer(buffer)
  in
  let clear (FFTData(real,imag)) =
    Array.fill real 0 fftsize 0.0;
    Array.fill imag 0 fftsize 0.0;
  in
  let clearbuffer (FFTBuffer(a)) =
    Array.iter (fun x -> clear x; ) a
  in

  let tempb = create_fft_buffer () in
  let blitreals (FFTData(real,imag)) reals start =
    Array.blit reals start real 0 fftsize;
    Array.fill imag 0 fftsize 0.0;
  in
  let normalize arr = 
    Array.iteri (fun i x -> arr.(i) <- arr.(i)/.32767.0) arr;
  in
  let denormalize arr = 
    Array.iteri (fun i x -> arr.(i) <- arr.(i)*.32767.0) arr;
  in
  let prepare buffer sound =
      let l = nffts in
      for i = 0 to (l-1) do
       let arr = (fftbufferarr(buffer)).(i) in

       blitreals arr sound (i*fftsize);
    (*   normalize (realof arr);  *)
       window (realof arr);        
      done;
      
  in
  let fft soundin fftout =
    prepare tempb soundin;
    debugf "prepared";
    for i = 0 to (nffts-1) do
      let fo = fftbufferget fftout i in
      clear fo;

       let tb = fftbufferget tempb i in

       fourier#fft_f (realof tb) (imagof tb) (realof fo) (imagof fo);
    done;
  in
  let ffti (FFTBuffer(fftin)) soundout =
    for i = 0 to (nffts-1) do
      clear temp;
      fourier#fft_i (realof (fftin.(i))) (imagof (fftin.(i))) (realof temp) (imagof temp);
      Array.blit (realof temp) 0 soundout (i*fftsize) fftsize;
(*      denormalize soundout;  *)
    done;
  in
  (create_fft_buffer,fft,ffti)
;;

(* expects a list of fft frames *)
(* max overlap of .5 *)
(* 
  |/--\|/--\|/--\|/--\|/--\|/--\|

  2
  |/--\|/--\|/--\|/--\|/--\|/--\|
  |  /-|-\/-|-\/-|-\/-|-\/-|-\/-|
   1    1     1    1    1   1
  3
  |/-\| /-|\ /|-\ |/-\| /-|
     /|-\ |/-\| /-|\ /|-\ |
    2   2   3   2   3   2 
   1   2   1   2  1  2 
*)
(*
   Ignore 1
   2: 1 1 1 1 1 1 1 1 1 1
   3: 1 2 1 2 1 2 1 2 1 2
   4: 1 3 2 1 3 2 1 3 2 1
   5: 1 4 3 2 1 4 3 2 1 4
*)


let lowpass bin (FFTBuffer(arr)) =
  let l = Array.length arr in
  let l2 = Array.length (realof arr.(0)) in
  for i = 0 to (l - 1) do
    let real = realof arr.(i) in
    for j = bin to (l2-1) do
      let mix = (float_of_int (j - bin))/.(float_of_int(l - bin)) in
      real.(j) <- real.(j) *. mix ;
    done;
  done;
;;
let half (FFTBuffer(arr)) =
  let l = Array.length arr in
  let l2 = Array.length (realof arr.(0)) in
  for i = 0 to (l - 1) do
    let real = realof arr.(i) in
    let imag = imagof arr.(i) in
    Array.fill real (l2/2) (l2/2) 0.0;
    Array.fill imag (l2/2) (l2/2) 0.0;
  done;

;;
let fftcopy (FFTBuffer(out)) (FFTBuffer(inb)) =
  let l = Array.length out in
  let l2 = Array.length (realof out.(0)) in
  for i = 0 to (l-1) do
    Array.blit (realof inb.(i)) 0 (realof out.(i)) 0 l2;
    Array.blit (imagof inb.(i)) 0 (imagof out.(i)) 0 l2;
  done;
;;
let fftclear (FFTBuffer(out)) =
  let l = Array.length out in
  let l2 = Array.length (realof out.(0)) in
  for i = 0 to (l-1) do
    Array.fill (realof out.(i)) 0 l2 0.0;
    Array.fill (imagof out.(i)) 0 l2 0.0;
  done;

;;