{-# LANGUAGE TypeOperators #-}
module Numeric.FFTW.Extra.Rank1Sketch where

import Numeric.FFTW.Extra.Rank1.Convolution (convolve, convolveLongToShort)

import qualified Numeric.FFTW.Batch as Batch
import qualified Numeric.FFTW.Shape as Spectrum
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 qualified Data.Complex as Complex
import Data.Complex (Complex((:+)))
import Data.Bits (Bits)

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



-- Internally all signals are appended with exact padding.
-- This way we need to pad to better size only at the end.
_convolveOneManyHeterogenous ::
   (Integral n, Bits n, Class.Real a) =>
   Array (Shape.ZeroBased n) (Complex a) ->
   [Array (Shape.ZeroBased n) (Complex a)] ->
   [Array (Shape.ZeroBased n) (Complex a)]
_convolveOneManyHeterogenous = undefined


sqr :: (Class.Floating a) => a -> a
sqr x = x*x

takeNonDivisorFourierTooLong, takeNonDivisorFourierPlain ::
   (Integral n, Bits n, Class.Real a) =>
   Batch.Sign ->
   n ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Shape.ZeroBased n) (Complex a)
takeNonDivisorFourierTooLong sign m x =
   let Shape.Cyclic n = Array.shape x in
   let cu = pi / fromIntegral n in
   let c = case sign of Batch.Forward -> cu; Batch.Backward -> -cu in
   let chirp = Array.sample (Shape.ZeroBased n)
                  (\k -> Complex.cis (c * sqr (fromIntegral k))) in
   let symmChirp = Vector.reverse chirp <> Vector.drop 1 chirp in
   Vector.take m $
   Vector.mulConj chirp $
   convolveLongToShort
      (Vector.mulConj chirp $
         Array.mapShape (Shape.ZeroBased . Shape.cyclicSize) x)
      symmChirp

takeNonDivisorFourierPlain sign m x =
   let Shape.Cyclic n = Array.shape x in
   let cu = pi / fromIntegral n in
   let c = case sign of Batch.Forward -> cu; Batch.Backward -> -cu in
   let chirp = Array.sample (Shape.ZeroBased n)
                  (\k -> Complex.cis (c * sqr (fromIntegral k))) in
   let symmChirp = Vector.reverse chirp <> Vector.drop 1 chirp in
   Vector.take m $
   Vector.mulConj chirp $
   Vector.take n $
   Vector.drop (n-1) $
   convolve
      (Vector.mulConj chirp $
         Array.mapShape (Shape.ZeroBased . Shape.cyclicSize) x)
      symmChirp


multiplyRealSpectra ::
   (Integral n, Class.Real a) =>
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Shape.Cyclic n) (Complex a) ->
   Array (Spectrum.Half n) (Complex a)
multiplyRealSpectra fx fy =
   Array.sample
      (Spectrum.Half (Shape.cyclicSize $ Array.shape fx))
      (\k -> (sqr(fx!k) - Complex.conjugate (sqr(fy!(-k)))) / (0:+4))