(*
 * Copyright (c) 1997-1999 Massachusetts Institute of Technology
 * Copyright (c) 2003, 2007-11 Matteo Frigo
 * Copyright (c) 2003, 2007-11 Massachusetts Institute of Technology
 *
 * 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-1307  USA
 *
 *)


(* This is the part of the generator that actually computes the FFT
   in symbolic form *)

open Complex
open Util

(* choose a suitable factor of n *)
let choose_factor n =
  (* first choice: i such that gcd(i, n / i) = 1, i as big as possible *)
  let choose1 n =
    let rec loop i f =
      if (i * i > n) then f
      else if ((n mod i) == 0 && gcd i (n / i) == 1) then loop (i + 1) i
      else loop (i + 1) f
    in loop 1 1

  (* second choice: the biggest factor i of n, where i < sqrt(n), if any *)
  and choose2 n =
    let rec loop i f =
      if (i * i > n) then f
      else if ((n mod i) == 0) then loop (i + 1) i
      else loop (i + 1) f
    in loop 1 1

  in let i = choose1 n in
  if (i > 1) then i
  else choose2 n

let is_power_of_two n = (n > 0) && ((n - 1) land n == 0)
  
let rec dft_prime sign n input = 
  let sum filter i =
    sigma 0 n (fun j ->
      let coeff = filter (exp n (sign * i * j))
      in coeff @* (input j)) in
  let computation_even = array n (sum identity)
  and computation_odd =
    let sumr = array n (sum real)
    and sumi = array n (sum ((times Complex.i) @@ imag)) in
    array n (fun i ->
      if (i = 0) then
	(* expose some common subexpressions *)
	input 0 @+ 
	sigma 1 ((n + 1) / 2) (fun j -> input j @+ input (n - j))
      else
	let i' = min i (n - i) in
	if (i < n - i) then 
	  sumr i' @+ sumi i'
	else
	  sumr i' @- sumi i') in
  if (n >= !Magic.rader_min) then
    dft_rader sign n input
  else if (n == 2) then
    computation_even
  else
    computation_odd 


