module Numeric.FFTW.Extra.Shape where

import qualified Data.Array.Comfort.Shape as Shape


{- $setup
>>> :set -XTypeFamilies
>>> import Numeric.FFTW.Extra.Shape
>>> import Test.Numeric.FFTW.Extra.Utility
>>>
>>> import qualified Data.Array.Comfort.Shape as Shape
>>>
>>> import qualified Test.QuickCheck as QC
-}


{- |
>>> convolve (Shape.ZeroBased 0) (Shape.ZeroBased (10::Int))
ZeroBased {... 0}

>>> convolve (Shape.ZeroBased 6) (Shape.ZeroBased (9::Int))
ZeroBased {... 14}

prop> :{
   QC.forAll genVectorShape $ \shapeX ->
   QC.forAll genVectorShape $ \shapeY ->

   convolve shapeX shapeY
   ==
   convolve shapeY shapeX
:}

prop> :{
   QC.forAll genVectorShape $ \shapeX ->
   QC.forAll genVectorShape $ \shapeY ->
   QC.forAll genVectorShape $ \shapeZ ->

   convolve shapeX (convolve shapeY shapeZ)
   ==
   convolve (convolve shapeX shapeY) shapeZ
:}

prop> :{
   QC.forAll genShiftedShape $ \shapeX ->
   QC.forAll genShiftedShape $ \shapeY ->

   convolve shapeX shapeY
   ==
   convolve shapeY shapeX
:}

prop> :{
   QC.forAll genShiftedShape $ \shapeX ->
   QC.forAll genShiftedShape $ \shapeY ->
   QC.forAll genShiftedShape $ \shapeZ ->

   convolve shapeX (convolve shapeY shapeZ)
   ==
   convolve (convolve shapeX shapeY) shapeZ
:}

prop> :{
   QC.forAll genIntervalShape $ \shapeX ->
   QC.forAll genIntervalShape $ \shapeY ->

   convolve shapeX shapeY
   ==
   convolve shapeY shapeX
:}

prop> :{
   QC.forAll genIntervalShape $ \shapeX ->
   QC.forAll genIntervalShape $ \shapeY ->
   QC.forAll genIntervalShape $ \shapeZ ->

   convolve shapeX (convolve shapeY shapeZ)
   ==
   convolve (convolve shapeX shapeY) shapeZ
:}

prop> :{
   QC.forAll genIntervalShape $ \shapeX ->
   QC.forAll genIntervalShape $ \shapeY ->

   convolveSize (size shapeX) (size shapeY)
   ==
   size (convolve shapeX shapeY)
:}

prop> :{
   QC.forAll genIntervalShape $ \shapeX ->
   QC.forAll genIntervalShape $ \shapeY ->
   QC.forAll genIntervalShape $ \shapeZ ->

   correlate shapeX (correlate shapeY shapeZ)
   ==
   correlate (convolve shapeX shapeY) shapeZ
:}

prop> :{
   QC.forAll genIntervalShape $ \shapeX ->
   QC.forAll genIntervalShape $ \shapeY ->

   correlateSize (size shapeX) (size shapeY)
   ==
   size (correlate shapeX shapeY)
:}
-}
class (Shape.C sh) => Convolve sh where
   convolve :: sh -> sh -> sh
   correlate :: sh -> sh -> sh
   size :: sh -> Shape.Index sh

convolveSize :: (Integral n) => n -> n -> n
convolveSize n m = if n==0 || m==0 then 0 else n+m-1

correlateSize :: (Integral n) => n -> n -> n
correlateSize n m = max 0 $ m+1-n

unshift :: (Convolve sh) => sh -> Shape.ZeroBased (Shape.Index sh)
unshift = Shape.ZeroBased . size

instance (Integral n) => Convolve (Shape.ZeroBased n) where
   convolve (Shape.ZeroBased n) (Shape.ZeroBased m) =
      Shape.ZeroBased $ convolveSize n m
   correlate (Shape.ZeroBased n) (Shape.ZeroBased m) =
      Shape.ZeroBased $ correlateSize n m
   size = Shape.zeroBasedSize

instance (Integral n) => Convolve (Shape.Shifted n) where
   convolve (Shape.Shifted fromX n) (Shape.Shifted fromY m) =
      Shape.Shifted (fromX+fromY) (convolveSize n m)
   correlate (Shape.Shifted fromX n) (Shape.Shifted fromY m) =
      Shape.Shifted (fromX+fromY) (correlateSize n m)
   size = Shape.shiftedSize

instance (Integral n) => Convolve (Shape.Interval n) where
   convolve (Shape.Interval fromX toX) (Shape.Interval fromY toY) =
      Shape.Interval (fromX+fromY) (toX+toY)
   correlate (Shape.Interval fromX toX) (Shape.Interval fromY toY) =
      Shape.Interval (fromY-fromX) (toY-toX)
   size (Shape.Interval from to) = max 0 $ to+1-from