{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Numeric.FFTW.Extra.Rank1.Transform (
   fourierRC2,
   spectrumIntervalFromHalf,

   takeDivisorFourier,
   takeDivisorFourierRC,
   takeDivisorFourierInterval,
   takeNonDivisorFourier,
   takeNonDivisorFourierRC,
   takeNonDivisorFourierInterval,

   extendDivisorFourier,
   extendDivisorFourierCR,
   extendDivisorFourierInterval,
   extendNonDivisorFourier,
   extendNonDivisorFourierCR,
   extendNonDivisorFourierInterval,

   shortTimeFourierForward,
   shortTimeFourierBackward,
   shortTimeFourierRC,
   shortTimeFourierCR,

   fromOverlappingBlocks,
   ) where

import Numeric.BLAS.Rank1.Unit
         (twiddleMatrix, twiddleMatrixRowsAuto,
          nonNegChirp, appendRevToInterval1,
          interval2FromChirp, interval3FromChirp)
import Numeric.FFTW.Extra.Rank1.Convolution
import Numeric.FFTW.Extra.Utility

import qualified Numeric.FFTW.Rank1 as Rank1
import qualified Numeric.FFTW.Batch as Batch
import qualified Numeric.FFTW.Shape as Spectrum
import qualified Numeric.BLAS.Matrix.RowMajor.Block as BlockMatrix
import qualified Numeric.BLAS.Matrix.RowMajor as Matrix
import qualified Numeric.BLAS.Vector.Chunk as Chunk
import qualified Numeric.BLAS.Vector.Symbolic as VectorSymb
import qualified Numeric.BLAS.Vector.Slice as VectorSlice
import qualified Numeric.BLAS.Vector as Vector
import qualified Numeric.BLAS.Subobject.Shape as Subshape
import qualified Numeric.BLAS.Subobject.View.Matrix as ViewMatrix
import qualified Numeric.BLAS.Subobject.Layout as Layout
import qualified Numeric.BLAS.Slice as Slice
import qualified Numeric.Netlib.Class as Class
import Numeric.BLAS.Matrix.RowMajor.Block ((&|||))
import Numeric.BLAS.Vector.Slice ((|+|), (|-|))

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array, (!))
import Data.Array.Comfort.Shape ((::+)((::+)))

