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

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.Slice as VectorSlice
import qualified Numeric.BLAS.Vector.Chunk as Chunk
import qualified Numeric.BLAS.Vector.Mutable as MutVector
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.View as View
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.Matrix.RowMajor (Matrix)
import Numeric.BLAS.Vector.Slice ((|+|), (|-|))

import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Storable.Mutable as MutArray
import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)
import Data.Array.Comfort.Shape ((::+)((::+)))
import Data.Complex (Complex)
-- import Data.Tuple.HT (mapFst)

-- import Control.Arrow ((&&&))
-- import Debug.Trace (trace)



-- ToDo: move to comfort-fftw
removeUnitDim ::
   (Shape.C sh, Shape.C cyclic) =>
   Array (sh, (), cyclic) a -> Array (sh, cyclic) a
removeUnitDim = Array.mapShape (\(sh, (), cyclic) -> (sh, cyclic))

addUnitDim ::
   (Shape.C sh, Shape.C cyclic) =>
   Array (sh, cyclic) a -> Array (sh, (), cyclic) a
addUnitDim = Array.mapShape (\(sh, cyclic) -> (sh, (), cyclic))

withExtraUnitDim ::
   (Shape.C sh, Shape.C sh1, Shape.C sh2) =>
   (Array (sh, (), sh1) a -> Array (sh, (), sh2) b) ->
   Array (sh, sh1) a -> Array (sh, sh2) b
withExtraUnitDim f = removeUnitDim . f . addUnitDim



acyclic :: (Integral n) =>
   Array (Shape.Cyclic n) a -> Array (Shape.ZeroBased n) a
acyclic = Array.mapShape (Shape.ZeroBased . Shape.cyclicSize)

acyclicWidth :: (Shape.C height, Integral n) =>
   Matrix height (Shape.Cyclic n) a -> Matrix height (Shape.ZeroBased n) a
acyclicWidth = Matrix.mapWidth (Shape.ZeroBased . Shape.cyclicSize)



splitShape ::
   (Shape.C sh0, Shape.C sh1) =>
   sh1 -> sh0 -> sh1 ::+ Slice.ShapeInt
splitShape sh1 sh0 = sh1 ::+ Shape.ZeroBased (Shape.size sh0 - Shape.size sh1)

takeShape ::
   (Shape.C sh0, Shape.C sh1, Class.Floating a) =>
   sh1 -> Array sh0 a -> VectorSlice.T Layout.Slice sh1 a
takeShape sh1 =
   VectorSlice.sliceVector (Slice.left . Slice.mapShape (splitShape sh1))

takeSubshape ::
   (Shape.C sh0, Shape.C sh1, Class.Floating a) =>
   sh1 -> Array sh0 a -> VectorSlice.T Layout.Subvector sh1 a
takeSubshape sh1 =
   VectorSlice.viewVector (View.left . View.mapShape (splitShape sh1))

_boxedTakeShape ::
   (Shape.C sh0, Shape.C sh1) =>
   sh1 -> BoxedArray.Array sh0 a -> BoxedArray.Array sh1 a
_boxedTakeShape sh1 =
   BoxedArray.takeLeft . BoxedArray.mapShape (splitShape sh1)



padShape ::
   (VectorSlice.C v, Shape.C sh0, Shape.C sh1, Class.Floating a) =>
   sh1 -> v sh0 a -> Array sh1 a
padShape sh1 x =
   Array.reshape sh1 $
   VectorSlice.append x $
   VectorSlice.zero
      (Shape.ZeroBased $ Shape.size sh1 - Shape.size (VectorSlice.shape x))

{-
It is a checked error if the target size is smaller than the input data size.
-}
padCyclic ::
   (VectorSlice.C v, Integral n, Class.Floating a) =>
   n ->
   v (Shape.ZeroBased n) a ->
   Array (Shape.Cyclic n) a
padCyclic n = padShape (Shape.Cyclic n)

spreadCyclic ::
   (Integral n, Class.Floating a) =>
   n ->
   Array (Shape.Interval n) a ->
   Array (Shape.Cyclic n) a
