{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
-- ToDo: name module consistent with FFTW.Extra.Rank1.Transformation
module Numeric.BLAS.Rank1.Fourier where

import Numeric.BLAS.Rank1.Unit (twiddleVectorAuto)

import qualified Numeric.FFTW.Rank1 as Rank1
import qualified Numeric.FFTW.Batch as Batch
import qualified Numeric.BLAS.Vector.Slice as VectorSlice
import qualified Numeric.BLAS.Vector as Vector
import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)

import Data.Complex (Complex)


{- $setup
>>> :set -XTypeFamilies
>>> import Numeric.BLAS.Rank1.Fourier
>>> import Test.Numeric.FFTW.Extra.Utility
>>> import DocTest.NumberType_.Numeric.FFTW.Extra.Rank1.Convolution (complex_)
>>> import Numeric.FFTW.Extra.Utility (padCyclic)
>>> import qualified Numeric.FFTW.Rank1 as Rank1
>>>
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>>
>>> import qualified Test.QuickCheck as QC
-}


{- |
prop> :{
   QC.forAll (genCyclic complex_) $ \x sign ->
   let n = Shape.size $ Array.shape x in
   QC.forAll (chooseLogarithmic n) $ \m ->
   let spectrumShape = Shape.ZeroBased m in

   takeFourier sign spectrumShape x
   =~=
   takeShape spectrumShape (Rank1.fourier sign x)
:}
-}
takeFourier ::
   (Shape.Indexed shape, Shape.Index shape ~ n, Integral n, Class.Real a) =>
   Rank1.Sign ->
   shape ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array shape (Complex a)
takeFourier sign shape x =
   let Shape.Cyclic n = Array.shape x in
   let cu = 2*pi / fromIntegral n in
   let c = case sign of Batch.Forward -> cu; Batch.Backward -> -cu in
   let xShape = Shape.ZeroBased n in
   Array.sample shape $ \j ->
      VectorSlice.inner
         (twiddleVectorAuto c n xShape j)
         (Array.reshape xShape x)


{- |
@extendFourier n@ computes @n@ spectral coefficients efficiently
from a shorter signal.

prop> :{
   QC.forAll genVectorShape $ \shape sign ->
   let n = Shape.zeroBasedSize shape in
   QC.forAll (chooseLogarithmic n) $ \m ->
   QC.forAll (genVectorForShape complex_ (Shape.ZeroBased m)) $ \x ->

   extendFourier sign n x
   =~=
   Rank1.fourier sign (padCyclic n x)
:}
-}
extendFourier ::
   (Shape.Indexed shape, Shape.Index shape ~ n, Integral n, Class.Real a) =>
   Rank1.Sign ->
   n ->
   Array shape (Complex a) ->
   Array (Shape.Cyclic n) (Complex a)
extendFourier sign n x =
   let cu = 2*pi / fromIntegral n in
   let c = case sign of Batch.Forward -> -cu; Batch.Backward -> cu in
   Array.reshape (Shape.Cyclic n) $
   foldl1 Vector.add $
   map (\(j,xj) ->
            VectorSlice.scale xj $
               twiddleVectorAuto c n (Shape.ZeroBased n) j) $
   Array.toAssociations x