-- import qualified Data.Foldable as Fold
-- import qualified Data.StorableVector.Lazy as SVL
-- import qualified Data.StorableVector as SV
import qualified Data.List.HT as ListHT
import qualified Data.List as List
import qualified Data.Complex as Complex
-- import Data.Traversable (Traversable)
-- import Data.Tuple.HT (mapFst, mapSnd)
-- import Data.Maybe.HT (toMaybe)
import Data.Bool.HT (if')
import Data.Complex (Complex((:+)))
import Data.Bits (Bits)


import GHC.Stack (HasCallStack)
-- import Debug.Trace (trace)


-- ToDo: define shapeInt

{- $setup
>>> :set -XTypeFamilies
>>> import Test.Numeric.FFTW.Extra.Utility
>>> import DocTest.NumberType_.Numeric.FFTW.Extra.Rank1.Convolution
>>>    (complex_, real_)
>>> import Numeric.FFTW.Extra.Rank1.Transform
>>> import Numeric.FFTW.Extra.Rank1.Convolution (arrangeMutable)
>>> import Numeric.FFTW.Extra.NumberTheory (divisors)
>>> import Numeric.FFTW.Extra.Utility (cyclicTake, padCyclic, spreadCyclic)
>>> import qualified Numeric.FFTW.Shape as Spectrum
>>> import qualified Numeric.FFTW.Rank1 as Rank1
>>> import qualified Numeric.BLAS.Vector as Vector
>>>
>>> -- import qualified Data.Array.Comfort.Storable.Dim2 as Array2
>>> import qualified Data.Array.Comfort.Storable as Array
>>> import qualified Data.Array.Comfort.Shape as Shape
>>> -- import qualified Data.Foldable as Fold
>>> -- import qualified Data.List.Match as Match
>>> import Data.Array.Comfort.Storable ((!), (//))
>>> import Data.Complex (Complex((:+)), cis)
>>>
>>> import qualified Test.QuickCheck as QC
>>> import Test.QuickCheck ((.&&.))
-}



{- |
Perform two Fourier transforms on two equal sized real vectors
for the price of one complex Fourier transform. Almost.

prop> :{
   QC.forAll (genCyclic complex_) $ \xy ->
   let (x,y) = Vector.unzipComplex xy in
   let (fx,fy) = fourierRC2 (x,y) in
   Rank1.fourierRC x =~= fx
   .&&.
   Rank1.fourierRC y =~= fy
:}
-}
fourierRC2 ::
   (Integral n, Bits n, Class.Real a) =>
   (Array (Shape.Cyclic n) a, Array (Shape.Cyclic n) a) ->
   (Array (Spectrum.Half n) (Complex a), Array (Spectrum.Half n) (Complex a))
fourierRC2 (x,y) =
   let Shape.Cyclic n = Array.shape x in
   let spectrum = Rank1.fourier Rank1.Forward $ Vector.zipComplex x y in
   let takeHalf = takeShape (Spectrum.Half n) in
   let a = takeHalf spectrum in
   let b = VectorSlice.conjugate $ takeHalf (Vector.cyclicReverse spectrum) in
   {-
   a = fx + i*fy
   b = conj fx + i * conj fy
   conj b = fx - i*fy
   solution:
   fx = (a + conj b) / 2
   fy = (a - conj b) / 2i
   -}
   (Vector.scaleReal (1/2) (a |+| b),
    Vector.scale (0 :+ (-1/2)) (a |-| b))


shapeCyclicBlock :: (Integral n) =>
   n -> Shape.Cyclic n -> (Shape.Cyclic n, Shape.ZeroBased n)
shapeCyclicBlock m (Shape.Cyclic n) =
   (Shape.Cyclic m, Shape.ZeroBased $ div n m)



applyTwiddleMatrixSimple ::
   (Shape.Indexed height, Shape.Index height ~ n, Eq height,
    Shape.Indexed width, Shape.Index width ~ n, Eq width,
    Integral n, Class.Real a) =>
   a -> Array (height, width) (Complex a) -> Array height (Complex a)
applyTwiddleMatrixSimple c a =
   Matrix.innerRowwise (twiddleMatrix c (Array.shape a)) a

applyTwiddleMatrix ::
   (Shape.IntegralRange height, Shape.Index height ~ n, Eq height,
    Shape.IntegralRange width, Shape.Index width ~ n, Eq width,
    Integral n, Class.Real a) =>
   a -> Array (height, width) (Complex a) -> Array height (Complex a)
applyTwiddleMatrix c a =
   Array.fromList (Matrix.height a) $
   zipWith VectorSlice.inner
      (twiddleMatrixRowsAuto c (Array.shape a))
      (VectorSlice.slicesVector Slice.rows a)

_extrudeTwiddleMatrixSimple ::
   (Shape.Indexed width, Shape.Index width ~ n, Eq width,
    Integral n, Class.Real a) =>
   a -> n ->
   Array width (Complex a) ->
   Array (Shape.ZeroBased n, width) (Complex a)
_extrudeTwiddleMatrixSimple c k x =
   Matrix.scaleColumns x $ twiddleMatrix c (Shape.ZeroBased k, Array.shape x)

extrudeTwiddleMatrix ::
   (Shape.Indexed width, Shape.Index width ~ n, Eq width,
    Integral n, Class.Real a) =>
   a -> n ->
   Array width (Complex a) ->
   Array (Shape.ZeroBased n, width) (Complex a)
extrudeTwiddleMatrix c k x =
   Matrix.reshapeHeight (Shape.ZeroBased k) $
   Matrix.fromRows (Array.shape x) $
   map (Vector.mul x) $ List.genericTake k $
   twiddleMatrixRowsAuto c (Shape.ZeroBased k, Array.shape x)


-- ToDo: deduplicate with Utility.takeCenter
takeCenterTransposed ::
   (Shape.C leftPad, Shape.C rightPad, Shape.C height, Shape.C width) =>
   (Shape.C sh, Class.Floating a) =>
   leftPad ::+ width ::+ rightPad ->
   Array (height, sh) a -> BlockMatrix.Matrix width height a
takeCenterTransposed shape x =
   BlockMatrix.transposedSubmatrix $
   Array.mapShape
      (Subshape.focus (ViewMatrix.left . ViewMatrix.right) .
       Subshape.submatrixFromMatrix) $
   Matrix.reshapeWidth shape x

{-# INLINE takeTransposed #-}
takeTransposed ::
   (Shape.C height, Shape.C width) =>
   (Shape.C sh, Class.Floating a) =>
   Either width width ->
   Array (height, sh) a -> BlockMatrix.Matrix width height a
takeTransposed ewidth x =
   let width0 = Matrix.width x in
   let pad_ w = Shape.ZeroBased $ Shape.size width0 - Shape.size w in
   takeCenterTransposed
      (case ewidth of
         Left width -> Shape.ZeroBased 0 ::+ width ::+ pad_ width
         Right width -> pad_ width ::+ width ::+ Shape.ZeroBased 0)
      x

cyclicTakeZeroBasedTransposed ::
   (Shape.C height, Eq height, Integral n, Class.Floating a) =>
   n ->
   Array (height, Shape.Cyclic n) a ->
   BlockMatrix.Matrix (Shape.ZeroBased n) height a
cyclicTakeZeroBasedTransposed newWidth x =
   let m = Shape.cyclicSize $ Matrix.width x in
   case divMod newWidth m of
      (q,r) ->
         BlockMatrix.reshapeHeight (Shape.ZeroBased newWidth) $
         BlockMatrix.fromBlockColumn (Matrix.height x) $
         List.genericReplicate q
            (BlockMatrix.transposed $
               Matrix.reshapeWidth (Shape.ZeroBased m) x) ++
         [takeTransposed (Left $ Shape.ZeroBased r) x]

cyclicTakeIntervalTransposed ::
   (Shape.C height, Eq height, Integral n, Class.Floating a) =>
   Shape.Interval n ->
   Array (height, Shape.Cyclic n) a ->
   BlockMatrix.Matrix (Shape.Interval n) height a
cyclicTakeIntervalTransposed rng@(Shape.Interval from to) x =
   let m = Shape.cyclicSize $ Matrix.width x in
   let (fromQ, fromR) = divMod from m in
   let (toQ,   toR)   = divMod to m in
   if fromQ == toQ
   then
      takeCenterTransposed
         (Shape.ZeroBased fromR ::+ rng ::+ Shape.ZeroBased (m - toR - 1))
         x
   else
      BlockMatrix.reshapeHeight rng $
      BlockMatrix.fromBlockColumn (Matrix.height x) $
      takeTransposed (Right $ Shape.ZeroBased (m-fromR)) x :
      List.genericReplicate (toQ-fromQ-1)
         (BlockMatrix.transposed $
            Matrix.reshapeWidth (Shape.ZeroBased m) x) ++
      takeTransposed (Left $ Shape.ZeroBased (toR+1)) x :
      []

blockCyclicFromHalf ::
   (Shape.C height, Eq height, Integral n, Class.Floating a) =>
   Array (height, Spectrum.Half n) a ->
   Array (height, Shape.Cyclic n) a
blockCyclicFromHalf x =
   let width@(Spectrum.Half m) = Matrix.width x in
   Matrix.reshapeWidth (Shape.Cyclic m) $ BlockMatrix.toMatrix $
   x
   &|||
   (let mirrorWidth = div (m-1) 2
        slc = Slice.sub 1 mirrorWidth .
               Slice.reshape (Shape.ZeroBased $ fromIntegral $ Shape.size width)
    in Matrix.fromRowArray (Shape.ZeroBased mirrorWidth) $
       fmap (Vector.conjugate . VectorSlice.reverse . VectorSlice.slice slc) $
       VectorSlice.slicesVector Slice.rowArray x)

halfTakeHalfTransposed ::
   (Shape.C height, Eq height, Integral n, Class.Floating a) =>
   n ->
   Array (height, Spectrum.Half n) a ->
   BlockMatrix.Matrix (Spectrum.Half n) height a
halfTakeHalfTransposed newWidth =
   let halfShape = Spectrum.Half newWidth in
   BlockMatrix.reshapeHeight halfShape .
   cyclicTakeZeroBasedTransposed (fromIntegral $ Shape.size halfShape) .
   blockCyclicFromHalf


{- |
The prefix length m must be a divisor of the data size n.

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

   takeDivisorFourier sign mcut m x
   =~=
   takeShape (Shape.ZeroBased mcut) (Rank1.fourier sign x)
:}

It works this way:
Perform the three steps of Cooley-Tukey in an optimized way.
The first step, namely column-wise transforms, is performed as normal.
The last step collapses to just computing the absolute terms,
which is just row-wise sums.
We combine it with the second step, namely apply twiddling factors.
-}
{-
ToDo:
it would be nice to have a columnwise-FFT,
this would save us an explicit matrix transposition
-}
{-
ToDo:
If you really take a divisor size,
we could make the result Shape.Cyclic and expose this function.
This would simplify the subsampling QuickCheck test.
-}
{-
ToDo:
turn mcut parameter into a shape parameter, with a shape type class:
Shape.ZeroBased: take frequencies starting from zero
Shape.Interval: take frequencies from an arbitrary interval
Shape.Cyclic: take frequencies, but only for a divisor of the original size
-}
takeDivisorFourier ::
   (Integral n, Class.Real a) =>
   Rank1.Sign ->
   n -> n ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Shape.ZeroBased n) (Complex a)
takeDivisorFourier sign mcut m 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
   applyTwiddleMatrix c .
   {-
   One transposition should be more efficient
   than many scattered inner products.
   I hope that we can replace the transpositions by one columnwise batched FFT.
   -}
   BlockMatrix.toMatrix .
   cyclicTakeZeroBasedTransposed mcut .
   Batch.fourier sign .
   Matrix.transpose .
   Array.mapShape (shapeCyclicBlock m)
   $ x


{- |
prop> :{
   QC.forAll (genCyclic complex_) $ \x sign ->
   let n = Shape.size $ Array.shape x in
   QC.forAll (QC.elements $ divisors n) $ \m ->
   QC.forAll (chooseLogarithmic n) $ \mcut from ->
   let shape = Shape.intervalFromShifted $ Shape.Shifted from mcut in

   takeDivisorFourierInterval sign shape m x
   =~=
   cyclicTake shape (Rank1.fourier sign x)
:}
-}
takeDivisorFourierInterval ::
   (Integral n, Class.Real a) =>
   Rank1.Sign ->
   Shape.Interval n -> n ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Shape.Interval n) (Complex a)
takeDivisorFourierInterval sign rng m 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
   -- ToDo: shall we merge BlockMatrix.toMatrix into applyTwiddleMatrix?
   applyTwiddleMatrix c .
   BlockMatrix.toMatrix . cyclicTakeIntervalTransposed rng .
   Batch.fourier sign .
   Matrix.transpose .
   Array.mapShape (shapeCyclicBlock m)
   $ x


{- |
@takeNonDivisorFourier n@ computes the first @n@ spectral coefficients
efficiently in case the data size has no or no appropriate divisors.
It employs the Bluestein chirp transform
optimized with 'convolveLongToShort'.

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

   takeDivisorFourier sign mcut m x
   =~=
   takeNonDivisorFourier sign mcut x
:}
-}
takeNonDivisorFourier ::
   (Integral n, Bits n, Class.Real a) =>
   Rank1.Sign ->
   n ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Shape.ZeroBased n) (Complex a)