spreadCyclic n x =
   let Shape.Interval from to = Array.shape x in
   let (fromQ, fromR) = divMod from n in
   let (toQ,   toR)   = divMod to n in
   if from>to
   then error "spreadCyclic: from>to"
   else
      if from+n<=to
      then error "spreadCyclic: interval too large"
      else
         Array.reshape (Shape.Cyclic n) $ Chunk.toVector $
         case toQ-fromQ of
            0 ->
               Chunk.zero fromR <>
               Chunk.simple (Array.reshape (Shape.ZeroBased (to-from+1)) x) <>
               Chunk.zero (n-toR-1)
            1 ->
               let split =
                     Shape.ZeroBased (n-fromR) ::+ Shape.ZeroBased (toR+1) in
               let xsplit = VectorSlice.fromVector $ Array.reshape split x in
               Chunk.simple (VectorSlice.takeRight xsplit) <>
               Chunk.zero (fromR-toR-1) <>
               Chunk.simple (VectorSlice.takeLeft xsplit)
            _ -> error "spreadCyclic: panic - broken arithmetic"



{- |
Periodic summation.

I.e. convert @ShapeInt@-Array to a cyclic one
by padding or wrapping around and accumulating.
This is like convolving the signal with an equidistant Dirac pulse train
and keep one cycle.
-}
toCyclic ::
   (Integral n, Class.Floating a) =>
   n -> Array (Shape.ZeroBased n) a -> Array (Shape.Cyclic n) a
toCyclic n x =
   let Shape.ZeroBased m = Array.shape x in
   let (q,r) = divMod m n in
   let (block,tl) =
         VectorSlice.split $ VectorSlice.fromVector $
         Array.reshape
            ((Shape.ZeroBased q, Shape.ZeroBased n) ::+ Shape.ZeroBased r) x in
   if q == 0
   then padCyclic n tl
   else
      Array.reshape (Shape.Cyclic n) $
      MutArray.runST (do
         y <-
            if q == 1
            then
               MutVector.thawSlice $
               -- ToDo: VectorSlice.mapShape snd block
               VectorSlice.slice (Slice.mapShape snd) block
            -- ToDo: generalize sumColumns to submatrix
            -- ToDo: let sumColumns write to Mutable Array
            else MutArray.thaw $ Matrix.sumColumns $ VectorSlice.toVector block
         MutVector.add (MutVector.sliceVector (Slice.take r) y) tl
         return y)


padBlock ::
   (Shape.C sh, Eq sh, Shape.C sh0, Shape.C sh1, Class.Floating a) =>
   sh1 ->
   Array (sh, sh0) a ->
   Array (sh, sh1) a
padBlock sh1 x =
   let (height, sh0) = Array.shape x in
   Array.reshape (height, sh1) $
   BlockMatrix.toMatrix $
   x
   &|||
   BlockMatrix.zero height (Shape.ZeroBased (Shape.size sh1 - Shape.size sh0))

padBlockCyclic ::
   (BlockMatrix.Block block, Integral n, Shape.C sh, Eq sh, Class.Floating a) =>
   n ->
   block (sh, Shape.ZeroBased n) a ->
   Array (sh, Shape.Cyclic n) a
padBlockCyclic n x =
   let (height, Shape.ZeroBased m) = BlockMatrix.shape x in
   Array.reshape (height, Shape.Cyclic n) $
   BlockMatrix.toMatrix $
   x &||| BlockMatrix.zero height (Shape.ZeroBased (n-m))

-- cf. spreadCyclic
spreadBlockCyclic ::
   (Integral n, Shape.C sh, Eq sh, Class.Floating a) =>
   n ->
   Array (sh, Shape.Interval n) a ->
   Array (sh, Shape.Cyclic n) a
spreadBlockCyclic n x =
   let (height, Shape.Interval from to) = Array.shape x in
   let (fromQ, fromR) = divMod from n in
   let (toQ,   toR)   = divMod to n in
   if from>to
   then error "spreadBlockCyclic: from>to"
   else
      if from+n<=to
      then error "spreadBlockCyclic: interval too large"
      else
         case toQ-fromQ of
            0 ->
               Matrix.reshapeWidth (Shape.Cyclic n) $ BlockMatrix.toMatrix $
               BlockMatrix.zero height (Shape.ZeroBased fromR) &|||
               x &|||
               BlockMatrix.zero height (Shape.ZeroBased (n-toR-1))
            1 ->
               Matrix.reshapeWidth (Shape.Cyclic n) $ BlockMatrix.toMatrix $
               let leftShape = Shape.ZeroBased (n-fromR) in
               let rightShape = Shape.ZeroBased (toR+1) in
               takeCenter (leftShape ::+ rightShape ::+ Shape.Zero) x &|||
               BlockMatrix.zero height (Shape.ZeroBased (fromR-toR-1)) &|||
               takeCenter (Shape.Zero ::+ leftShape ::+ rightShape) x
            _ -> error "spreadBlockCyclic: panic - broken arithmetic"


