module Numeric.BLAS.Rank1.Convolution where

import qualified Numeric.FFTW.Extra.Shape as CShape
import Numeric.FFTW.Extra.Utility
         (takeShape, blockToCyclic, acyclic, acyclicWidth, padBlock)

import qualified Numeric.BLAS.Matrix.RowMajor.Block as BlockMatrix
import qualified Numeric.BLAS.Matrix.RowMajor.Mutable as MutMatrix
import qualified Numeric.BLAS.Matrix.RowMajor as Matrix
import qualified Numeric.BLAS.Vector.Mutable as MutVector
import qualified Numeric.BLAS.Vector.Chunk as Chunk
import qualified Numeric.BLAS.Vector.Slice as VectorSlice
import qualified Numeric.BLAS.Vector as Vector
import qualified Numeric.BLAS.Scalar as Scalar
import qualified Numeric.BLAS.Slice as Slice
import qualified Numeric.BLAS.Subobject.Shape as Subshape
-- import qualified Numeric.BLAS.Subobject.View.Matrix as ViewMatrix
import qualified Numeric.BLAS.Subobject.View as View
import qualified Numeric.Netlib.Class as Class
import Numeric.BLAS.Matrix.RowMajor.Block ((&|||))
import Numeric.BLAS.Vector.Slice ((|+|), (|-|))
-- import Numeric.Netlib.Modifier (Conjugation(NonConjugated))

import qualified Data.Array.Comfort.Storable.Mutable as MutArray
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.List.HT as ListHT
-- import qualified Data.List as List
import qualified Data.Foldable as Fold
import Data.Foldable (for_)
-- import Data.Tuple.HT (double)

import Control.Monad.ST (ST)
import Control.Monad (when)


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