takeNonDivisorFourier sign m x =
   let Shape.Cyclic n = Array.shape x in
   let chirp = nonNegChirp sign n n in
   let chirpCut = takeShape (Shape.ZeroBased m) chirp in
   VectorSlice.mul (VectorSymb.conjugate chirpCut) $
   convolveLongToShort
      -- ToDo: Shape.zeroBasedFromCyclic
      (Vector.mulConj chirp $ Array.reshape (Shape.ZeroBased n) x)
      (Array.reshape (Shape.ZeroBased (n+m-1)) $ Chunk.toVector $
       Chunk.reversed chirp <>
         Chunk.simple (VectorSlice.slice (Slice.drop 1) chirpCut))

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

   takeNonDivisorFourierInterval sign shape x
   =~=
   cyclicTake shape (Rank1.fourier sign x)
:}
-}
takeNonDivisorFourierInterval ::
   HasCallStack =>
   (Integral n, Bits n, Class.Real a) =>
   Rank1.Sign ->
   Shape.Interval n ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Shape.Interval n) (Complex a)
takeNonDivisorFourierInterval sign rng x =
   let Shape.Cyclic n = Array.shape x in
   let rngNorm@(Shape.Interval from to) = normalizeInterval n rng in
   let chirp = nonNegChirp sign n (n+1) in
   Array.reshape rng $
   VectorSlice.mul (VectorSymb.conjugate $ interval2FromChirp rngNorm chirp) $
   convolveLongToShortInterval
      (let shapeN = Shape.Interval 0 (n-1) in
       VectorSlice.mul
         (VectorSymb.conjugate $ takeShape shapeN chirp)
         (Array.reshape shapeN x))
      (interval3FromChirp (Shape.Interval (from-n+1) to) chirp)