and dft_rader sign p input =
  let half = 
    let one_half = inverse_int 2 in
    times one_half

  and make_product n a b =
    let scale_factor = inverse_int n in
    array n (fun i -> a i @* (scale_factor @* b i)) in

  (* generates a convolution using ffts.  (all arguments are the
     same as to gen_convolution, below) *)
  let gen_convolution_by_fft n a b addtoall =
    let fft_a = dft 1 n a
    and fft_b = dft 1 n b in 

    let fft_ab = make_product n fft_a fft_b
    and dc_term i = if (i == 0) then addtoall else zero in

    let fft_ab1 = array n (fun i -> fft_ab i @+ dc_term i)
    and sum = fft_a 0 in
    let conv = dft (-1) n fft_ab1 in
    (sum, conv)

  (* alternate routine for convolution.  Seems to work better for
     small sizes.  I have no idea why. *)
  and gen_convolution_by_fft_alt n a b addtoall =
    let ap = array n (fun i -> half (a i @+ a ((n - i) mod n)))
    and am = array n (fun i -> half (a i @- a ((n - i) mod n)))
    and bp = array n (fun i -> half (b i @+ b ((n - i) mod n)))
    and bm = array n (fun i -> half (b i @- b ((n - i) mod n)))
    in

    let fft_ap = dft 1 n ap
    and fft_am = dft 1 n am
    and fft_bp = dft 1 n bp
    and fft_bm = dft 1 n bm in

    let fft_abpp = make_product n fft_ap fft_bp
    and fft_abpm = make_product n fft_ap fft_bm
    and fft_abmp = make_product n fft_am fft_bp
    and fft_abmm = make_product n fft_am fft_bm 
    and sum = fft_ap 0 @+ fft_am 0
    and dc_term i = if (i == 0) then addtoall else zero in

    let fft_ab1 = array n (fun i -> (fft_abpp i @+ fft_abmm i) @+ dc_term i)
    and fft_ab2 = array n (fun i -> fft_abpm i @+ fft_abmp i) in
    let conv1 = dft (-1) n fft_ab1 
    and conv2 = dft (-1) n fft_ab2 in
    let conv = array n (fun i ->
      conv1 i @+ conv2 i) in
    (sum, conv) 

    (* generator of assignment list assigning conv to the convolution of
       a and b, all of which are of length n.  addtoall is added to
       all of the elements of the result.  Returns (sum, convolution) pair
       where sum is the sum of the elements of a. *)

  in let gen_convolution = 
    if (p <= !Magic.alternate_convolution) then 
      gen_convolution_by_fft_alt
    else
      gen_convolution_by_fft

  (* fft generator for prime n = p using Rader's algorithm for
     turning the fft into a convolution, which then can be
     performed in a variety of ways *)
  in  
    let g = find_generator p in
    let ginv = pow_mod g (p - 2) p in
    let input_perm = array p (fun i -> input (pow_mod g i p))
    and omega_perm = array p (fun i -> exp p (sign * (pow_mod ginv i p)))
    and output_perm = array p (fun i -> pow_mod ginv i p)
    in let (sum, conv) = 
      (gen_convolution (p - 1)  input_perm omega_perm (input 0))
    in array p (fun i ->
      if (i = 0) then
	input 0 @+ sum
      else
	let i' = suchthat 0 (fun i' -> i = output_perm i')
	in conv i')

(* our modified version of the conjugate-pair split-radix algorithm,
   which reduces the number of multiplications by rescaling the 
   sub-transforms (power-of-two n's only) *)
and newsplit sign n input =
  let rec s n k = (* recursive scale factor *)
    if n <= 4 then
      one
    else 
      let k4 = (abs k) mod (n / 4) in
      let k4' = if k4 <= (n / 8) then k4 else (n/4 - k4) in
      (s (n / 4) k4') @* (real (exp n k4'))
			  
  and sinv n k = (* 1 / s(n,k) *)
    if n <= 4 then
      one
    else 
      let k4 = (abs k) mod (n / 4) in
      let k4' = if k4 <= (n / 8) then k4 else (n/4 - k4) in
      (sinv (n / 4) k4') @* (sec n k4')

  in let sdiv2 n k = (s n k) @* (sinv (2*n) k) (* s(n,k) / s(2*n,k) *)
  and sdiv4 n k = (* s(n,k) / s(4*n,k) *)
    let k4 = (abs k) mod n in
    sec (4*n) (if k4 <= (n / 2) then k4 else (n - k4))
      
  in let t n k = (exp n k) @* (sdiv4 (n/4) k)

  and dft1 input = input
  and dft2 input = array 2 (fun k -> (input 0) @+ ((input 1) @* exp 2 k))

  in let rec newsplit0 sign n input =
    if (n == 1) then dft1 input
    else if (n == 2) then dft2 input
    else let u = newsplit0 sign (n / 2) (fun i -> input (i*2))
    and z = newsplitS sign (n / 4) (fun i -> input (i*4 + 1))
    and z' = newsplitS sign (n / 4) (fun i -> input ((n + i*4 - 1) mod n)) 
    and twid = array n (fun k -> s (n/4) k @* exp n (sign * k)) in
    let w = array n (fun k -> twid k @* z (k mod (n / 4)))
    and w' = array n (fun k -> conj (twid k) @* z' (k mod (n / 4))) in
    let ww = array n (fun k -> w k @+ w' k) in
    array n (fun k -> u (k mod (n / 2)) @+ ww k)
      
  and newsplitS sign n input =
    if (n == 1) then dft1 input
    else if (n == 2) then dft2 input
    else let u = newsplitS2 sign (n / 2) (fun i -> input (i*2))
    and z = newsplitS sign (n / 4) (fun i -> input (i*4 + 1))
    and z' = newsplitS sign (n / 4) (fun i -> input ((n + i*4 - 1) mod n)) in
    let w = array n (fun k -> t n (sign * k) @* z (k mod (n / 4)))
    and w' = array n (fun k -> conj (t n (sign * k)) @* z' (k mod (n / 4))) in
    let ww = array n (fun k -> w k @+ w' k) in
    array n (fun k -> u (k mod (n / 2)) @+ ww k)
      
  and newsplitS2 sign n input =
    if (n == 1) then dft1 input
    else if (n == 2) then dft2 input
    else let u = newsplitS4 sign (n / 2) (fun i -> input (i*2))
    and z = newsplitS sign (n / 4) (fun i -> input (i*4 + 1))
    and z' = newsplitS sign (n / 4) (fun i -> input ((n + i*4 - 1) mod n)) in
    let w = array n (fun k -> t n (sign * k) @* z (k mod (n / 4)))
    and w' = array n (fun k -> conj (t n (sign * k)) @* z' (k mod (n / 4))) in
    let ww = array n (fun k -> (w k @+ w' k) @* (sdiv2 n k)) in
    array n (fun k -> u (k mod (n / 2)) @+ ww k)
      
  and newsplitS4 sign n input =
    if (n == 1) then dft1 input
    else if (n == 2) then 
      let f = dft2 input
      in array 2 (fun k -> (f k) @* (sinv 8 k))
    else let u = newsplitS2 sign (n / 2) (fun i -> input (i*2))
    and z = newsplitS sign (n / 4) (fun i -> input (i*4 + 1))
    and z' = newsplitS sign (n / 4) (fun i -> input ((n + i*4 - 1) mod n)) in
    let w = array n (fun k -> t n (sign * k) @* z (k mod (n / 4)))
    and w' = array n (fun k -> conj (t n (sign * k)) @* z' (k mod (n / 4))) in
    let ww = array n (fun k -> w k @+ w' k) in
    array n (fun k -> (u (k mod (n / 2)) @+ ww k) @* (sdiv4 n k))
      
  in newsplit0 sign n input
 
and dft sign n input =
  let rec cooley_tukey sign n1 n2 input =
    let tmp1 = 
      array n2 (fun i2 -> 
	dft sign n1 (fun i1 -> input (i1 * n2 + i2))) in
    let tmp2 =  
      array n1 (fun i1 ->
	array n2 (fun i2 ->
	  exp n (sign * i1 * i2) @* tmp1 i2 i1)) in
    let tmp3 = array n1 (fun i1 -> dft sign n2 (tmp2 i1)) in
    (fun i -> tmp3 (i mod n1) (i / n1))

  (*
   * This is "exponent -1" split-radix by Dan Bernstein.
   *)
  and split_radix_dit sign n input =
    let f0 = dft sign (n / 2) (fun i -> input (i * 2))
    and f10 = dft sign (n / 4) (fun i -> input (i * 4 + 1))
    and f11 = dft sign (n / 4) (fun i -> input ((n + i * 4 - 1) mod n)) in
    let g10 = array n (fun k ->
      exp n (sign * k) @* f10 (k mod (n / 4)))
    and g11 = array n (fun k ->
      exp n (- sign * k) @* f11 (k mod (n / 4))) in
    let g1 = array n (fun k -> g10 k @+ g11 k) in
    array n (fun k -> f0 (k mod (n / 2)) @+ g1 k)

  and split_radix_dif sign n input =
    let n2 = n / 2 and n4 = n / 4 in
    let x0 = array n2 (fun i -> input i @+ input (i + n2))
    and x10 = array n4 (fun i -> input i @- input (i + n2))
    and x11 = array n4 (fun i ->
	input (i + n4) @- input (i + n2 + n4)) in
    let x1 k i = 
      exp n (k * i * sign) @* (x10 i @+ exp 4 (k * sign) @* x11 i) in
    let f0 = dft sign n2 x0 
    and f1 = array 4 (fun k -> dft sign n4 (x1 k)) in
    array n (fun k ->
      if k mod 2 = 0 then f0 (k / 2)
      else let k' = k mod 4 in f1 k' ((k - k') / 4))

  and prime_factor sign n1 n2 input =
    let tmp1 = array n2 (fun i2 ->
      dft sign n1 (fun i1 -> input ((i1 * n2 + i2 * n1) mod n)))
    in let tmp2 = array n1 (fun i1 ->
      dft sign n2 (fun k2 -> tmp1 k2 i1))
    in fun i -> tmp2 (i mod n1) (i mod n2)

  in let algorithm sign n =
    let r = choose_factor n in
    if List.mem n !Magic.rader_list then
      (* special cases *)
      dft_rader sign n
    else if (r == 1) then  (* n is prime *)
      dft_prime sign n
    else if (gcd r (n / r)) == 1 then
      prime_factor sign r (n / r)
    else if (n mod 4 = 0 && n > 4) then
      if !Magic.newsplit && is_power_of_two n then
	newsplit sign n
      else if !Magic.dif_split_radix then
	split_radix_dif sign n
      else
	split_radix_dit sign n
    else 
      cooley_tukey sign r (n / r)
  in
  array n (algorithm sign n input)