(*  Title:      Fast Fourier Transform

     Author:     Clemens Ballarin <ballarin at in.tum.de>, started 12 April 2005

     Maintainer: Clemens Ballarin <ballarin at in.tum.de>

 *)

 theory FFT

 imports Complex_Main

 begin

 text {* We formalise a functional implementation of the FFT algorithm

   over the complex numbers, and its inverse.  Both are shown

   equivalent to the usual definitions

   of these operations through Vandermonde matrices.  They are also

   shown to be inverse to each other, more precisely, that composition

   of the inverse and the transformation yield the identity up to a

   scalar.

   The presentation closely follows Section 30.2 of Cormen \textit{et

   al.}, \emph{Introduction to Algorithms}, 2nd edition, MIT Press,

   2003. *}

 section {* Preliminaries *}

 lemma of_nat_cplx:

   "of_nat n = Complex (of_nat n) 0"

   by (induct n) (simp_all add: complex_one_def)

 text {* The following two lemmas are useful for experimenting with the

   transformations, at a vector length of four. *}

 lemma Ivl4:

   "{0..<4::nat} = {0, 1, 2, 3}"

 proof -

   have "{0..<4::nat} = {0..<Suc (Suc (Suc (Suc 0)))}" by (simp add: eval_nat_numeral)

   also have "... = {0, 1, 2, 3}"

     by (simp add: atLeastLessThanSuc eval_nat_numeral insert_commute)

   finally show ?thesis .

 qed

 lemma Sum4:

   "(\<Sum>i=0..<4::nat. x i) = x 0 + x 1 + x 2 + x 3"

   by (simp add: Ivl4 eval_nat_numeral)

 text {* A number of specialised lemmas for the summation operator,

   where the index set is the natural numbers *}

 lemma setsum_add_nat_ivl_singleton:

   assumes less: "m < (n::nat)"

   shows "f m + setsum f {m<..<n} = setsum f {m..<n}"

 proof -

   have "f m + setsum f {m<..<n} = setsum f ({m} \<union> {m<..<n})"

     by (simp add: setsum_Un_disjoint ivl_disj_int)

   also from less have "... = setsum f {m..<n}"

     by (simp only: ivl_disj_un)

   finally show ?thesis .

 qed

 lemma setsum_add_split_nat_ivl_singleton:

   assumes less: "m < (n::nat)"

     and g: "!!i. [| m < i; i < n |] ==> g i = f i"

   shows "f m + setsum g {m<..<n} = setsum f {m..<n}"

   using less g

   by(simp add: setsum_add_nat_ivl_singleton cong: strong_setsum_cong)

 lemma setsum_add_split_nat_ivl:

   assumes le: "m <= (k::nat)" "k <= n"

     and g: "!!i. [| m <= i; i < k |] ==> g i = f i"

     and h: "!!i. [| k <= i; i < n |] ==> h i = f i"

   shows "setsum g {m..<k} + setsum h {k..<n} = setsum f {m..<n}"

   using le g h by (simp add: setsum_add_nat_ivl cong: strong_setsum_cong)

 lemma ivl_splice_Un:

   "{0..<2*n::nat} = (op * 2 ` {0..<n}) \<union> ((%i. Suc (2*i)) ` {0..<n})"

   apply (unfold image_def Bex_def)

   apply auto

   apply arith

   done

 lemma ivl_splice_Int:

   "(op * 2 ` {0..<n}) \<inter> ((%i. Suc (2*i)) ` {0..<n}) = {}"

   by auto arith

 lemma double_inj_on:

   "inj_on (%i. 2*i::nat) A"

   by (simp add: inj_onI)

 lemma Suc_double_inj_on:

   "inj_on (%i. Suc (2*i)) A"

   by (rule inj_onI) simp

 lemma setsum_splice:

   "(\<Sum>i::nat = 0..<2*n. f i) = (\<Sum>i = 0..<n. f (2*i)) + (\<Sum>i = 0..<n. f (2*i+1))"

 proof -

   have "(\<Sum>i::nat = 0..<2*n. f i) =

     setsum f (op * 2 ` {0..<n}) + setsum f ((%i. 2*i+1) ` {0..<n})"

     by (simp add: ivl_splice_Un ivl_splice_Int setsum_Un_disjoint)

   also have "... = (\<Sum>i = 0..<n. f (2*i)) + (\<Sum>i = 0..<n. f (2*i+1))"

     by (simp add: setsum_reindex [OF double_inj_on]

       setsum_reindex [OF Suc_double_inj_on])

   finally show ?thesis .

 qed

 section {* Complex Roots of Unity *}

 text {* The function @{term cis} from the complex library returns the

   point on the unity circle corresponding to the argument angle.  It

   is the base for our definition of @{text root}.  The main property,

   De Moirve's formula is already there in the library. *}

 definition root :: "nat => complex" where

   "root n == cis (2*pi/(real (n::nat)))"

 lemma sin_periodic_pi_diff [simp]: "sin (x - pi) = - sin x"

   by (simp add: sin_diff)

 lemma sin_cos_between_zero_two_pi:

   assumes 0: "0 < x" and pi: "x < 2 * pi"

   shows "sin x \<noteq> 0 \<or> cos x \<noteq> 1"

 proof -

   { assume "0 < x" and "x < pi"

     then have "sin x \<noteq> 0" by (auto dest: sin_gt_zero_pi) }

   moreover

   { assume "x = pi"

     then have "cos x \<noteq> 1" by simp }

   moreover

   { assume pi1: "pi < x" and pi2: "x < 2 * pi"

     then have "0 < x - pi" and "x - pi < pi" by arith+

     then have "sin (x - pi) \<noteq> 0" by (auto dest: sin_gt_zero_pi)

     with pi1 pi2 have "sin x \<noteq> 0" by simp }

   ultimately show ?thesis using 0 pi by arith

 qed

 subsection {* Basic Lemmas *}

 lemma root_nonzero:

   "root n ~= 0"

   apply (unfold root_def)

   apply (unfold cis_def)

   apply auto

   apply (drule sin_zero_abs_cos_one)

   apply arith

   done

 lemma root_unity:

   "root n ^ n = 1"

   apply (unfold root_def)

   apply (simp add: DeMoivre)

   apply (simp add: cis_def)

   done

 lemma root_cancel:

   "0 < d ==> root (d * n) ^ (d * k) = root n ^ k"

   apply (unfold root_def)

   apply (simp add: DeMoivre)

   done

 lemma root_summation:

   assumes k: "0 < k" "k < n"

   shows "(\<Sum>i=0..<n. (root n ^ k) ^ i) = 0"

 proof -

   from k have real0: "0 < real k * (2 * pi) / real n"

     by (simp add: zero_less_divide_iff

       mult_strict_right_mono [where a = 0, simplified])

   from k mult_strict_right_mono [where a = "real k" and

     b = "real n" and c = "2 * pi / real n", simplified]

   have realk: "real k * (2 * pi) / real n < 2 * pi"

     by (simp add: zero_less_divide_iff)

   txt {* Main part of the proof *}

   have "(\<Sum>i=0..<n. (root n ^ k) ^ i) =

     ((root n ^ k) ^ n - 1) / (root n ^ k - 1)"

     apply (rule geometric_sum)

     apply (unfold root_def)

     apply (simp add: DeMoivre)

     using real0 realk sin_cos_between_zero_two_pi 

     apply (auto simp add: cis_def complex_one_def)

     done

   also have "... = ((root n ^ n) ^ k - 1) / (root n ^ k - 1)"

     by (simp add: power_mult [THEN sym] mult_ac)

   also have "... = 0"

     by (simp add: root_unity)

   finally show ?thesis .

 qed

 lemma root_summation_inv:

   assumes k: "0 < k" "k < n"

   shows "(\<Sum>i=0..<n. ((1 / root n) ^ k) ^ i) = 0"

 proof -

   from k have real0: "0 < real k * (2 * pi) / real n"

     by (simp add: zero_less_divide_iff

       mult_strict_right_mono [where a = 0, simplified])

   from k mult_strict_right_mono [where a = "real k" and

     b = "real n" and c = "2 * pi / real n", simplified]

   have realk: "real k * (2 * pi) / real n < 2 * pi"

     by (simp add: zero_less_divide_iff)

   txt {* Main part of the proof *}

   have "(\<Sum>i=0..<n. ((1 / root n) ^ k) ^ i) =

     (((1 / root n) ^ k) ^ n - 1) / ((1 / root n) ^ k - 1)"

     apply (rule geometric_sum)

     apply (simp add: nonzero_inverse_eq_divide [THEN sym] root_nonzero)

     apply (unfold root_def)

     apply (simp add: DeMoivre)

     using real0 realk sin_cos_between_zero_two_pi

     apply (auto simp add: cis_def complex_one_def)

     done

   also have "... = (((1 / root n) ^ n) ^ k - 1) / ((1 / root n) ^ k - 1)"

     by (simp add: power_mult [THEN sym] mult_ac)

   also have "... = 0"

     by (simp add: power_divide root_unity)

   finally show ?thesis .

 qed

 lemma root0 [simp]:

   "root 0 = 1"

   by (simp add: root_def cis_def)

 lemma root1 [simp]:

   "root 1 = 1"

   by (simp add: root_def cis_def)

 lemma root2 [simp]:

   "root 2 = Complex -1 0"

   by (simp add: root_def cis_def)

 lemma root4 [simp]:

   "root 4 = ii"

   by (simp add: root_def cis_def)

 subsection {* Derived Lemmas *}

 lemma root_cancel1:

   "root (2 * m) ^ (i * (2 * j)) = root m ^ (i * j)"

 proof -

   have "root (2 * m) ^ (i * (2 * j)) = root (2 * m) ^ (2 * (i * j))"

     by (simp add: mult_ac)

   also have "... = root m ^ (i * j)"

     by (simp add: root_cancel)

   finally show ?thesis .

 qed

 lemma root_cancel2:

   "0 < n ==> root (2 * n) ^ n = - 1"

   txt {* Note the space between @{text "-"} and @{text "1"}. *}

   using root_cancel [where n = 2 and k = 1]

   apply (simp only: mult_ac)

   apply (simp add: complex_one_def)

   done

 section {* Discrete Fourier Transformation *}

 text {*

   We define operations  @{text DFT} and @{text IDFT} for the discrete

   Fourier Transform and its inverse.  Vectors are simply functions of

   type @{text "nat => complex"}. *}

 text {*

   @{text "DFT n a"} is the transform of vector @{text a}

   of length @{text n}, @{text IDFT} its inverse. *}

 definition DFT :: "nat => (nat => complex) => (nat => complex)" where

   "DFT n a == (%i. \<Sum>j=0..<n. (root n) ^ (i * j) * (a j))"

 definition IDFT :: "nat => (nat => complex) => (nat => complex)" where

   "IDFT n a == (%i. (\<Sum>k=0..<n. (a k) / (root n) ^ (i * k)))"

 schematic_lemma "map (DFT 4 a) [0, 1, 2, 3] = ?x"

   by(simp add: DFT_def Sum4)

 text {* Lemmas for the correctness proof. *}

 lemma DFT_lower:

   "DFT (2 * m) a i =

   DFT m (%i. a (2 * i)) i +

   (root (2 * m)) ^ i * DFT m (%i. a (2 * i + 1)) i"

 proof (unfold DFT_def)

   have "(\<Sum>j = 0..<2 * m. root (2 * m) ^ (i * j) * a j) =

     (\<Sum>j = 0..<m. root (2 * m) ^ (i * (2 * j)) * a (2 * j)) +

     (\<Sum>j = 0..<m. root (2 * m) ^ (i * (2 * j + 1)) * a (2 * j + 1))"

     (is "?s = _")

     by (simp add: setsum_splice)

   also have "... = (\<Sum>j = 0..<m. root m ^ (i * j) * a (2 * j)) +

     root (2 * m) ^ i *

     (\<Sum>j = 0..<m. root m ^ (i * j) * a (2 * j + 1))"

     (is "_ = ?t")

     txt {* First pair of sums *}

     apply (simp add: root_cancel1)

     txt {* Second pair of sums *}

     apply (simp add: setsum_right_distrib)

     apply (simp add: power_add)

     apply (simp add: root_cancel1)

     apply (simp add: mult_ac)

     done

   finally show "?s = ?t" .

 qed

 lemma DFT_upper:

   assumes mbound: "0 < m" and ibound: "m <= i"

   shows "DFT (2 * m) a i =

     DFT m (%i. a (2 * i)) (i - m) -

     root (2 * m) ^ (i - m) * DFT m (%i. a (2 * i + 1)) (i - m)"

 proof (unfold DFT_def)

   have "(\<Sum>j = 0..<2 * m. root (2 * m) ^ (i * j) * a j) =

     (\<Sum>j = 0..<m. root (2 * m) ^ (i * (2 * j)) * a (2 * j)) +

     (\<Sum>j = 0..<m. root (2 * m) ^ (i * (2 * j + 1)) * a (2 * j + 1))"

     (is "?s = _")

     by (simp add: setsum_splice)

   also have "... =

     (\<Sum>j = 0..<m. root m ^ ((i - m) * j) * a (2 * j)) -

     root (2 * m) ^ (i - m) *

     (\<Sum>j = 0..<m. root m ^ ((i - m) * j) * a (2 * j + 1))"

     (is "_ = ?t")

     txt {* First pair of sums *}

     apply (simp add: root_cancel1)

     apply (simp add: root_unity ibound root_nonzero power_diff power_mult)

     txt {* Second pair of sums *}

     apply (simp add: mbound root_cancel2)

     apply (simp add: setsum_right_distrib)

     apply (simp add: power_add)

     apply (simp add: root_cancel1)

     apply (simp add: power_mult)

     apply (simp add: mult_ac)

     done

   finally show "?s = ?t" .

 qed

 lemma IDFT_lower:

   "IDFT (2 * m) a i =

   IDFT m (%i. a (2 * i)) i +

   (1 / root (2 * m)) ^ i * IDFT m (%i. a (2 * i + 1)) i"

 proof (unfold IDFT_def)

   have "(\<Sum>j = 0..<2 * m. a j / root (2 * m) ^ (i * j)) =

     (\<Sum>j = 0..<m. a (2 * j) / root (2 * m) ^ (i * (2 * j))) +

     (\<Sum>j = 0..<m. a (2 * j + 1) / root (2 * m) ^ (i * (2 * j + 1)))"

     (is "?s = _")

     by (simp add: setsum_splice)

   also have "... = (\<Sum>j = 0..<m. a (2 * j) / root m ^ (i * j)) +

     (1 / root (2 * m)) ^ i *

     (\<Sum>j = 0..<m. a (2 * j + 1) / root m ^ (i * j))"

     (is "_ = ?t")

     txt {* First pair of sums *}

     apply (simp add: root_cancel1)

     txt {* Second pair of sums *}

     apply (simp add: setsum_right_distrib)

     apply (simp add: power_add)

     apply (simp add: nonzero_power_divide root_nonzero)

     apply (simp add: root_cancel1)

     done

   finally show "?s = ?t" .

 qed

 lemma IDFT_upper:

   assumes mbound: "0 < m" and ibound: "m <= i"

   shows "IDFT (2 * m) a i =

     IDFT m (%i. a (2 * i)) (i - m) -

     (1 / root (2 * m)) ^ (i - m) *

     IDFT m (%i. a (2 * i + 1)) (i - m)"

 proof (unfold IDFT_def)

   have "(\<Sum>j = 0..<2 * m. a j / root (2 * m) ^ (i * j)) =

     (\<Sum>j = 0..<m. a (2 * j) / root (2 * m) ^ (i * (2 * j))) +

     (\<Sum>j = 0..<m. a (2 * j + 1) / root (2 * m) ^ (i * (2 * j + 1)))"

     (is "?s = _")

     by (simp add: setsum_splice)

   also have "... =

     (\<Sum>j = 0..<m. a (2 * j) / root m ^ ((i - m) * j)) -

     (1 / root (2 * m)) ^ (i - m) *

     (\<Sum>j = 0..<m. a (2 * j + 1) / root m ^ ((i - m) * j))"

     (is "_ = ?t")

     txt {* First pair of sums *}

     apply (simp add: root_cancel1)

     apply (simp add: root_unity ibound root_nonzero power_diff power_mult)

     txt {* Second pair of sums *}

     apply (simp add: nonzero_power_divide root_nonzero)

     apply (simp add: mbound root_cancel2)

     apply (simp add: setsum_divide_distrib)

     apply (simp add: power_add)

     apply (simp add: root_cancel1)

     apply (simp add: power_mult)

     apply (simp add: mult_ac)

     done

   finally show "?s = ?t" .

 qed

 text {* @{text DFT} und @{text IDFT} are inverses. *}

 declare divide_divide_eq_right [simp del]

   divide_divide_eq_left [simp del]

 lemma power_diff_inverse:

   assumes nz: "(a::'a::field) ~= 0"

   shows "m <= n ==> (inverse a) ^ (n-m) = (a^m) / (a^n)"

   apply (induct n m rule: diff_induct)

     apply (simp add: nonzero_power_inverse

       nonzero_inverse_eq_divide [THEN sym] nz)

    apply simp

   apply (simp add: nz)

   done

 lemma power_diff_rev_if:

   assumes nz: "(a::'a::field) ~= 0"

   shows "(a^m) / (a^n) = (if n <= m then a ^ (m-n) else (1/a) ^ (n-m))"

 proof (cases "n <= m")

   case True with nz show ?thesis

     by (simp add: power_diff)

 next

   case False with nz show ?thesis

     by (simp add: power_diff_inverse nonzero_inverse_eq_divide [THEN sym])

 qed

 lemma power_divides_special:

   "(a::'a::field) ~= 0 ==>

   a ^ (i * j) / a ^ (k * i) = (a ^ j / a ^ k) ^ i"

   by (simp add: nonzero_power_divide power_mult [THEN sym] mult_ac)

 theorem DFT_inverse:

   assumes i_less: "i < n"

   shows  "IDFT n (DFT n a) i = of_nat n * a i"

   using [[linarith_split_limit = 0]]

   apply (unfold DFT_def IDFT_def)

   apply (simp add: setsum_divide_distrib)

   apply (subst setsum_commute)

   apply (simp only: times_divide_eq_left [THEN sym])

   apply (simp only: power_divides_special [OF root_nonzero])

   apply (simp add: power_diff_rev_if root_nonzero)

   apply (simp add: setsum_divide_distrib [THEN sym]

     setsum_left_distrib [THEN sym])

   proof -

     from i_less have i_diff: "!!k. i - k < n" by arith

     have diff_i: "!!k. k < n ==> k - i < n" by arith

     let ?sum = "%i j n. setsum (op ^ (if i <= j then root n ^ (j - i)

                   else (1 / root n) ^ (i - j))) {0..<n} * a j"

     let ?sum1 = "%i j n. setsum (op ^ (root n ^ (j - i))) {0..<n} * a j"

     let ?sum2 = "%i j n. setsum (op ^ ((1 / root n) ^ (i - j))) {0..<n} * a j"

     from i_less have "(\<Sum>j = 0..<n. ?sum i j n) =

       (\<Sum>j = 0..<i. ?sum2 i j n) + (\<Sum>j = i..<n. ?sum1 i j n)"

       (is "?s = _")

       by (simp add: root_summation_inv nat_dvd_not_less

         setsum_add_split_nat_ivl [where f = "%j. ?sum i j n"])

     also from i_less i_diff

     have "... = (\<Sum>j = i..<n. ?sum1 i j n)"

       by (simp add: root_summation_inv nat_dvd_not_less)

     also from i_less have "... =

       (\<Sum>j\<in>{i} \<union> {i<..<n}. ?sum1 i j n)"

       by (simp only: ivl_disj_un)

     also have "... =

       (?sum1 i i n + (\<Sum>j\<in>{i<..<n}. ?sum1 i j n))"

       by (simp add: setsum_Un_disjoint ivl_disj_int)

     also from i_less diff_i have "... = ?sum1 i i n"

       by (simp add: root_summation nat_dvd_not_less)

     also from i_less have "... = of_nat n * a i" (is "_ = ?t")

       by (simp add: of_nat_cplx)

     finally show "?s = ?t" .

   qed

 section {* Discrete, Fast Fourier Transformation *}

 text {* @{text "FFT k a"} is the transform of vector @{text a}

   of length @{text "2 ^ k"}, @{text IFFT} its inverse. *}

 primrec FFT :: "nat => (nat => complex) => (nat => complex)" where

   "FFT 0 a = a"

 | "FFT (Suc k) a =

      (let (x, y) = (FFT k (%i. a (2*i)), FFT k (%i. a (2*i+1)))

       in (%i. if i < 2^k

             then x i + (root (2 ^ (Suc k))) ^ i * y i

             else x (i- 2^k) - (root (2 ^ (Suc k))) ^ (i- 2^k) * y (i- 2^k)))"

 primrec IFFT :: "nat => (nat => complex) => (nat => complex)" where

   "IFFT 0 a = a"

 | "IFFT (Suc k) a =

      (let (x, y) = (IFFT k (%i. a (2*i)), IFFT k (%i. a (2*i+1)))

       in (%i. if i < 2^k

             then x i + (1 / root (2 ^ (Suc k))) ^ i * y i

             else x (i - 2^k) -

               (1 / root (2 ^ (Suc k))) ^ (i - 2^k) * y (i - 2^k)))"

 text {* Finally, for vectors of length @{text "2 ^ k"},

   @{text DFT} and @{text FFT}, and @{text IDFT} and

   @{text IFFT} are equivalent. *}

 theorem DFT_FFT:

   "!!a i. i < 2 ^ k ==> DFT (2 ^ k) a i = FFT k a i"

 proof (induct k)

   case 0

   then show ?case by (simp add: DFT_def)

 next

   case (Suc k)

   assume i: "i < 2 ^ Suc k"

   show ?case proof (cases "i < 2 ^ k")

     case True

     then show ?thesis apply simp apply (simp add: DFT_lower)

       apply (simp add: Suc) done

   next

     case False

     from i have "i - 2 ^ k < 2 ^ k" by simp

     with False i show ?thesis apply simp apply (simp add: DFT_upper)

       apply (simp add: Suc) done

   qed

 qed

 theorem IDFT_IFFT:

   "!!a i. i < 2 ^ k ==> IDFT (2 ^ k) a i = IFFT k a i"

 proof (induct k)

   case 0

   then show ?case by (simp add: IDFT_def)

 next

   case (Suc k)

   assume i: "i < 2 ^ Suc k"

   show ?case proof (cases "i < 2 ^ k")

     case True

     then show ?thesis apply simp apply (simp add: IDFT_lower)

       apply (simp add: Suc) done

   next

     case False

     from i have "i - 2 ^ k < 2 ^ k" by simp

     with False i show ?thesis apply simp apply (simp add: IDFT_upper)

       apply (simp add: Suc) done

   qed

 qed

 schematic_lemma "map (FFT (Suc (Suc 0)) a) [0, 1, 2, 3] = ?x"

   by simp

 end