takeDivisorFourierRC ::
   (Integral n, Class.Real a) =>
   n -> n ->
   Array (Shape.Cyclic n) a ->
   Array (Spectrum.Half n) (Complex a)
takeDivisorFourierRC mcut m x =
   let Shape.Cyclic n = Array.shape x in
   let c = 2*pi / fromIntegral n in
   -- ToDo: instance Shape.IntegralRange Spectrum.Half - this would also enable access to halfSize
   -- alternative: make Shape.toShifted conversion a parameter
   applyTwiddleMatrixSimple c . BlockMatrix.toMatrix .
   halfTakeHalfTransposed mcut .
   withExtraUnitDim Batch.fourierRC .
   Matrix.transpose .
   Array.mapShape (shapeCyclicBlock m)
   $ x

-- ToDo: comfort-fftw:Shape should export halfSize, such that we do not need fromIntegral . Shape.size
{- |
prop> :{
   QC.forAll (genCyclic real_) $ \x ->
   let n = Shape.cyclicSize $ Array.shape x in
   QC.forAll (QC.elements $ divisors n) $ \m ->
   QC.forAll (chooseLogarithmic n) $ \mcut ->

   takeDivisorFourierRC mcut m x
   =~=
   takeNonDivisorFourierRC mcut x
:}
-}
takeNonDivisorFourierRC ::
   (Integral n, Bits n, Class.Real a) =>
   n ->
   Array (Shape.Cyclic n) a ->
   Array (Spectrum.Half n) (Complex a)