takeCenter ::
   (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 height width a
takeCenter shape x =
   BlockMatrix.submatrix $
   Array.mapShape
      (Subshape.focus (ViewMatrix.left . ViewMatrix.right) .
       Subshape.submatrixFromMatrix) $
   Matrix.reshapeWidth shape x

splitBlocksSideBySide ::
   (Integral n, Shape.C height, Shape.C width, Shape.C left, Shape.C right) =>
   (Class.Floating a) =>
   Array (height, left ::+ (Shape.ZeroBased n, width) ::+ right) a ->
   [Array (height, width) a]
splitBlocksSideBySide x =
   let (left ::+ (Shape.ZeroBased n, width) ::+ right) = Matrix.width x in
   map (BlockMatrix.toMatrix . flip takeCenter x) $
   map
      (\k ->
         (left ::+ (Shape.ZeroBased k, width)) ::+
         width ::+
         ((Shape.ZeroBased (n-k-1), width) ::+ right)) $
   takeWhile (<n) $ iterate (1+) 0

blockToCyclic ::
   (Integral n, Shape.C sh, Eq sh, Class.Floating a) =>
   n -> Array (sh, Shape.ZeroBased n) a -> Array (sh, Shape.Cyclic n) a
blockToCyclic n x =
   let (height, Shape.ZeroBased m) = Array.shape x in
   let (q,r) = divMod m n in
   let blockShape = (Shape.ZeroBased q, Shape.Cyclic n) in
   let tl = takeCenter (blockShape ::+ Shape.ZeroBased r ::+ Shape.Zero) x in
   foldl VectorSlice.add (padBlockCyclic n tl) $
   splitBlocksSideBySide $
   Array.reshape (height, Shape.Zero ::+ blockShape ::+ Shape.ZeroBased r) x

-- cf. spreadBlockCyclic
intervalBlockToCyclic ::
   (Integral n, Shape.C sh, Eq sh, Class.Floating a) =>
   n -> Array (sh, Shape.Interval n) a -> Array (sh, Shape.Cyclic n) a
intervalBlockToCyclic n x =
   let (height, Shape.Interval from to) = Array.shape x in
   let (fromQ, fromR) = divMod from n in
   let (toQ,   toR)   = divMod to n in
   if from>to
   then error "intervalBlockToCyclic: from>to"
   else
      case toQ-fromQ of
         0 ->
            Matrix.reshapeWidth (Shape.Cyclic n) $ BlockMatrix.toMatrix $
            BlockMatrix.zero height (Shape.ZeroBased fromR) &|||
            x &|||
            BlockMatrix.zero height (Shape.ZeroBased (n-toR-1))
         _ ->
            let blockShape = (Shape.ZeroBased (toQ-fromQ-1), Shape.Cyclic n) in
            let leftShape = Shape.ZeroBased (n-fromR) in
            let rightShape = Shape.ZeroBased (toR+1) in
            let hd =
                  flip takeCenter x $
                  Shape.Zero ::+ leftShape ::+ (blockShape ::+ rightShape) in
            let tl =
                  flip takeCenter x $
                  (leftShape ::+ blockShape) ::+ rightShape ::+ Shape.Zero in
            let gap = fromR-toR-1 in
            let remd =
                  if gap >= 0
                  then
                     Matrix.reshapeWidth (Shape.Cyclic n) $
                     BlockMatrix.toMatrix $
                     tl
                     &|||
                     BlockMatrix.zero height (Shape.ZeroBased gap)
                     &|||
                     hd
                  else
                     (Matrix.reshapeWidth (Shape.Cyclic n) $
                      BlockMatrix.toMatrix $
                        BlockMatrix.zero height (Shape.ZeroBased fromR) &||| hd)
                     |+|
                     padBlockCyclic n tl
                  in
            foldl VectorSlice.add remd $
            splitBlocksSideBySide $
            Array.reshape (height, leftShape ::+ blockShape ::+ rightShape) x

halfBlockToHalf ::
   (Integral n, Shape.C sh, Eq sh, Class.Real a) =>
   n ->
   Array (sh, Spectrum.Half n) (Complex a) ->
   Array (sh, Spectrum.Half n) (Complex a)
halfBlockToHalf m x =
--   let (height, width@(Spectrum.Half mcut)) = Array.shape x in
   let width@(Spectrum.Half mcut) = Matrix.width x in
   let mcutHalf = fromIntegral $ Shape.size width in
   let halfShape = Spectrum.Half m in
   if mcut <= m - mod (m+1) 2
   then padBlock halfShape x
   else
      let xzb = Matrix.reshapeWidth (Shape.ZeroBased mcutHalf) x in
      if mcutHalf <= m
      then foldBlockToHalf m xzb
      else foldBlockToHalf m $
           Matrix.reshapeWidth (Shape.ZeroBased m) $
           BlockMatrix.toMatrix $
           let takeRealColumn ix =
                  VectorSlice.realPart .
                     VectorSlice.sliceVector (Slice.column ix) in
           (\block ->
               -- ToDo: make this more efficient by a mutable update
               (Matrix.singleColumn $ Vector.fromReal $
                  VectorSlice.scale 2 (takeRealColumn 0 block)
                  |-|
                  takeRealColumn 0 x)
               &|||
               -- ToDo: takeTransposed  without transposed would save a copy here
               Matrix.takeRight
                  (Matrix.reshapeWidth (() ::+ Shape.ZeroBased (m-1)) block)) $
           blockToCyclic m xzb

-- ToDo: comfort-blas:Matrix.mapRows - apply a vector function to each row
foldBlockToHalf ::
   (Integral n, Shape.C sh, Eq sh, Class.Floating a) =>
   n -> Array (sh, Shape.ZeroBased n) a -> Array (sh, Spectrum.Half n) a
foldBlockToHalf m x =
   let (Shape.ZeroBased mcutHalf) = Matrix.width x in
   let halfShape = Spectrum.Half m in
   let mHalf = fromIntegral $ Shape.size halfShape in
   Matrix.fromRowArray halfShape $
   fmap (\row ->
      VectorSlice.slice
         (Slice.left . Slice.mapShape (splitShape halfShape)) row
      |+|
      Array.reshape halfShape
         (VectorSlice.append
            (VectorSlice.zero (Shape.ZeroBased (m+1-mcutHalf)))
            (VectorSlice.conjugate $ VectorSlice.reverse $
               VectorSlice.slice
                  (Slice.sub (m+1-mHalf) (mHalf+mcutHalf-(m+1)))
                  row))) $
   VectorSlice.slicesVector Slice.rowArray x


assembleInterval1 ::
   (Integral n) =>
   Shape.ZeroBased n -> Shape.ZeroBased n -> Shape.Interval n
assembleInterval1 (Shape.ZeroBased m) (Shape.ZeroBased n) =
   Shape.Interval (1-m) n

normalizeInterval :: (Integral n) => n -> Shape.Interval n -> Shape.Interval n
normalizeInterval n (Shape.Interval from0 to0) =
   let to = mod to0 n in Shape.Interval (from0 + to-to0) to

cyclicTakeNormalized ::
   (Integral n, Class.Floating a) =>
   Shape.Interval n -> Array (Shape.Cyclic n) a -> Array (Shape.Interval n) a
cyclicTakeNormalized rng@(Shape.Interval from to) x =
   let Shape.Cyclic n = Array.shape x in
   -- let y = Array.mapShape Shape.zeroBasedFromCyclic x in
   let y = Array.reshape (Shape.ZeroBased n) x in
   if from<0
   then
      Array.reshape rng $
      VectorSlice.append
         (VectorSlice.sliceVector (Slice.drop (n+from)) y)
         (VectorSlice.sliceVector (Slice.take (to+1)) y)
   else
      VectorSlice.extract (Slice.subInterval rng) $
      Array.mapShape Shape.intervalFromZeroBased y

cyclicTake ::
   (Integral n, Class.Floating a) =>
   Shape.Interval n -> Array (Shape.Cyclic n) a -> Array (Shape.Interval n) a
cyclicTake rng x =
   Array.reshape rng $
   cyclicTakeNormalized
      (normalizeInterval (Shape.cyclicSize $ Array.shape x) rng)
      x