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

import qualified Numeric.BLAS.Matrix.RowMajor.Block as BlockMatrix
import qualified Numeric.BLAS.Matrix.RowMajor as Matrix
import qualified Numeric.BLAS.Vector as Vector
import qualified Numeric.BLAS.Scalar as Scalar
import qualified Numeric.Netlib.Class as Class
import Numeric.BLAS.Matrix.RowMajor.Block ((&|||))
import Numeric.BLAS.Vector (Vector)
import Numeric.BLAS.Scalar (RealOf)

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Control.Monad (replicateM)
import Control.Applicative (liftA2)
import Foreign.Storable (Storable)
import Data.Array.Comfort.Shape ((::+)((::+)))
import Data.Complex (Complex((:+)))

import qualified Test.QuickCheck as QC
import System.Random (Random)


type ComplexFloat = Complex Float
type ComplexDouble = Complex Double


genReal :: (Class.Real a) => QC.Gen a
genReal = Class.switchReal (QC.choose (-1,1)) (QC.choose (-1,1))

genNumber :: (Class.Floating a) => QC.Gen a
genNumber =
   Class.switchFloating
      genReal genReal
      (liftA2 (:+) genReal genReal)
      (liftA2 (:+) genReal genReal)

-- ToDo: forShape with included Shrink
genVectorShape :: QC.Gen (Shape.ZeroBased Int)
genVectorShape = fmap (Shape.ZeroBased . length) $ QC.listOf1 (pure ())

genVector :: (Storable a) => QC.Gen a -> QC.Gen (Vector (Shape.ZeroBased Int) a)
genVector genElem = fmap Vector.autoFromList $ QC.listOf1 genElem

genShiftedShape :: QC.Gen (Shape.Shifted Int)
genShiftedShape =
   liftA2 Shape.Shifted QC.arbitrary $
   fmap Shape.zeroBasedSize genVectorShape

genIntervalShape :: QC.Gen (Shape.Interval Int)
genIntervalShape = fmap Shape.intervalFromShifted genShiftedShape

genCyclicShape :: QC.Gen (Shape.Cyclic Int)
genCyclicShape = fmap (Shape.Cyclic . length) $ QC.listOf1 (pure ())

genCyclic :: (Storable a) => QC.Gen a -> QC.Gen (Vector (Shape.Cyclic Int) a)
genCyclic genElem =
   fmap (Array.mapShape (\(Shape.ZeroBased n) -> Shape.Cyclic n)) $
   genVector genElem

genVectorForShape :: (Shape.C shape, Storable a) =>
   QC.Gen a -> shape -> QC.Gen (Vector shape a)
genVectorForShape genElem shape =
   fmap (Array.fromList shape) $ replicateM (Shape.size shape) genElem

chooseLogarithmic :: Int -> QC.Gen Int
chooseLogarithmic n =
   fmap (round . (2**)) $
   QC.choose (0, logBase 2 (fromIntegral n :: Double))

forChoose :: (Show n, Random n, Integral n, QC.Testable prop) =>
   (n, n) -> (n -> prop) -> QC.Property
forChoose (lower,upper) =
   QC.forAllShrink
      (QC.choose (lower,upper))
      (\n ->
         filter (lower<) [div (lower+n) 2]
         ++
         if lower<n then [pred n] else [])

takeShape ::
   (Shape.C sh0, Shape.C sh1, Storable a) =>
   sh1 -> Vector sh0 a -> Vector sh1 a
takeShape sh1 =
   Array.takeLeft .
   Array.mapShape
      (\sh0 -> sh1 ::+ Shape.ZeroBased (Shape.size sh0 - Shape.size sh1))

-- ToDo: can make use of conversion (Shape.Cyclic, Shape.ZeroBased) -> Shape.Cyclic
-- ToDo: we can just call padBlockCyclic
upsample :: (Integral n, Class.Floating a) =>
   n -> Vector (Shape.Cyclic n) a -> Vector (Shape.Cyclic n) a
upsample m x =
   let shape = Array.shape x
       n = Shape.cyclicSize shape
   in Array.reshape (Shape.Cyclic (n*m)) $
      BlockMatrix.toMatrix $
      Matrix.singleColumn x
      &|||
      BlockMatrix.zero shape (Shape.ZeroBased (m-1))



infixl 9 &::

(&::) :: f a -> g a -> f a
(&::) = const


infix 4 =~=

-- ToDo: generalize to VectorSlice.C
(=~=) ::
   (Shape.InvIndexed sh, Show sh, Eq sh, Shape.Index sh ~ ix, Show ix) =>
   (Show a, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Vector sh a ->
   Vector sh a ->
   QC.Property
x=~=y =
   QC.counterexample (show x) $
   QC.counterexample (show y) $
   QC.counterexample (show (Vector.argAbsMaximum (Vector.sub x y))) $
   Vector.normInf (Vector.sub x y) < Scalar.selectReal 1e-3 1e-8