takeNonDivisorFourierRC m x =
   let Shape.Cyclic n = Array.shape x in
   let chirp = nonNegChirp Rank1.Forward n n in
   let halfShape = Spectrum.Half m in
   let chirpCut = takeShape halfShape chirp in
   let halfSize = fromIntegral $ Shape.size halfShape in
   VectorSlice.mul (VectorSymb.conjugate chirpCut) $
   Array.reshape halfShape $
   convolveLongToShort
      (Vector.mulRealConj (Array.reshape (Shape.ZeroBased n) x) chirp)
      (Array.reshape (Shape.ZeroBased (n+halfSize-1)) $ Chunk.toVector $
       Chunk.reversed chirp
       <>
       Chunk.simple
         (VectorSlice.slice
            (Slice.drop 1 . Slice.reshape (Shape.ZeroBased halfSize))
            chirpCut))


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

   Rank1.fourier sign (padCyclic n x)
   =~=
   extendDivisorFourier sign n m x
:}
-}
extendDivisorFourier ::
   (Integral n, Class.Real a) =>
   Rank1.Sign ->
   n -> n ->
   Array (Shape.ZeroBased n) (Complex a) ->
   Array (Shape.Cyclic n) (Complex a)
extendDivisorFourier sign n m 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) .
   Matrix.transpose .
   Batch.fourier sign .
   blockToCyclic m .
   extrudeTwiddleMatrix c (div n m) $ x


extendDivisorFourierCR ::
   (Integral n, Class.Real a) =>
   n -> n ->
   Array (Spectrum.Half n) (Complex a) ->
   Array (Shape.Cyclic n) a
extendDivisorFourierCR n m x =
   let c = 2*pi / fromIntegral n in
   Array.reshape (Shape.Cyclic n) .
   Matrix.transpose .
   withExtraUnitDim Batch.fourierCR .
   halfBlockToHalf m .
   extrudeTwiddleMatrix c (div n m) $ x


spectrumIntervalFromHalf ::
   (Integral n, Class.Real a) =>
   Array (Spectrum.Half n) (Complex a) ->
   Array (Shape.Interval n) (Complex a)
spectrumIntervalFromHalf xh =
   let mh = fromIntegral $ Shape.size $ Array.shape xh in
   let xz = VectorSlice.sliceVector (Slice.drop 1) $
            Array.reshape (Shape.ZeroBased mh) xh in
   Array.reshape (Shape.Interval (1-mh) (mh-1)) $ Chunk.toVector $
      Chunk.reversed (VectorSlice.conjugate xz)
      <>
      Chunk.singleton (Complex.realPart(xh!0):+0)
      <>
      Chunk.simple xz

{- |
prop> :{
   QC.forAll genVectorShape $ \shape sign ->
   let n = Shape.zeroBasedSize shape in
   QC.forAll (QC.elements $ divisors n) $ \m ->
   QC.forAll (chooseLogarithmic n) $ \mcut from ->
   let rng = Shape.intervalFromShifted $ Shape.Shifted from mcut in
   QC.forAll (genVectorForShape complex_ rng) $ \x ->

   extendDivisorFourierInterval sign n m x
   =~=
   Rank1.fourier sign (spreadCyclic n x)
:}
-}
extendDivisorFourierInterval ::
   (Integral n, Class.Real a) =>
   Rank1.Sign ->
   n -> n ->
   Array (Shape.Interval n) (Complex a) ->
   Array (Shape.Cyclic n) (Complex a)
extendDivisorFourierInterval sign n m 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) .
   Matrix.transpose .
   Batch.fourier sign .
   intervalBlockToCyclic m .
   extrudeTwiddleMatrix c (div n m) $ x



{- |
@extendNonDivisorFourier n@ computes @n@ spectral coefficients efficiently
from a shorter signal
in case @n@ has no or no appropriate divisors.
It employs the Bluestein chirp transform optimized with 'convolveShortLong'.

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 ->

   Rank1.fourier sign (padCyclic n x)
   =~=
   extendNonDivisorFourier sign n x
:}

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 ->

   takeNonDivisorFourier (Rank1.flipSign sign) m
      (extendNonDivisorFourier sign n x)
   =~=
   Vector.scaleReal (fromIntegral n) x