-- ToDo: generalize to VectorSlice.C
{- |
It is more efficient if the first operand is the shorter one.
-}
convolve ::
   (Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
convolve xs ys = MutArray.runST (do
   zs <-
      MutVector.new
         (CShape.convolve (Array.shape xs) (Array.shape ys)) Scalar.zero
   addConvolve xs ys $ MutVector.fromVector zs
   return zs)

addConvolve ::
   (VectorSlice.C v, Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   v (Shape.ZeroBased n) a ->
   MutVector.T (ST s) (Shape.ZeroBased n) a ->
   ST s ()
addConvolve xs ys zs = do
   let (Shape.ZeroBased m) = VectorSlice.shape ys
   for_ (Array.toAssociations xs) $ \(k,x) ->
      MutVector.mac (MutVector.slice (Slice.sub k m) zs) x ys

{- |
If memory is not an issue,
this function computes the full multiplication table explicitly
and then sums up all columns.
For small matrices this may pay off.

prop> :{
   QC.forAll (genVector number_) $ \x ->
   QC.forAll (genVector number_) $ \y ->

   convolve x y
   =~=
   convolveViaTable x y
:}
-}
convolveViaTable ::
   (Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
convolveViaTable xs ys =
   let shx = Array.shape xs in
   let shy = Array.shape ys in
   Matrix.sumColumns $ VectorSlice.toVector $
   takeShape (shx, CShape.convolve shx shy) $
   BlockMatrix.toMatrix $
   Matrix.tensorProduct xs ys
   &|||
   BlockMatrix.zero shx shx

{-
Vector.dot is based on gemv and writes the result to memory, anyway.
We could use a new MutVector.dot.
-}
correlateLong ::
   (Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
correlateLong x y =
   let shx@(Shape.ZeroBased m) = Array.shape x in
   Array.sample (CShape.correlate shx (Array.shape y)) $ \k ->
      VectorSlice.dot x $ VectorSlice.sliceVector (Slice.sub k m) y

{- |
ToDo:
For really long vectors 'correlateLong' will certainly be more efficient.
Find a better function name? convolveCut?

prop> :{
   QC.forAll (genVector number_) $ \x0 ->
   QC.forAll (genVector number_) $ \y0 ->
   let (x,y) =
         if Shape.size (Array.shape x0) < Shape.size (Array.shape y0)
         then (x0,y0) else (y0,x0) in

   correlateLong x y
   =~=
   convolveLongToShort (Vector.reverse x) y
:}
-}
convolveLongToShort ::
   (Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
convolveLongToShort xs ys = MutArray.runST (do
   let n = Shape.zeroBasedSize $ Array.shape xs
   let m = Shape.zeroBasedSize $ Array.shape ys
   let k = CShape.correlateSize n m
   arr <- MutVector.new (Shape.ZeroBased k) Scalar.zero
   when (k>0) $ addConvolveCut xs ys arr
   return arr)

addConvolveCut ::
   (MutVector.C v, Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   v (ST s) (Shape.ZeroBased n) a ->
   ST s ()
addConvolveCut x y z = do
   let n = Shape.zeroBasedSize $ Array.shape x
   let k = Shape.zeroBasedSize $ MutVector.shape z
   for_ (Array.toAssociations x) $ \(j,xi) ->
      MutVector.mac z xi $ VectorSlice.sliceVector (Slice.sub (n-1-j) k) y

convolveCyclic ::
   (Integral n, Class.Floating a) =>
   Array (Shape.Cyclic n) a ->
   Array (Shape.Cyclic n) a ->
   Array (Shape.Cyclic n) a
convolveCyclic x y =
   Array.reshape (Array.shape x) $
   convolveLongToShort (acyclic x) (replicateCyclic y)

replicateCyclic ::
   (VectorSlice.Mono v, Integral n, Class.Floating a) =>
   v (Shape.Cyclic n) a -> Array (Shape.ZeroBased n) a
replicateCyclic x_ =
   let x = VectorSlice.fromAny x_ in
   let n = Shape.cyclicSize $ VectorSlice.shape x in
   Array.reshape (Shape.ZeroBased (2*n-1)) $
   VectorSlice.append
      (VectorSlice.slice
         (Slice.drop 1 . Slice.mapShape (Shape.ZeroBased . Shape.cyclicSize)) x)
      x
--   VectorSlice.append (VectorSlice.slice (Slice.drop 1) (acyclic x)) x


convolveOneMany_ ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array     (Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a
convolveOneMany_ x y = MutArray.runST (do
   let shx = Array.shape x
   let shy = Matrix.width y
   z <- MutVector.new (Matrix.height y, CShape.convolve shx shy) Scalar.zero
   let subshape =
         (Subshape.submatrixFromMatrix $ MutArray.shape z) {
            Subshape.shape = (Matrix.height y, shx)
         }
   for_ (zip [0..] $ VectorSlice.slicesVector Slice.columns y) $ \(k,yi) ->
      MutMatrix.addTensorProduct 1 yi x $
         MutMatrix.wrap $ MutArray.reshape (subshape{Subshape.start = k}) z
   return z)

convolveOneMany ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array     (Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a
convolveOneMany x y = MutArray.runST (do
   let shx = Array.shape x
   let shy = Matrix.width y
   z <- MutVector.new (Matrix.height y, CShape.convolve shx shy) Scalar.zero
{-
   for_ (BoxedArray.toAssociations $ Matrix.toColumnArray y) $ \(k,yi) ->
      MutMatrix.addTensorProduct 1 yi x $
-}
   for_ (zip [0..] $ VectorSlice.slicesVector Slice.columns y) $ \(k,yi) ->
      MutMatrix.addTensorProduct 1 yi x $
         MutMatrix.viewMatrix (View.subWidth k (Shape.zeroBasedSize shx)) z
   return z)

{- |
Same considerations as for 'convolveViaTable'.

prop> :{
   QC.forAll (genVector number_) $ \x ->
   QC.forAll genVectorShape $ \height ->
   QC.forAll genVectorShape $ \width ->
   QC.forAll (genVectorForShape number_ (height, width)) $ \ys ->

   convolveOneMany x ys
   =~=
   convolveOneManyViaTable x ys
:}
-}
{-
**__**__**___
**__**__**___
**__**__**___

**__**__**__
_**__**__**_
__**__**__**   ___
-}
-- ToDo: convolveCyclicOneManyViaTable - periodize the result of convolveOneManyViaTable
convolveOneManyViaTable ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array     (Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a
convolveOneManyViaTable x ys =
   let         shx  = Array.shape x  in
   let (height,shy) = Array.shape ys in
   let shz = CShape.convolve shx shy in
   Matrix.sumColumns $ VectorSlice.toVector $
   takeShape (shx, (height, shz)) $
   BlockMatrix.toMatrix $
   Matrix.tensorProduct x (padBlock shz ys) &||| BlockMatrix.zero shx ()


convolveCyclicOneMany ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array     (Shape.Cyclic n) a ->
   Array (sh, Shape.Cyclic n) a ->
   Array (sh, Shape.Cyclic n) a
convolveCyclicOneMany x y = MutArray.runST (do
   let (Shape.Cyclic n) = Array.shape x
   z <- MutVector.new (Matrix.height y, Shape.ZeroBased n) Scalar.zero
   let xExt = replicateCyclic x
   for_ (zip [0..] $ VectorSlice.slicesVector Slice.columns y) $ \(k,yi) ->
      MutMatrix.addTensorProduct 1 yi
         (VectorSlice.sliceVector (Slice.sub (n-1-k) n) xExt)
         (MutMatrix.fromMatrix z)
   return $ MutArray.reshape (Array.shape y) z)

convolveCyclicOneManyViaTable ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array     (Shape.Cyclic n) a ->
   Array (sh, Shape.Cyclic n) a ->
   Array (sh, Shape.Cyclic n) a
convolveCyclicOneManyViaTable x y =
   blockToCyclic (Shape.cyclicSize $ Array.shape x) $
   convolveOneManyViaTable (acyclic x) (acyclicWidth y)


convolveManyMany ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array (sh, Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a ->
   Array (sh, Shape.ZeroBased n) a
convolveManyMany xs ys = MutArray.runST (do
   let (height,widthX) = Array.shape xs
   let         widthY  = Matrix.width ys
   zs <- MutVector.new (height, CShape.convolve widthX widthY) Scalar.zero
   Fold.sequence_ $
      BoxedArray.zipWith3 addConvolve
         (VectorSlice.extractMany Slice.rowArray xs)
         (VectorSlice.slicesVector Slice.rowArray ys)
         (MutVector.slicesVector Slice.rowArray zs)
   return zs)

convolveCyclicManyMany ::
   (Shape.C sh, Eq sh, Integral n, Class.Floating a) =>
   Array (sh, Shape.Cyclic n) a ->
   Array (sh, Shape.Cyclic n) a ->
   Array (sh, Shape.Cyclic n) a
convolveCyclicManyMany xs ys = MutArray.runST (do
   let (height,widthX) = Array.shape xs
--   let         widthY  = Matrix.width ys
   let width = Shape.ZeroBased $ Shape.cyclicSize widthX
   zs <- MutVector.new (height, width) Scalar.zero
   Fold.sequence_ $
      BoxedArray.zipWith3 addConvolveCut
         (fmap acyclic $ VectorSlice.extractMany Slice.rowArray xs)
         (fmap replicateCyclic $ VectorSlice.slicesVector Slice.rowArray ys)
         (MutVector.slicesVector Slice.rowArray zs)
   return $ MutArray.reshape (Array.shape xs) zs)



-- ToDo: subconvolution should be a parameter
{- ToDo:
Handle case of different sizes.
I think we should split the larger vector
such that only operands of equal length are convolved with Karatsuba.
The remaining convolution can be performed with the subconvolution.

Should we have a karatsubaOneMany?
-}
{- |
Size of vectors must match and must be even.

prop> :{
   QC.forAll genVectorShape $ \(Shape.ZeroBased n0) ->
   let shape = Shape.ZeroBased $ n0 + mod n0 2 in
   QC.forAll (genVectorForShape number_ shape) $ \x ->
   QC.forAll (genVectorForShape number_ shape) $ \y ->
   convolve x y =~= karatsuba x y
:}
-}
karatsuba ::
   (Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
karatsuba xs ys =
   let n = Shape.zeroBasedSize $ Array.shape xs in
   let n2 = div n 2 in
   let doubleShape = Shape.ZeroBased n2 ::+ Shape.ZeroBased n2 in
   let (x0,x1) = Array.split $ Array.reshape doubleShape xs in
   let (y0,y1) = Array.split $ Array.reshape doubleShape ys in
   let xy00 = convolve x0 y0 in
   let xy11 = convolve x1 y1 in
   let xym = convolve (x0|+|x1) (y0|+|y1) |-| (xy00 |+| xy11) in
   (xy00 <> Vector.zero (Shape.ZeroBased 1) <> xy11)
   |+|
   (Vector.zero (Shape.ZeroBased n2)
    <>
    xym
    <>
    Vector.zero (Shape.ZeroBased n2))

{- |
Size of vectors must match and must be even.

prop> :{
   QC.forAll genVectorShape $ \(Shape.ZeroBased n0) ->
   let shape = Shape.ZeroBased $ n0 + mod n0 2 in
   QC.forAll (genVectorForShape number_ shape) $ \x ->
   QC.forAll (genVectorForShape number_ shape) $ \y ->
   convolve x y =~= karatsubaViaMutable x y
:}
-}
karatsubaViaMutable ::
   (Integral n, Class.Floating a) =>
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a ->
   Array (Shape.ZeroBased n) a
karatsubaViaMutable xs ys =
   let n = Shape.zeroBasedSize $ Array.shape xs in
   let n2 = div n 2 in
   let doubleShape = Shape.ZeroBased n2 ::+ Shape.ZeroBased n2 in
   let (x0,x1) = Array.split $ Array.reshape doubleShape xs in
   let (y0,y1) = Array.split $ Array.reshape doubleShape ys in
   let xy00 = convolve x0 y0 in
   let xy11 = convolve x1 y1 in
   MutArray.runST (do
      z <-
         fmap (MutArray.reshape (Shape.ZeroBased $ 2*n-1)) $
         MutVector.fromChunk $
            Chunk.simple xy00 <>
            Chunk.singleton Scalar.zero <>
            Chunk.simple xy11
      let zPart = MutVector.sliceVector (Slice.sub n2 (n-1)) z
      MutVector.sub zPart xy00
      MutVector.sub zPart xy11
      addConvolve (x0|+|x1) (y0|+|y1) zPart
      return z)