:}
-}
extendNonDivisorFourier ::
   (Show n, Show a) =>
   (Integral n, Bits n, Class.Real a) =>
   Rank1.Sign ->
   n ->
   Array (Shape.ZeroBased n) (Complex a) ->
   Array (Shape.Cyclic n) (Complex a)
extendNonDivisorFourier sign n x =
   let chirp = nonNegChirp sign n n in
   let chirpM = takeShape (Array.shape x) chirp in
   Array.reshape (Shape.Cyclic n) $
   VectorSlice.mul
      (VectorSymb.conjugate $
       Array.mapShape Shape.intervalFromZeroBased chirp) $
   convolveShortLongCut
      (Array.mapShape Shape.intervalFromZeroBased $
         VectorSlice.mul (VectorSymb.conjugate chirpM) x)
      (appendRevToInterval1 chirpM
         (VectorSlice.sliceVector (Slice.drop 1) chirp))

{- |
-- ToDo: Can we get rid of the n>=2 constraint?
prop> :{
   QC.forAll genCyclicShape $ \shape ->
   let n = Shape.cyclicSize shape in
   n>=2 QC.==>
   QC.forAll (chooseLogarithmic (n-1)) $ \m ->
   let halfShape = Spectrum.Half m in
   QC.forAll (genVectorForShape complex_ halfShape) $ \xh ->

   extendNonDivisorFourierInterval Rank1.Backward n
      (spectrumIntervalFromHalf xh)
   =~=
   Vector.fromReal (extendNonDivisorFourierCR n xh)
:}

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

   Vector.realPart
      (Rank1.fourier Rank1.Backward
         (padCyclic n $ Array.mapShape (Shape.ZeroBased . Shape.size) $
            Vector.scaleReal 2 x // [(0, x!0)]))
   =~=
   extendNonDivisorFourierCR n x
:}

prop> :{
   QC.forAll genVectorShape $ \shape ->
   let n = Shape.zeroBasedSize shape in
   QC.forAll (QC.elements $ divisors n) $ \m ->
   QC.forAll (chooseLogarithmic n) $ \mcut ->
   QC.forAll (genVectorForShape complex_ (Spectrum.Half mcut)) $ \x ->

   extendDivisorFourierCR n m x
   =~=
   extendNonDivisorFourierCR n x
:}

-- ToDo: handle last coefficient consistently with extendFourierCR
prop> :{
   QC.forAll genCyclicShape $ \shape ->
   let n = Shape.cyclicSize shape in
   forChoose (0, div n 2) $ \f a b ->
   let c = -2*pi / fromIntegral n in
   let (a2,b2) = if f==0 then (a,b) else (2*a,2*b) in

   extendNonDivisorFourierCR n
      (Array.fromAssociations 0 (Spectrum.Half (2*f))
         [(f, a:+b :: Complex Float)])
   =~=
   Array.sample shape
      (\t ->
         a2 * cos (c * fromIntegral (mod (t*f) n)) +
         b2 * sin (c * fromIntegral (mod (t*f) n)))
:}
-}
extendNonDivisorFourierCR ::
   (Show n, Show a) =>
   (Integral n, Bits n, Class.Real a) =>
   n ->
   Array (Spectrum.Half n) (Complex a) ->
   Array (Shape.Cyclic n) a
extendNonDivisorFourierCR n xh =
   let m = fromIntegral (Shape.size (Array.shape xh)) in
   let xz = VectorSlice.sliceVector (Slice.drop 1) $
               Array.reshape (Shape.ZeroBased m) xh in
   let x = Array.reshape (Shape.Interval (1-m) (m-1)) $ Chunk.toVector $
               Chunk.conjugate (Chunk.reversed xz)
               <>
               Chunk.singleton (Complex.realPart(xh!0):+0)
               <>
               Chunk.simple xz in
-- ToDo: the chirp is 2*n-periodic - can we make use of this in a Shape.Cyclic?
-- ToDo: Can we make the convolution symmetric + odd-symmetric and save some operations by this trick?
   let chirp = nonNegChirp Rank1.Backward n (n+1) in
   let chirpM = takeShape (Shape.ZeroBased m) chirp in
   let chirpM2 =
         appendRevToInterval1 chirpM
            (VectorSlice.sliceVector (Slice.sub 1 (m-1)) chirp) in
   let chirpR =
         Array.reshape (Shape.Interval (1-m) (n+m-2)) $ Chunk.toVector $
         Chunk.reversed chirpM
         <>
         Chunk.simple (VectorSlice.sliceVector (Slice.sub 1 (n-1)) chirp)
         <>
         Chunk.reversed (VectorSlice.sliceVector (Slice.drop (n-m+2)) chirp)
         in
   Array.reshape (Shape.Cyclic n) $
   realPartOfMulConj (takeShape (Shape.Interval 0 (n-1)) chirp) $
   convolveShortLongCut (Vector.mulConj chirpM2 x) chirpR

realPartOfMulConj ::
   (Shape.C sh, Class.Real a, Eq sh) =>
   VectorSlice.T Layout.Slice sh (Complex a) ->
   VectorSlice.T Layout.Slice sh (Complex a) ->
   Array sh a
realPartOfMulConj x y =
   Matrix.dotRowwise
      (Vector.decomplex (VectorSlice.toVector x))
      (Vector.decomplex (VectorSlice.toVector y))

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

   extendNonDivisorFourier sign n x
   =~=
   extendNonDivisorFourierInterval sign n (Array.reshape (Shape.Interval 0 (m-1)) x)
:}

prop> :{
   QC.forAll genCyclicShape $ \shape sign f a ->
   let n = Shape.cyclicSize shape in
   let cu = 2*pi / fromIntegral n in
   let c = case sign of Rank1.Forward -> -cu; Rank1.Backward -> cu in

   extendNonDivisorFourierInterval sign n
      (Array.replicate (Shape.Interval f f) (a :: Complex Float))
   =~=
   Array.sample shape (\t -> a * cis (c * fromIntegral (mod (t*f) n)))
:}

prop> :{
   QC.forAll genCyclicShape $ \shape sign ->
   let n = Shape.cyclicSize shape in
   n>=2 QC.==>
   forChoose (1,n-1) $ \d f0 a b ->
   let f1 = f0+d in
   let cu = 2*pi / fromIntegral n in
   let c = case sign of Rank1.Forward -> -cu; Rank1.Backward -> cu in

   extendNonDivisorFourierInterval sign n
      (Array.fromAssociations 0 (Shape.Interval f0 f1)
         [(f0,a), (f1,b :: Complex Float)])
   =~=
   Array.sample shape
      (\t ->
         a * cis (c * fromIntegral (mod (t*f0) n)) +
         b * cis (c * fromIntegral (mod (t*f1) n)))
:}
-}
extendNonDivisorFourierInterval ::
   (Show n, Show a) =>
   (Integral n, Bits n, Class.Real a) =>
   Rank1.Sign ->
   n ->
   Array (Shape.Interval n) (Complex a) ->
   Array (Shape.Cyclic n) (Complex a)
extendNonDivisorFourierInterval sign n x0 =
   let x = Array.mapShape (normalizeInterval n) x0 in
   let rng@(Shape.Interval from to) = Array.shape x in
--   let m = to0 - from0 + 1 in
   {-
   We could make the function work for larger Ranges,
   but the purpose of the function is efficiency,
   which cannot be achieved for large input Ranges.
   -}
   if' (to-from>=n)
         (error $ "extendNonDivisorFourierInterval: " ++
            "input range larger than output size") $
-- ToDo: extendNonDivisorFourier: maybe convolveLongToShort is better here, too, since we cut away stuff at the end
-- or a correlateLongToShort
-- ToDo: the chirp is 2*n-periodic
   let chirp = nonNegChirp sign n (n+1) in
   Array.reshape (Shape.Cyclic n) $
   VectorSlice.mul
      (VectorSymb.conjugate $ takeShape (Shape.Interval 0 (n-1)) chirp) $
   convolveShortLongCut
      (VectorSlice.mul (VectorSymb.conjugate $ interval2FromChirp rng chirp) x)
      (interval3FromChirp (Shape.Interval (-to) (n-1-from)) chirp)



{-
ToDo:
Move to Experimental module
as we might change the semantics to something
that handles borders better.
-}
{- ToDo:
A cyclic border handling would nicely allow for invertibility
but does not fit the usual application of STFT.
-}
shortTimeFourierForward ::
   (Integral n, Class.Real a) =>
   Array (Shape.ZeroBased n) a -> n ->
   Array (Shape.ZeroBased n) (Complex a) ->
   Array (Shape.ZeroBased n, Shape.Cyclic n) (Complex a)
shortTimeFourierForward window overlap x =
   Batch.fourier Batch.Forward $
   Matrix.mapWidth (Shape.Cyclic . Shape.zeroBasedSize) $
   Matrix.scaleColumnsReal window $
   toOverlappingBlocks (Array.shape window) overlap x

shortTimeFourierRC ::
   (Integral n, Class.Real a) =>
   Array (Shape.ZeroBased n) a -> n ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n, Spectrum.Half n) (Complex a)
shortTimeFourierRC window overlap x =
   withExtraUnitDim Batch.fourierRC $
   Matrix.mapWidth (Shape.Cyclic . Shape.zeroBasedSize) $
   Matrix.scaleColumns window $
   toOverlappingBlocks (Array.shape window) overlap x

{- ToDo:
could be implemented without an inner Haskell loop
by writing to mutable submatrices
But the access pattern would not be super-cache friendly.
-}
toOverlappingBlocks ::
   (Integral n, Class.Floating a) =>
   Shape.ZeroBased n -> n ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n, Shape.ZeroBased n) a
toOverlappingBlocks windowShape overlap x =
   let blockSize = Shape.zeroBasedSize windowShape in
   let stepSize = div blockSize overlap in
   let n = Shape.zeroBasedSize $ Array.shape x in
   Matrix.mapHeight (fmap fromIntegral) $
   Matrix.fromRows windowShape $
   map (\k -> VectorSlice.sliceVector (Slice.sub k blockSize) x) $
   takeWhile (<= n-blockSize) $ iterate (stepSize+) 0

{- |
For invertibility we must make sure
that the elementwise product of windows of forward and backward transform
sum up to constant one when rotated.

prop> :{
   QC.forAll genVectorShape $ \shapeShort ->
   forChoose (1,10) $ \stepSize ->
   forChoose (2,10) $ \overlap ->
   let m = stepSize * overlap in
   let window = Array.sample (Shape.ZeroBased m)
                     (\i -> sin (pi * fromIntegral i / fromIntegral m)) in
   let shape = fmap (\n -> n + mod (-n) (stepSize*overlap)) shapeShort in
   QC.forAll (genVectorForShape complex_ shape) $ \x ->
   let zeroPad = Vector.zero $ Shape.ZeroBased $ (overlap-1)*stepSize in
   let xpad = zeroPad <> x <> zeroPad in

   Vector.scaleReal (fromIntegral (m*overlap) / 2) xpad
   =~=
   shortTimeFourierBackward window overlap
      (shortTimeFourierForward window overlap xpad)
:}

The factor @m@ comes from the backward Fourier transform,
whereas @overlap/2@ is @overlap@ times DC offset of the squared window.
-}
shortTimeFourierBackward ::
   (Integral n, Class.Real a) =>
   Array (Shape.ZeroBased n) a -> n ->
   Array (Shape.ZeroBased n, Shape.Cyclic n) (Complex a) ->
   Array (Shape.ZeroBased n) (Complex a)
shortTimeFourierBackward window overlap x =
   fromOverlappingBlocks overlap $
   Matrix.scaleColumnsReal window $
   Matrix.mapWidth (Shape.ZeroBased . Shape.cyclicSize) $
   Batch.fourier Batch.Backward x

shortTimeFourierCR ::
   (Integral n, Class.Real a) =>
   Array (Shape.ZeroBased n) a -> n ->
   Array (Shape.ZeroBased n, Spectrum.Half n) (Complex a) ->
   Array (Shape.ZeroBased n) a
shortTimeFourierCR window overlap x =
   fromOverlappingBlocks overlap $
   Matrix.scaleColumns window $
   Matrix.mapWidth (Shape.ZeroBased . Shape.cyclicSize) $
   withExtraUnitDim Batch.fourierCR x

{- |
prop> :{
   QC.forAll (genVector real_) $ \dat ->
   let n = Shape.zeroBasedSize $ Array.shape dat in
   forChoose (1, min n 4) $ \overlap ->
   QC.forAll (chooseLogarithmic (div n overlap)) $ \stepSize ->
   let width = overlap * stepSize in
   let height = div n width in
   let xs = takeShape (Shape.ZeroBased height, Shape.ZeroBased width) dat in

   fromOverlappingBlocks overlap xs
   =~=
   arrangeMutable
      ((height-1 + overlap) * stepSize)
      stepSize xs (Vector.zero (Shape.ZeroBased 0))
:}
-}
fromOverlappingBlocks ::
   (Integral n, Class.Floating a) =>
   n ->
   Array (Shape.ZeroBased n, Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
fromOverlappingBlocks overlap x =
   let stepSize = div (Shape.zeroBasedSize $ Matrix.width x) overlap in
   Array.mapShape (fmap fromIntegral) $
   foldl1 Vector.add $
   map (Chunk.toVector . mconcat) $
   ListHT.sliceHorizontal (fromIntegral overlap) $
   (\chunks ->
      let filling =
            map Chunk.zero $
            List.genericTake overlap $ iterate (stepSize+) 0 in
      filling ++ map Chunk.simple chunks ++ reverse filling) $
   BoxedArray.toList $
   VectorSlice.slicesVector Slice.rowArray x