{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.BatchNormalisation
Description : Batch normalization layer
Copyright   : (c) Theo Charalambous, 2020
License     : BSD2
Stability   : experimental

The layer follows the implementation as described in the paper:
Sergey Ioffe and Christian Szegedy. Batch normalization:  Accelerating deep network training by reducing internal covariate shift
<http://arxiv.org/abs/1502.03167>
-}

module Grenade.Layers.BatchNormalisation
  (
  -- * Layer Definition
    BatchNorm (..)
  , BatchNormParams(..)

  -- * Differentation Types
  , BatchNormTape(..)
  , BatchNormGrad(..)

  -- * Helper functions
  , initBatchNorm
  )
where

import           Control.DeepSeq
import           Data.Kind                         (Type)
import           Data.List                         (transpose, zipWith5)
import           Data.Maybe                        (fromJust)
import           Data.Proxy
import           Data.Serialize
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra             as LA
import           Numeric.LinearAlgebra.Static      hiding (Seed)

import           Grenade.Core
import           Grenade.Layers.Internal.BatchNorm
import           Grenade.Layers.Internal.Update
import           Grenade.Onnx
import           Grenade.Types
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

import           Lens.Micro

data BatchNormTape :: Nat -- The number of channels of the tensor
                   -> Nat -- The number of rows of the tensor
                   -> Nat -- The number of columns of the tensor
                   -> Type where
  -- | there is no back propogation when using network for testing, so
  --   so the tape for automatic differentiation is empty
  TestBatchNormTape  :: ( KnownNat channels
                     , KnownNat rows
                     , KnownNat columns)
                     => BatchNormTape channels rows columns -- ^ empty tape

  -- | prevents having to calculate information twice during forward and backwards pass
  TrainBatchNormTape :: ( KnownNat channels
                     , KnownNat rows
                     , KnownNat columns)
                     => [R (channels * rows * columns)]  -- ^ xnorm
                     -> R channels                       -- ^ std
                     -> R channels                       -- ^ running mean
                     -> R channels                       -- ^ running variance
                     -> BatchNormTape channels rows columns

data BatchNorm :: Nat -- The number of channels of the tensor
               -> Nat -- The number of rows of the tensor
               -> Nat -- The number of columns of the tensor
               -> Nat -- momentum
               -> Type where
  BatchNorm :: ( KnownNat channels
               , KnownNat rows
               , KnownNat columns
               , KnownNat momentum)
            => Bool                                 -- ^ True: training, False: testing
            -> BatchNormParams channels             -- ^ gamma and beta
            -> R channels                           -- ^ running mean
            -> R channels                           -- ^ running variance
            -> RealNum                              -- ^ epsilon (used for numerical stability)
            -> ListStore (BatchNormParams channels) -- ^ momentum store
            -> BatchNorm channels rows columns momentum

data BatchNormParams :: Nat -> Type where
  BatchNormParams :: KnownNat channels
                  => R channels -- gamma
                  -> R channels -- beta
                  -> BatchNormParams channels

data BatchNormGrad :: Nat -- The number of channels of the tensor
                   -> Nat -- The number of rows of the tensor
                   -> Nat -- The number of columns of the tensor
                   -> Type where
  BatchNormGrad :: ( KnownNat channels
                   , KnownNat rows
                   , KnownNat columns)
            => R channels -- ^ running mean
            -> R channels -- ^ running variance
            -> R channels -- ^ derivative of loss wrt gamma
            -> R channels -- ^ derivative of loss wrt beta
            -> BatchNormGrad channels rows columns

-- | NFData instances
instance NFData (BatchNormTape channels rows columns) where
  rnf TestBatchNormTape = ()
  rnf (TrainBatchNormTape xnorm std mean var) = rnf xnorm `seq` rnf std `seq` rnf mean `seq` rnf var

instance NFData (BatchNormParams flattenSize) where
  rnf (BatchNormParams gamma beta)
    = rnf gamma `seq` rnf beta

instance NFData (BatchNorm channels rows columns momentum) where
  rnf (BatchNorm training bnparams mean var eps store)
    = rnf training `seq` rnf bnparams `seq` rnf mean `seq` rnf var `seq` rnf eps `seq` rnf store

-- | Show instance

instance Show (BatchNorm channels rows columns momentum) where
  show _ = "Batch Normalization"

-- Serialize instances
instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => Serialize (BatchNorm channels rows columns momentum) where

  put (BatchNorm training bnparams mean var ε store) = do
    put training
    put bnparams
    putListOf put . LA.toList . extract $ mean
    putListOf put . LA.toList . extract $ var
    put ε
    put store

  get = do
    training <- get
    bnparams <- get
    mean     <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
    var      <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
    ε        <- get
    store    <- get

    return $ BatchNorm training bnparams mean var ε store

instance KnownNat flattenSize => Serialize (BatchNormParams flattenSize) where
  put (BatchNormParams gamma beta) = do
    putListOf put . LA.toList . extract $ gamma
    putListOf put . LA.toList . extract $ beta

  get = do
    gamma <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
    beta  <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get

    return $ BatchNormParams gamma beta

-- | Neural network operations
instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat mom, KnownNat (channels * rows * columns))
  => UpdateLayer (BatchNorm channels rows columns mom) where

  type Gradient (BatchNorm channels rows columns mom) = BatchNormGrad channels rows columns
  type MomentumStore (BatchNorm channels rows columns mom) = ListStore (BatchNormParams channels)

  reduceGradient = undefined

  runUpdate opt@OptSGD{} x@(BatchNorm training (BatchNormParams oldGamma oldBeta) _ _ ε store) (BatchNormGrad runningMean runningVar dGamma dBeta)
    = let BatchNormParams oldGammaMomentum oldBetaMomentum = getData opt x store
          VectorResultSGD newGamma newGammaMomentum        = descendVector opt (VectorValuesSGD oldGamma dGamma oldGammaMomentum)
          VectorResultSGD newBeta  newBetaMomentum         = descendVector opt (VectorValuesSGD oldBeta dBeta oldBetaMomentum)
          newStore                                         = setData opt x store (BatchNormParams newGammaMomentum newBetaMomentum)
      in  BatchNorm training (BatchNormParams newGamma newBeta) runningMean runningVar ε newStore

  runUpdate opt@OptAdam{} x@(BatchNorm training (BatchNormParams oldGamma oldBeta) _ _ ε store) (BatchNormGrad runningMean runningVar dGamma dBeta)
    = let [BatchNormParams oldMGamma oldMBeta, BatchNormParams oldVGamma oldVBeta] = getData opt x store
          VectorResultAdam newGamma newMGamma newVGamma                            = descendVector opt (VectorValuesAdam (getStep store) oldGamma dGamma oldMGamma oldVGamma)
          VectorResultAdam newBeta  newMBeta  newVBeta                             = descendVector opt (VectorValuesAdam (getStep store) oldBeta  dBeta  oldMBeta  oldVBeta)
          newStore                                                                 = setData opt x store [BatchNormParams newMGamma newMBeta, BatchNormParams newVGamma newVBeta]
      in  BatchNorm training (BatchNormParams newGamma newBeta) runningMean runningVar ε newStore

  runSettingsUpdate NetworkSettings{trainingActive=training} (BatchNorm _ bnparams mean var ε store) = BatchNorm training bnparams mean var ε store

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat mom)
  => LayerOptimizerData (BatchNorm channels rows columns mom) (Optimizer 'SGD) where

  type MomentumDataType (BatchNorm channels rows columns mom) (Optimizer 'SGD) = BatchNormParams channels
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = BatchNormParams (konst 0) (konst 0)

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat mom)
  => LayerOptimizerData (BatchNorm channels rows columns mom) (Optimizer 'Adam) where

  type MomentumDataType (BatchNorm channels rows columns mom) (Optimizer 'Adam) = BatchNormParams channels
  type MomentumExpOptResult (BatchNorm channels rows columns mom) (Optimizer 'Adam) = [BatchNormParams channels]
  getData     = getListStore
  setData     = setListStore
  newData _ _ = BatchNormParams (konst 0) (konst 0)

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => RandomLayer (BatchNorm channels rows columns momentum) where
  createRandomWith _ _ = pure initBatchNorm

-- | Initialize a batch norm layer with gamma set to zero for each channel and beta
--   set to one for each channel. Running mean is 0 and running var is 1. The default
--   value for epsilon is 0.00001
initBatchNorm :: forall channels rows columns momentum.
  (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => BatchNorm channels rows columns momentum
initBatchNorm =
  let ch     = fromIntegral $ natVal (Proxy :: Proxy channels)
      zeroes = replicate ch 0
      ones   = replicate ch 1
      gamma  = vector ones   :: R channels
      beta   = vector zeroes :: R channels
      mean   = vector zeroes :: R channels
      var    = vector ones   :: R channels
      ε      = 0.00001
  in BatchNorm True (BatchNormParams gamma beta) mean var ε mkListStore

instance (KnownNat rows, KnownNat momentum)
  => Layer (BatchNorm 1 1 rows momentum) ('D1 rows) ('D1 rows) where

  type Tape (BatchNorm 1 1 rows momentum) ('D1 rows) ('D1 rows) = BatchNormTape 1 1 rows

  runForwards (BatchNorm True _ _ _ _ _) _
    = error "Cannot train use batch size of 1 with BatchNorm layer during training"

  runForwards (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S1D x)
    = let gamma'       = extract gamma LA.! 0
          beta'        = extract beta LA.! 0
          runningMean' = extract runningMean LA.! 0
          runningVar'  = extract runningVar LA.! 0

          std = sqrt $ runningVar' + ε

          y = dvmap (\x -> ((x - runningMean') / std ) * gamma' + beta') x

      in (TestBatchNormTape, S1D y)

  runBatchForwards bn@(BatchNorm False _ _ _ _ _) xs
    = let outs = map (snd . runForwards bn) xs
      in ([TestBatchNormTape], outs)

  runBatchForwards (BatchNorm True (BatchNormParams gamma beta) runningMean runningVar ε _) xs
    = let [m]             = vectorToList runningMean :: [RealNum]
          [v]             = vectorToList runningVar  :: [RealNum]
          [g]             = vectorToList gamma       :: [RealNum]
          [b]             = vectorToList beta        :: [RealNum]
          mom             = (/ 100) $ fromIntegral $ natVal (Proxy :: Proxy momentum)

          xs'             = map extractV xs

          sample_mean     = batchNormMean xs'                 :: RealNum
          sample_var      = batchNormVariance xs'             :: RealNum

          m'              = mom * m + (1 - mom) * sample_mean :: RealNum
          v'              = mom * v + (1 - mom) * sample_var  :: RealNum
          std             = sqrt $ sample_var + ε             :: RealNum

          x_extracted     = map (\(S1D x) -> x) xs
          x_normalised    = map (dvmap (\a -> (a - sample_mean) / std)) x_extracted
          scaledShifted   = map (dvmap (\a -> g * a + b)) x_normalised
          out             = map S1D scaledShifted

          stdV            = listToVector [std] :: R 1
          runningMeanV    = listToVector [m']  :: R 1
          runningVarV     = listToVector [v']  :: R 1

      in ([TrainBatchNormTape x_normalised stdV runningMeanV runningVarV], out)

  runBatchBackwards (BatchNorm True (BatchNormParams _ _) _ _ _ _) _ _
    = undefined

  runBatchBackwards _ _ _
    = undefined

instance (KnownNat rows, KnownNat columns, KnownNat momentum)
  => Layer (BatchNorm 1 rows columns momentum) ('D2 rows columns) ('D2 rows columns) where

  type Tape (BatchNorm 1 rows columns momentum) ('D2 rows columns) ('D2 rows columns) = BatchNormTape 1 rows columns

  runForwards (BatchNorm True _ _ _ _ _) _
    = error "Cannot train use batch size of 1 with BatchNorm layer during training"

  runForwards (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S2D x)
    = let rows     = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns  = fromIntegral $ natVal (Proxy :: Proxy columns)

          gamma'       = extract gamma
          beta'        = extract beta
          runningMean' = extract runningMean
          runningVar'  = extract runningVar
          mat          = extract x

          y  = batchnorm 1 rows columns ε mat gamma' beta' runningMean' runningVar'
          y' = fromJust . create $ y

      in (TestBatchNormTape, S2D y')

  runBatchForwards bn@(BatchNorm False _ _ _ _ _) xs
    = let outs = map (snd . runForwards bn) xs
      in ([TestBatchNormTape], outs)

  runBatchForwards (BatchNorm True (BatchNormParams gamma beta) runningMean runningVar ε _) xs
    = let [m]             = vectorToList runningMean :: [RealNum]
          [v]             = vectorToList runningVar  :: [RealNum]
          [g]             = vectorToList gamma       :: [RealNum]
          [b]             = vectorToList beta        :: [RealNum]
          mom             = (/ 100) $ fromIntegral $ natVal (Proxy :: Proxy momentum)

          xs'             = map (sflatten . extractM2D) xs

          sample_mean     = batchNormMean xs'                 :: RealNum
          sample_var      = batchNormVariance xs'             :: RealNum

          m'              = mom * m + (1 - mom) * sample_mean :: RealNum
          v'              = mom * v + (1 - mom) * sample_var  :: RealNum
          std             = sqrt $ sample_var + ε             :: RealNum

          x_extracted     = map (\(S2D x) -> x) xs
          x_normalised    = map (dmmap (\a -> (a - sample_mean) / std)) x_extracted
          scaledShifted   = map (dmmap (\a -> g * a + b)) x_normalised
          out             = map S2D scaledShifted

          x_normalised'   = map sflatten x_normalised
          stdV            = listToVector [std] :: R 1
          runningMeanV    = listToVector [m']  :: R 1
          runningVarV     = listToVector [v']  :: R 1

      in ([TrainBatchNormTape x_normalised' stdV runningMeanV runningVarV], out)

  runBatchBackwards (BatchNorm True (BatchNormParams _ _) _ _ _ _) _ _
    = undefined

  runBatchBackwards _ _ _
    = undefined

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum)
  => Layer (BatchNorm channels rows columns momentum) ('D3 rows columns channels) ('D3 rows columns channels) where

  type Tape (BatchNorm channels rows columns momentum) ('D3 rows columns channels) ('D3 rows columns channels) = BatchNormTape channels rows columns

  runForwards (BatchNorm True _ _ _ _ _) _
    = error "Cannot train use batch size of 1 with BatchNorm layer during training"

  runForwards (BatchNorm False (BatchNormParams gamma beta) runningMean runningVar ε _) (S3D x)
    = let rows     = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns  = fromIntegral $ natVal (Proxy :: Proxy columns)
          channels = fromIntegral $ natVal (Proxy :: Proxy channels)

          gamma'       = extract gamma
          beta'        = extract beta
          runningMean' = extract runningMean
          runningVar'  = extract runningVar
          mat          = extract x

          y  = batchnorm channels rows columns ε mat gamma' beta' runningMean' runningVar'
          y' = fromJust . create $ y

      in (TestBatchNormTape, S3D y')

  runBatchForwards bn@(BatchNorm False _ _ _ _ _) xs
    = let outs = map (snd . runForwards bn) xs
      in ([TestBatchNormTape], outs)

  runBatchForwards (BatchNorm True (BatchNormParams gamma beta) runningMean runningVar ε _) xs
    = let ms     = vectorToList runningMean
          vs     = vectorToList runningVar
          gs     = vectorToList gamma
          bs     = vectorToList beta

          cs     = map splitChannels xs :: [[S ('D2 rows columns)]]
          cs'    = transpose cs

          f c g b m v = let gs' = listToVector [g] :: R 1
                            bs' = listToVector [b] :: R 1
                            ms' = listToVector [m] :: R 1
                            vs' = listToVector [v] :: R 1
                            bn' = BatchNorm True (BatchNormParams gs' bs') ms' vs' ε undefined :: BatchNorm 1 rows columns momentum
                        in  runBatchForwards bn' c

          (tapes, outs) = unzip $ zipWith5 f cs' gs bs ms vs
          outs' = transpose outs
      in (combineTapes tapes, map combineChannels outs')
    where
      combineTapes :: [[BatchNormTape 1 i j]] -> [BatchNormTape k i j]
      combineTapes = undefined

  runBatchBackwards (BatchNorm True (BatchNormParams _ _) _ _ _ _) _ _
    = undefined

  runBatchBackwards _ _ _ = undefined

instance OnnxOperator (BatchNorm channels rows columns momentum) where
  onnxOpTypeNames _ = ["BatchNormalization"]

instance (KnownNat channels, KnownNat rows, KnownNat columns, KnownNat momentum) => OnnxLoadable (BatchNorm channels rows columns momentum) where
  loadOnnxNode inits node = case node ^. #input of
    [_, scale, b, mean, var] -> do
      epsilon     <- readFloatAttributeToRealNum "epsilon" node
      loadedScale <- readInitializerVector inits scale
      loadedB     <- readInitializerVector inits b
      loadedMean  <- readInitializerVector inits mean
      loadedVar   <- readInitializerVector inits var

      return $ BatchNorm False (BatchNormParams loadedScale loadedB) loadedMean loadedVar epsilon mkListStore
    _               -> onnxIncorrectNumberOfInputs


{-# LANGUAGE DeriveFunctor                            #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE NoStarIsType                             #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilyDependencies                   #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Numeric.Wavelet.Discrete (
    DWavelet(..)
  , DWT(..)
  , cdwt
  , dwt
  , idwt
  , haar
    -- DWD(..)
  -- , dwdApprox, dwdDetail
  -- , flattenDWD
  -- , denseDWD
  -- , haar
  -- , unHaar
  ) where

import           Data.Complex
import           Data.Finite
import           Data.Finite.Internal
import           Data.Kind
import           Data.Proxy
import           Data.Vector.Generic.Sized    (Vector)
import           GHC.Generics
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Numeric.Wavelet.Internal.FFT
import qualified Data.Vector                  as UV
import qualified Data.Vector.Generic          as UVG
import qualified Data.Vector.Generic.Sized    as VG
import qualified Data.Vector.Sized            as V

data DWavelet v l a = DW
    { dwHiDecomp :: VG.Vector v l a      -- ^ High-pass decomposition filter
    , dwLoDecomp :: VG.Vector v l a      -- ^ Low-pass decomposition filter
    , dwHiRecon  :: VG.Vector v l a      -- ^ High-pass reconstruction filter
    , dwLoRecon  :: VG.Vector v l a      -- ^ Low-pass reconstruction filter
    }
  deriving (Show, Eq, Ord, Functor, Generic)

data DWT v n a = DWT
    { dwtApprox :: VG.Vector v n a
    , dwtDetail :: VG.Vector v n a
    }
  deriving (Show, Eq, Ord, Functor, Generic)

cdwt
    :: (UVG.Vector v (Complex a), KnownNat l, KnownNat n, 1 <= l, 1 <= n, FFTWReal a)
    => DWavelet  v l (Complex a)
    -> VG.Vector v n (Complex a)
    -> DWT v ((n + l - 1) `Div` 2) (Complex a)
cdwt DW{..} x = DWT{..}
  where
    dwtApprox = downsamp $ convolve x dwLoDecomp
    dwtDetail = downsamp $ convolve x dwHiDecomp

dwt
    :: (UVG.Vector v (Complex a), UVG.Vector v a, KnownNat l, KnownNat n, 1 <= l, 1 <= n, FFTWReal a)
    => DWavelet  v l a
    -> VG.Vector v n a
    -> DWT v ((n + l - 1) `Div` 2) a
dwt DW{..} x = DWT{..}
  where
    dwtApprox = downsamp $ rconvolve x dwLoDecomp
    dwtDetail = downsamp $ rconvolve x dwHiDecomp

idwt
    :: forall v l n a.
     ( UVG.Vector v (Complex a)
     , UVG.Vector v a
     , KnownNat l
     , KnownNat n
     , FFTWReal a
     , Div ((((Div ((n + l) - 1) 2 * 2) + l) - 1) - n) 2 <= ((((Div ((n + l) - 1) 2 * 2) + l) - 1) - n)
     )
    => DWavelet  v l a
    -> DWT v ((n + l - 1) `Div` 2) a
    -> VG.Vector v n a      -- todo: remove the choice
idwt DW{..} DWT{..} = VG.slice @_ @((((n+l-1)`Div`2)*2+l-1-n)`Div`2) @n @(((n+l-1)`Div`2)*2+l-1-n-((((n+l-1)`Div`2)*2+l-1-n)`Div`2)) Proxy $
                VG.zipWith (+) x y
  where
    x = rconvolve (upsamp dwtApprox) dwLoRecon
    y = rconvolve (upsamp dwtDetail) dwHiRecon

downsamp :: (UVG.Vector v a, KnownNat n) => VG.Vector v n a -> VG.Vector v (n `Div` 2) a
downsamp v = VG.generate $ \i -> v `VG.index` doubleFinite i
  where
    doubleFinite :: Finite (n `Div` 2) -> Finite n
    doubleFinite (Finite x) = Finite (x * 2 + 1)

upsamp :: (UVG.Vector v a, KnownNat n, Num a) => VG.Vector v n a -> VG.Vector v (n * 2) a
upsamp v = VG.generate $ \i -> case separateProduct @2 i of
    (j, k)
      | j == 0    -> v `VG.index` k
      | otherwise -> 0

-- upsamp :: forall v n a. (UVG.Vector v a, KnownNat n, Num a) => VG.Vector v (n `Div` 2) a -> VG.Vector v n a
-- upsamp v = VG.generate $ \(Finite i) -> case i `divMod` 2 of
--     (d, m)
--       | m == 0    -> v `VG.index` Finite d
--       | otherwise -> 0

    -- let (j, k) = case separateProduct @2 i
    -- in  if j == 0
    --       then

-- unhaarPass
--     :: (UVG.Vector v a, KnownNat n, Num a)
--     => Vector v (2 ^ n) a
--     -> Vector v (2 ^ n) a
--     -> Vector v (2 ^ (n + 1)) a
-- unhaarPass app det = VG.generate $ \i ->
--     let (j, k) = separateProduct @2 i
--         combiner
--           | j == 0    = (+)
--           | otherwise = (-)
--     in  combiner (app `VG.index` k) (det `VG.index` k)

haar :: (UVG.Vector v a, Floating a) => DWavelet v 2 a
haar = DW{..}
  where
    srtot = sqrt 2 / 2
    dwHiDecomp = VG.fromTuple (-srtot, srtot)
    dwLoDecomp = VG.fromTuple ( srtot, srtot)
    dwHiRecon  = VG.fromTuple ( srtot,-srtot)
    dwLoRecon  = VG.fromTuple ( srtot, srtot)

-- data DWD :: (Type -> Type) -> Nat -> Type -> Type where
--     DWDZ :: { _dwdDetailZ :: !a
--            , _dwdApprox  :: !a
--            } -> DWD v 0 a
--     DWDS :: { _dwdDetailS :: !(Vector  v (2 ^ (n + 1)) a)
--            , _dwdNext    :: !(DWD v n             a)
--            } -> DWD v (n + 1) a

-- deriving instance (Show (v a), Show a) => Show (DWD v n a)

-- dwdDetail :: UVG.Vector v a => DWD v n a -> Vector v (2 ^ n) a
-- dwdDetail = \case
--     DWDZ d _ -> VG.singleton d
--     DWDS d _ -> d

-- dwdApprox :: DWD v n a -> a
-- dwdApprox = \case
--     DWDZ _ a -> a
--     DWDS _ a -> dwdApprox a

-- testDec :: DWD UV.Vector 4 Double
-- testDec = DWDS (0 :: V.Vector 16 Double)
--         . DWDS (0 :: V.Vector  8 Double)
--         . DWDS (0 :: V.Vector  4 Double)
--         . DWDS (0 :: V.Vector  2 Double)
--         $ DWDZ 0 0

-- flattenDWD :: UVG.Vector v a => DWD v n a -> Vector v (2 ^ (n + 1)) a
-- flattenDWD = \case
--     DWDZ d  a  -> VG.fromTuple (a, d)
--     -- whoops this is O(n^2)
--     DWDS ds as -> flattenDWD as VG.++ ds

-- denseDWD
--     :: forall v n a. (KnownNat n, UVG.Vector v a)
--     => DWD v n a
--     -> V.Vector (n + 1) (Vector v (2 ^ n) a)
-- denseDWD = \case
--     DWDZ d  _  -> VG.singleton $ VG.singleton d
--     -- whoops this is O(n^3) or something
--     DWDS ds as -> (densify @_ @_ @(n - 1) <$> denseDWD as) `VG.snoc` ds

-- densify
--     :: (UVG.Vector v a, KnownNat n)
--     => Vector v (2 ^ n) a
--     -> Vector v (2 ^ (n + 1)) a
-- densify xs = VG.generate $ \i -> xs `VG.index` snd (separateProduct @2 i)

-- haar
--     :: forall v n a. (UVG.Vector v a, KnownNat n, Fractional a)
--     => Vector  v (2 ^ (n + 1)) a
--     -> DWD v n a
-- haar xs = case pNat (Proxy @n) of
--     PZ -> let x = xs `VG.index` 0
--               y = xs `VG.index` 1
--           in  DWDZ ((x - y) / 2) ((x + y) / 2)
--     PS -> let (a, d) = haarPass xs
--           in  DWDS d (haar a)

-- unHaar
--     :: (UVG.Vector v a, KnownNat n, Num a)
--     => DWD v n a
--     -> Vector v (2 ^ (n + 1)) a
-- unHaar = \case
--     DWDZ d  a  -> VG.fromTuple (a + d, a - d)
--     DWDS ds as -> unhaarPass (unHaar as) ds

-- haarPass
--     :: (UVG.Vector v a, KnownNat n, Fractional a)
--     => Vector v (2 ^ (n + 1)) a
--     -> (Vector v (2 ^ n) a, Vector v (2 ^ n) a)
-- haarPass xs = (app, det)
--   where
--     app = VG.generate $ \i -> ( xs `VG.index` combineProduct (0, i)
--                               + xs `VG.index` combineProduct (1, i)
--                               ) / 2
--     det = VG.generate $ \i -> ( xs `VG.index` combineProduct (0, i)
--                               - xs `VG.index` combineProduct (1, i)
--                               ) / 2

-- unhaarPass
--     :: (UVG.Vector v a, KnownNat n, Num a)
--     => Vector v (2 ^ n) a
--     -> Vector v (2 ^ n) a
--     -> Vector v (2 ^ (n + 1)) a
-- unhaarPass app det = VG.generate $ \i ->
--     let (j, k) = separateProduct @2 i
--         combiner
--           | j == 0    = (+)
--           | otherwise = (-)
--     in  combiner (app `VG.index` k) (det `VG.index` k)


{-# language BangPatterns        #-}
{-# language LambdaCase          #-}
{-# language NoImplicitPrelude   #-}
{-# language ScopedTypeVariables #-}

-- | This module exposes functions for performing
--   Fast Fourier Transform (FFT) and Inverse Fast Fourier Transform (IFFT)
--   over 'Contiguous' data structures.
module Data.Primitive.Contiguous.FFT
  ( fft
  , ifft
  , mfft
  ) where

import qualified Prelude

import Control.Applicative (pure)
import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad(..))
import Control.Monad.ST (runST)
import Data.Bits (shiftR,shiftL,(.&.),(.|.))
import Data.Bool (Bool,otherwise)
import Data.Complex (Complex(..),conjugate)
import Data.Eq (Eq(..))
import Data.Function (($),(.))
import Data.Ord (Ord(..))
import Data.Primitive.Contiguous (Contiguous,Element,Mutable)
import Data.Semiring (negate,(+),(*),(-))
import GHC.Exts
import GHC.Real ((/))
import qualified Data.Primitive.Contiguous as Contiguous

{-# RULES
"fft/ifft" forall x. fft (ifft x) = x
"ifft/fft" forall x. ifft (fft x) = x
  #-}

-- | Radix-2 decimation-in-time fast Fourier Transform.
--   The given array must have a length that is a power of two.
fft :: forall arr. (Contiguous arr, Element arr (Complex Double))
  => arr (Complex Double)
  -> arr (Complex Double)
{-# inlinable [1] fft #-}
fft arr = if arrOK arr
  then runST $ do {
      marr <- copyWhole arr
    ; mfft marr
    ; Contiguous.unsafeFreeze marr
  }
  else Prelude.error "Data.Primitive.Contiguous.FFT.fft: bad array length"
-- | Inverse fast Fourier transform.
ifft :: forall arr. (Contiguous arr, Element arr (Complex Double))
  => arr (Complex Double)
  -> arr (Complex Double)
{-# inlinable [1] ifft #-}
ifft arr = if arrOK arr
  then
    let lenComplex = intToComplexDouble (Contiguous.size arr)
    in cmap ((/lenComplex) . conjugate) . fft . cmap conjugate $ arr
  else Prelude.error "Data.Primitive.Contiguous.FFT.ifft: bad vector length"

copyWhole :: forall arr m a. (PrimMonad m, Contiguous arr, Element arr a)
  => arr a
  -> m (Mutable arr (PrimState m) a)
{-# inline copyWhole #-}
copyWhole arr = do
  let len = Contiguous.size arr
  marr <- Contiguous.new len
  Contiguous.copy marr 0 arr 0 len
  pure marr

arrOK :: forall arr a. (Contiguous arr, Element arr a)
  => arr a
  -> Bool
{-# inline arrOK #-}
arrOK arr =
  let n = Contiguous.size arr
  in (1 `shiftL` log2 n) == n

-- | Radix-2 decimation-in-time fast Fourier Transform.
--   The given array must have a length that is a power of two,
--   though this property is not checked.
mfft :: forall arr m. (PrimMonad m, Contiguous arr, Element arr (Complex Double))
  => Mutable arr (PrimState m) (Complex Double)
  -> m ()
mfft mut = do {
    len <- Contiguous.sizeMutable mut
  ; let bitReverse !i !j = do {
          ; if i == len - 1
              then stage 0 1
              else do {
                  when (i < j) $ swap mut i j
                ; let inner k l = if k <= l
                        then inner (k `shiftR` 1) (l - k)
                        else bitReverse (i + 1) (l + k)
                ; inner (len `shiftR` 1) j
              }
        }
        stage l l1 = if l == (log2 len)
          then pure ()
          else do {
              let !l2 = l1 `shiftL` 1
                  !e = (negate twoPi) / (intToDouble l2)
                  flight j !a = if j == l1
                    then stage (l + 1) l2
                    else do {
                        let butterfly i = if i >= len
                              then flight (j + 1) (a + e)
                              else do {
                                  let i1 = i + l1
                                ; xi1 :+ yi1 <- Contiguous.read mut i1
                                ; let !c = Prelude.cos a
                                      !s' = Prelude.sin a
                                      d = (c*xi1 - s'*yi1) :+ (s'*xi1 + c*yi1)
                                ; ci <- Contiguous.read mut i
                                ; Contiguous.write mut i1 (ci - d)
                                ; Contiguous.write mut i (ci + d)
                                ; butterfly (i + l2)
                              }
                      ; butterfly j
                    }
            ; flight 0 0
         }
  ; bitReverse 0 0
}

-- wildcard cases should never happen. if they do, really bad things will happen.
b,s :: Int -> Int
b = \case { 0 -> 0x02; 1 -> 0x0c; 2 -> 0xf0; 3 -> 0xff00; 4 -> wordToInt 0xffff0000; 5 -> wordToInt 0xffffffff00000000; _ -> 0; }
s = \case { 0 -> 1; 1 -> 2; 2 -> 4; 3 -> 8; 4 -> 16; 5 -> 32; _ -> 0; }
{-# inline b #-}
{-# inline s #-}

log2 :: Int -> Int
log2 v0 = if v0 <= 0
  then Prelude.error $ "Data.Primitive.Contiguous.FFT: nonpositive input, got " Prelude.++ Prelude.show v0
  else go 5 0 v0
  where
    go !i !r !v
      | i == -1 = r
      | v .&. b i /= 0 =
          let si = s i
          in go (i - 1) (r .|. si) (v `shiftR` si)
      | otherwise = go (i - 1) r v

swap :: forall arr m x. (PrimMonad m, Contiguous arr, Element arr x)
  => Mutable arr (PrimState m) x
  -> Int
  -> Int
  -> m ()
{-# inline swap #-}
swap mut i j = do
  atI <- Contiguous.read mut i
  atJ <- Contiguous.read mut j
  Contiguous.write mut i atJ
  Contiguous.write mut j atI

twoPi :: Double
{-# inline twoPi #-}
twoPi = 6.283185307179586

intToDouble :: Int -> Double
{-# inline intToDouble #-}
intToDouble = Prelude.fromIntegral

wordToInt :: Word -> Int
{-# inline wordToInt #-}
wordToInt = Prelude.fromIntegral

intToComplexDouble :: Int -> Complex Double
{-# inline intToComplexDouble #-}
intToComplexDouble = Prelude.fromIntegral

cmap :: (Contiguous arr, Element arr (Complex Double))
  => (Complex Double -> Complex Double)
  -> arr (Complex Double)
  -> arr (Complex Double)
{-# inline cmap #-}
cmap = Contiguous.map


------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral.SmithNF
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Compute Smith normal form
--
------------------------------------------------

module Numeric.Matrix.Integral.SmithNF where

import Control.Monad (when, unless, forM_, guard, join)
import Control.Monad.ST
import Control.Monad.Loops (whileM_, whileJust_)

import Data.List as L
import Data.STRef

import Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Devel

--import Numeric.Matrix.Integral.HNFLLL (hnfLLLST)
import Numeric.Matrix.Integral.NormalForms (hermiteNFST)

-- | Find off-diagonal entries and calculate min of min{row,col} among them
findDiagSz :: LA.Matrix LA.Z -> Maybe Int
findDiagSz ma =
  let offdiag = filter (uncurry (/=)) $ LA.find (/=0) ma
  in L.foldl' bin Nothing offdiag
  where
    bin Nothing (i,j) = Just (min i j)
    bin (Just k) (i,j) = Just $ L.foldl' min k [i,j]

-- | Compute "pre-Smith" form with a function calculating Hermite normal forms
-- "Pre-Smith" means diag(c_1,c_2,..) with not necessarily c_i|c_{i+1}
preSmithNFST :: STMatrix s LA.Z -> STMatrix s LA.Z ->  STMatrix s LA.Z -> ST s ()
preSmithNFST stMatUL stMatA stMatURt = do
  whileJust_ (liftSTMatrix findDiagSz stMatA) $ \diagSz -> do
    -- Extract a submatrix containing all the off-diagonal entries.
    stMatAex <- thawMatrix =<< extractMatrix stMatA (FromRow diagSz) (FromCol diagSz)
    stMatULex <- thawMatrix =<< extractMatrix stMatUL (FromRow diagSz) AllCols
    stMatURtex <- thawMatrix =<< extractMatrix stMatURt (FromRow diagSz) AllCols
    -- Compute the Hermite normal forms of the extracted submatrix and of the transpose of the result.
    --hnfLLLST stMatULex stMatAex
    join $ hermiteNFST <$> newMatrix 0 0 0 <*> pure stMatULex <*> pure stMatAex
    stMatAext <- thawMatrix =<< (LA.tr' <$> freezeMatrix stMatAex)
    --hnfLLLST stMatURtex stMatAext
    join $ hermiteNFST <$> newMatrix 0 0 0 <*> pure stMatURtex <*> pure stMatAext
    -- Write out the result
    setMatrix stMatA diagSz diagSz =<< (LA.tr' <$> freezeMatrix stMatAext)
    setMatrix stMatUL diagSz 0 =<< freezeMatrix stMatULex
    setMatrix stMatURt diagSz 0 =<< freezeMatrix stMatURtex

isHeadGCD :: LA.Vector LA.Z -> Bool
isHeadGCD = isHeadGCD' . LA.toList
  where
    isHeadGCD' [] = True
    isHeadGCD' (x:xs) = L.foldl' (\b y -> b && rem y x == 0) True xs

-- Compute the submatrix containing all the non-zero diagonal entries
extractMaxNZDiag :: LA.Matrix LA.Z -> (Int,LA.Matrix LA.Z)
extractMaxNZDiag mx = runST $ do
  dRef <- newSTRef 0
  let p d = d < uncurry min (LA.size mx) && (mx!d!d) /= 0
  whileM_ (p <$> readSTRef dRef) $ modifySTRef' dRef (+1)
  d <- readSTRef dRef
  return (d,LA.subMatrix (0,0) (d,d) mx)

-- | Normalize diagonals to form a divisible chain; i.e. c_1 | c_2 | c_3 | ...
normalizeST :: STMatrix s LA.Z -> STMatrix s LA.Z ->  STMatrix s LA.Z -> ST s ()
normalizeST stMatUL stMatA stMatURt = do
  (dlen,matD) <- liftSTMatrix extractMaxNZDiag stMatA
  stMatD <- thawMatrix matD
  stMatULex <- thawMatrix =<< extractMatrix stMatUL (RowRange 0 (dlen-1)) AllCols
  stMatURtex <- thawMatrix =<< extractMatrix stMatURt (RowRange 0 (dlen-1)) AllCols
  normalizeST' stMatULex stMatD stMatURtex
  setMatrix stMatA 0 0 =<< freezeMatrix stMatD
  setMatrix stMatUL 0 0 =<< freezeMatrix stMatULex
  setMatrix stMatURt 0 0 =<< freezeMatrix stMatURtex
  where
    normalizeST' stMatULex stMatD stMatURtex = do
      b <- (isHeadGCD . LA.takeDiag) <$> freezeMatrix stMatD
      n <- liftSTMatrix LA.rows stMatD
      unless b $ do
        forM_ [1..(n-1)] $ \i -> do
          readMatrix stMatD i i >>= writeMatrix stMatD i 0
          rowOper (AXPY 1 i 0 AllCols) stMatURtex
        preSmithNFST stMatULex stMatD stMatURtex
      when (n > 1) $ do
        stMatDex <- thawMatrix =<< extractMatrix stMatD (FromRow 1) (FromCol 1)
        stMatULexex <- thawMatrix =<< extractMatrix stMatULex (FromRow 1) AllCols
        stMatURtexex <- thawMatrix =<< extractMatrix stMatURtex (FromRow 1) AllCols
        normalizeST' stMatULexex stMatDex stMatURtexex
        setMatrix stMatD 1 1 =<< freezeMatrix stMatDex
        setMatrix stMatULex 1 0 =<< freezeMatrix stMatULexex
        setMatrix stMatURtex 1 0 =<< freezeMatrix stMatURtexex

-- | Compute Smith normal form together with transform unimodular matrices.
smithNF :: LA.Matrix LA.Z -> (LA.Matrix LA.Z, LA.Matrix LA.Z, LA.Matrix LA.Z)
smithNF matA = runST $ do
  stMatA <- thawMatrix matA
  stMatUL <- thawMatrix $ LA.ident (LA.rows matA)
  stMatURt <- thawMatrix $ LA.ident (LA.cols matA)
  preSmithNFST stMatUL stMatA stMatURt
  normalizeST stMatUL stMatA stMatURt
  resMatS <- freezeMatrix stMatA
  resMatUL <- freezeMatrix stMatUL
  resMatURt <- freezeMatrix stMatURt
  return (resMatUL, resMatS, LA.tr' resMatURt)

-- | For debug.
-- Compute "pre-Smith form."
preSmithNF :: LA.Matrix LA.Z -> (LA.Matrix LA.Z, LA.Matrix LA.Z, LA.Matrix LA.Z)
preSmithNF matA = runST $ do
  stMatA <- thawMatrix matA
  stMatUL <- thawMatrix $ LA.ident (LA.rows matA)
  stMatURt <- thawMatrix $ LA.ident (LA.cols matA)
  preSmithNFST stMatUL stMatA stMatURt
  resMatS <- freezeMatrix stMatA
  resMatUL <- freezeMatrix stMatUL
  resMatURt <- freezeMatrix stMatURt
  return (resMatUL, resMatS, LA.tr' resMatURt)


{-# LANGUAGE ViewPatterns #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Signal
-- Copyright   : David Johnson (c) 2019-2020
-- License     : BSD 3
-- Maintainer  : David Johnson <djohnson.m@gmail.com>
-- Stability   : Experimental
-- Portability : GHC
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm)
--
-- Functions performing signal processing on 'Array'
--
-- >>> input = vector 3 [10,20,30]
-- >>> positions = vector 5 [0.0, 0.5, 1.0, 1.5, 2.0]
-- >>> approx1 @Double input positions Cubic 0.0
-- ArrayFire Array
-- [5 1 1 1]
--    10.0000
--    13.7500
--    20.0000
--    26.2500
--    30.0000
--
--------------------------------------------------------------------------------
module ArrayFire.Signal where

import Data.Complex

import ArrayFire.FFI
import ArrayFire.Internal.Signal
import ArrayFire.Internal.Types

-- | 'approx1' interpolates data along the first dimensions
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__approx1.htm)
--
-- Interpolation is performed assuming input data is equally spaced with indices in the range [0, n). The positions are sampled with respect to data at these locations.
--
-- >>> input = vector 3 [10,20,30]
-- >>> positions = vector 5 [0.0, 0.5, 1.0, 1.5, 2.0]
-- >>> approx1 @Double input positions Cubic 0.0
-- ArrayFire Array
-- [5 1 1 1]
--    10.0000
--    13.7500
--    20.0000
--    26.2500
--    30.0000
approx1
  :: AFType a
  => Array a
  -- ^ the input array
  -> Array a
  -- ^ array contains the interpolation locations
  -> InterpType
  -- ^ is the interpolation type, it can take one of the values defined by 'InterpType'
  -> Float
  -- ^ is the value that will set in the output array when certain index is out of bounds
  -> Array a
  -- ^ is the array with interpolated values
approx1 arr1 arr2 (fromInterpType -> i1) f =
  op2 arr1 arr2 (\p x y -> af_approx1 p x y i1 f)

-- | approx2 performs interpolation on data along the first and second dimensions.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__approx2.htm)
--
-- Interpolation is performed assuming input data is equally spaced with indices in the range [0, n) along each dimension. The positions are sampled with respect to data at these locations.
--
-- >>> input = matrix @Double (3,3) [ [ 1.0,1.0,1.0 ], [ 2.0, 2.0, 2.0 ], [ 3.0,3.0,3.0 ] ]
-- >>> positions1 = matrix @Double (2,2) [ [ 0.5,1.5 ],[ 0.5,1.5 ] ]
-- >>> positions2 = matrix @Double (2,2) [ [ 0.5,0.5 ],[ 1.5,1.5 ] ]
-- >>> approx2 @Double input positions1 positions2 Cubic 0.0
-- ArrayFire Array
-- [2 2 1 1]
--     1.3750     2.6250
--     1.3750     2.6250
--
approx2
  :: AFType a
  => Array a
  -- ^ is the input array
  -> Array a
  -- ^ array contains the interpolation locations for first dimension
  -> Array a
  -- ^ array contains the interpolation locations for second dimension
  -> InterpType
  -- ^ is the interpolation type, it can take one of the values defined by 'InterpType'
  -> Float
  -- ^ is the value that will set in the output array when certain index is out of bounds
  -> Array a
  -- ^	is the array with interpolated values
approx2 arr1 arr2 arr3 (fromInterpType -> i1) f =
  op3 arr1 arr2 arr3 (\p x y z -> af_approx2 p x y z i1 f)

-- DMJ: Where did these functions go? Were they removed?
-- http://arrayfire.org/docs/group__approx__mat.htm
-- approx1Uniform
--   :: AFType a
--   => Array a
--   -> Array a
--   -> Int
--   -> Double
--   -> Double
--   -> InterpType
--   -> Float
--   -> Array a
-- approx1Uniform arr1 arr2 (fromIntegral -> i1) d1 d2 (fromInterpType -> interp) f =
--   op2 arr1 arr2 (\p x y -> af_approx1_uniform p x y i1 d1 d2 interp f)

-- approx2Uniform
--   :: AFType a
--   => Array a
--   -> Array a
--   -> Int
--   -> Double
--   -> Double
--   -> Array a
--   -> Int
--   -> Double
--   -> Double
--   -> InterpType
--   -> Float
--   -> Array a
-- approx2Uniform arr1 arr2 (fromIntegral -> i1) d1 d2 arr3 (fromIntegral -> i2) d3 d4 (fromInterpType -> interp) f =
--   op3 arr1 arr2 arr3 (\p x y z -> af_approx2_uniform p x y i1 d1 d2 z i2 d3 d4 interp f)

-- | Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm#ga64d0db9e59c9410ba738591ad146a884)
--
-- The Fast Fourier Transform (FFT) is an efficient algorithm to compute the discrete Fourier transform (DFT) of a signal or array. This is most commonly used to convert data in the time (or space) domain to the frequency domain, Then, the inverse FFT (iFFT) is used to return the data to the original domain.
--
-- >>> fft (vector @Double 10 [1..]) 2.0 10
-- ArrayFire Array
-- [10 1 1 1]
--          (110.0000,0.0000)
--          (-10.0000,30.7768)
--          (-10.0000,13.7638)
--          (-10.0000,7.2654)
--          (-10.0000,3.2492)
--          (-10.0000,0.0000)
--          (-10.0000,-3.2492)
--          (-10.0000,-7.2654)
--          (-10.0000,-13.7638)
--          (-10.0000,-30.7768)
fft
  :: (AFType a, Fractional a)
  => Array a
  -- ^ input 'Array'
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals - used to either truncate or pad the input signals.
  -> Array a
  -- ^ is the transformed array
fft a d (fromIntegral -> x) =
  op1 a (\j k -> af_fft j k d x)

-- | Fast Fourier Transform (in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm#gaa2f03c9ee1cb80dc184c0b0a13176da1)
--
-- C Interface for fast fourier transform on one dimensional signals.
--
-- *Note* The input in must be a complex array
--
fftInPlace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
fftInPlace a d = a `inPlace` (flip af_fft_inplace d)

-- | Fast Fourier Transform (2-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft2.htm#gaab3fb1ed398e208a615036b4496da611)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
fft2
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
fft2 a d x y =
  op1 a (\j k -> af_fft2 j k d (fromIntegral x) (fromIntegral y))

-- | Fast Fourier Transform (2-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft2.htm#gacdeebb3f221ae698833dc4900a172b8c)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
-- *Note* The input in must be a complex array
--
fft2_inplace
  :: (Fractional a, AFType a)
  => Array (Complex a)
  -- ^  input array on entry and the output of 2D forward fourier transform on exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
fft2_inplace a d = a `inPlace` (flip af_fft2_inplace d)

-- | Fast Fourier Transform (3-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft3.htm#ga5138ef1740ece0fde2c796904d733c12)
--
-- C Interface for fast fourier transform on three dimensional signals.
--
fft3
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along third dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
fft3 a d x y z =
  op1 a (\j k -> af_fft3 j k d (fromIntegral x) (fromIntegral y) (fromIntegral z))

-- | Fast Fourier Transform (3-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft2.htm#gacdeebb3f221ae698833dc4900a172b8c)
--
-- C Interface for fast fourier transform on three dimensional signals.
--
-- *Note* The input in must be a complex array
--
fft3_inplace
  :: (Fractional a, AFType a)
  => Array (Complex a)
  -- ^  input array on entry and the output of 3D forward fourier transform on exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
fft3_inplace a d = a `inPlace` (flip af_fft3_inplace d)

-- | Inverse Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft.htm#ga2d62c120b474b3b937b0425c994645fe)
--
-- C Interface for inverse fast fourier transform on one dimensional signals.
--
ifft
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^  is the length of output signals - used to either truncate or pad the input signals
  -> Array a
  -- ^ the transformed array
ifft a d x =
  op1 a (\j k -> af_ifft j k d (fromIntegral x))

-- | Inverse Fast Fourier Transform (in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft.htm#ga827379bef0e2cadb382c1b6301c91429)
--
-- C Interface for fast fourier transform on one dimensional signals.
--
ifft_inplace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
ifft_inplace a d = a `inPlace` (flip af_ifft_inplace d)

-- | Inverse Fast Fourier Transform (2-dimensional signals)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft2.htm#ga7cd29c6a35c19240635b62cc5c30dc4f)
--
-- C Interface for inverse fast fourier transform on two dimensional signals.
--
ifft2
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
ifft2 a d x y =
  op1 a (\j k -> af_ifft2 j k d (fromIntegral x) (fromIntegral y))

-- | Inverse Fast Fourier Transform (2-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft2.htm#ga9e6a165d44306db4552a56d421ce5d05)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
ifft2_inplace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
ifft2_inplace a d = a `inPlace` (flip af_ifft2_inplace d)

-- | Inverse Fast Fourier Transform (3-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft3.htm)
--
-- C Interface for inverse fast fourier transform on three dimensional signals.
--
ifft3
  :: AFType a
  => Array a
  -- ^ the input array
  -> Double
  -- ^ the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along third dimension - used to either truncate/pad the input
  -> Array a
  -- ^ the transformed array
ifft3 a d x y z =
  op1 a (\j k -> af_ifft3 j k d (fromIntegral x) (fromIntegral y) (fromIntegral z))

-- | Inverse Fast Fourier Transform (3-dimensional, in-place)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__ifft3.htm#ga439a7a49723bc6cf77cf4fe7f8dfe334)
--
-- C Interface for fast fourier transform on two dimensional signals.
--
ifft3_inplace
  :: (AFType a, Fractional a)
  => Array (Complex a)
  -- ^ is the input array on entry and the output of 1D forward fourier transform at exit
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> IO ()
ifft3_inplace a d = a `inPlace` (flip af_ifft3_inplace d)

-- | Real to Complex Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__r2c.htm#ga7486f342182a18e773f14cc2ab4cb551)
--
-- C Interface for real to complex fast fourier transform for one dimensional signals.
--
-- The first dimension of the output will be of size (pad0 / 2) + 1. The second dimension of the output will be pad1. The third dimension of the output will be pad 2.
--
fftr2c
  :: (Fractional a, AFType a)
  => Array a
  -- ^ is a real array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Array a
  -- ^ is a complex array containing the non redundant parts of in.
fftr2c a d x =
  op1 a (\j k -> af_fft_r2c j k d (fromIntegral x))

-- | Real to Complex Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__r2c.htm#ga7486f342182a18e773f14cc2ab4cb551)
--
-- C Interface for real to complex fast fourier transform for two dimensional signals.
--
-- The first dimension of the output will be of size (pad0 / 2) + 1. The second dimension of the output will be pad1. The third dimension of the output will be pad 2.
--
fft2r2c
  :: (Fractional a, AFType a)
  => Array a
  -- ^ is a real array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Array a
  -- ^ is a complex array containing the non redundant parts of in.
fft2r2c a d x y =
  op1 a (\j k -> af_fft2_r2c j k d (fromIntegral x) (fromIntegral y))

-- | Real to Complex Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__r2c.htm#gab4ca074b54218b74d8cfbda63d38be51)
--
-- C Interface for real to complex fast fourier transform for three dimensional signals.
--
-- The first dimension of the output will be of size (pad0 / 2) + 1. The second dimension of the output will be pad1. The third dimension of the output will be pad 2.
--
fft3r2c
  :: (Fractional a, AFType a)
  => Array a
  -- ^ is a real array
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Int
  -- ^ is the length of output signals along first dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along second dimension - used to either truncate/pad the input
  -> Int
  -- ^ is the length of output signals along third dimension - used to either truncate/pad the input
  -> Array a
  -- ^ is a complex array containing the non redundant parts of in.
fft3r2c a d x y z =
  op1 a (\j k -> af_fft3_r2c j k d (fromIntegral x) (fromIntegral y) (fromIntegral z))

-- | Complex to Real Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__c2r.htm#gaa5efdfd84213a4a07d81a5d534cde5ac)
--
-- C Interface for complex to real fast fourier transform for one dimensional signals.
--
-- The first dimension of the output will be 2 * dim0 - 1 if is_odd is true else 2 * dim0 - 2 where dim0 is the first dimension of the input. The remaining dimensions are unchanged.
--
fftc2r
  :: AFType a
  => Array a
  -- ^ is a complex array containing only the non redundant parts of the signals.
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Bool
  -- ^ is a flag signifying if the output should be even or odd size
  -> Array a
  -- ^ is a real array containing the output of the transform.
fftc2r a cm (fromIntegral . fromEnum -> cd) = op1 a (\x y -> af_fft_c2r x y cm cd)

-- | Complex to Real Fast Fourier Transform (2-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__c2r.htm#gaaa7da16f226cacaffced631e08da4493)
--
-- C Interface for complex to real fast fourier transform for two dimensional signals.
--
-- The first dimension of the output will be 2 * dim0 - 1 if is_odd is true else 2 * dim0 - 2 where dim0 is the first dimension of the input. The remaining dimensions are unchanged.
--
fft2C2r
  :: AFType a
  => Array a
  -- ^ is a complex array containing only the non redundant parts of the signals.
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Bool
  -- ^ is a flag signifying if the output should be even or odd size
  -> Array a
  -- ^ is a real array containing the output of the transform.
fft2C2r a cm (fromIntegral . fromEnum -> cd) = op1 a (\x y -> af_fft2_c2r x y cm cd)

-- | Complex to Real Fast Fourier Transform (3-dimensional)
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__c2r.htm#gaa9b3322d9ffab15268919e1f114bed24)
--
-- C Interface for complex to real fast fourier transform for three dimensional signals.
--
-- The first dimension of the output will be 2 * dim0 - 1 if is_odd is true else 2 * dim0 - 2 where dim0 is the first dimension of the input. The remaining dimensions are unchanged.
--
fft3C2r
  :: AFType a
  => Array a
  -- ^ is a complex array containing only the non redundant parts of the signals.
  -> Double
  -- ^ is the normalization factor with which the input is scaled after the transformation is applied
  -> Bool
  -- ^ is a flag signifying if the output should be even or odd size
  -> Array a
  -- ^ is a real array containing the output of the transform.
fft3C2r a cm (fromIntegral . fromEnum -> cd) = op1 a (\x y -> af_fft3_c2r x y cm cd)

-- | Convolution Integral for one dimensional data
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve1.htm#ga25d77b794463b5cd72cd0b7f4af140d7)
--
-- C Interface for convolution on one dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve1
  :: AFType a
  => Array a
  -- ^ the input signal
  -> Array a
  -- ^ the signal that shall be flipped for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> ConvDomain
  -- ^ specifies if the convolution should be performed in frequency os spatial domain
  -> Array a
  -- ^ convolved array
convolve1 a b (toConvMode -> cm) (fromConvDomain -> cd) = op2 a b (\x y z -> af_convolve1 x y z cm cd)

-- | Convolution Integral for two dimensional data
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve2.htm#ga25d77b794463b5cd72cd0b7f4af140d7)
--
-- C Interface for convolution on two dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve2
  :: AFType a
  => Array a
  -- ^ the input signal
  -> Array a
  -- ^ the signal that shall be flipped for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> ConvDomain
  -- ^ specifies if the convolution should be performed in frequency os spatial domain
  -> Array a
  -- ^ convolved array
convolve2 a b (toConvMode -> cm) (fromConvDomain -> cd) = op2 a b (\x y z -> af_convolve2 x y z cm cd)

-- | Convolution Integral for three dimensional data
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve3.htm#ga25d77b794463b5cd72cd0b7f4af140d7)
--
-- C Interface for convolution on three dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve3
  :: AFType a
  => Array a
  -- ^ the input signal
  -> Array a
  -- ^ the signal that shall be flipped for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> ConvDomain
  -- ^ specifies if the convolution should be performed in frequency os spatial domain
  -> Array a
  -- ^ convolved array
convolve3 a b (toConvMode -> cm) (fromConvDomain -> cd) =
  op2 a b (\x y z -> af_convolve3 x y z cm cd)

-- | C Interface for separable convolution on two dimensional signals.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__convolve.htm#gaeb6ba88155cf3ef29d93f97b147e372f)
--
-- C Interface for separable convolution on two dimensional signals.
--
-- *Note* The default parameter of domain, AF_CONV_AUTO, heuristically switches between frequency and spatial domain.
--
convolve2Sep
  :: AFType a
  => Array a
  -- ^ filter that has to be applied along the coloumns
  -> Array a
  -- ^ filter that has to be applied along the rows
  -> Array a
  -- ^ the input array
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> Array a
  -- ^ convolved array
convolve2Sep a b c (toConvMode -> d) = op3 a b c (\x y z j -> af_convolve2_sep x y z j d)

-- DMJ: did this get removed? Can't find in latest docs
-- fftConvolve1
--   :: AFType a
--   => Array a
--   -> Array a
--   -> ConvMode
--   -> Array a
-- fftConvolve1 a b (toConvMode -> c) = op2 a b (\x y z -> af_fft_convolve1 x y z c)

-- | 2D Convolution using Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__convolve2.htm#gab52ebe631d8358cdef1b5c8a95550556)
--
-- A convolution is a common operation between a source array, a, and a filter (or kernel) array b. The answer to the convolution is the same as computing the coefficients in polynomial multiplication, if a and b are the coefficients.
--
-- C Interface for FFT-based convolution on two dimensional signals.
--
fftConvolve2
  :: AFType a
  => Array a
  -- ^ is the input signal
  -> Array a
  --  ^ is the signal that shall be used for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> Array a
  -- ^  is convolved array
fftConvolve2 a b (toConvMode -> c) = op2 a b (\x y z -> af_fft_convolve2 x y z c)

-- | 3D Convolution using Fast Fourier Transform
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft__convolve3.htm)
--
-- A convolution is a common operation between a source array, a, and a filter (or kernel) array b. The answer to the convolution is the same as computing the coefficients in polynomial multiplication, if a and b are the coefficients.
--
-- C Interface for FFT-based convolution on three dimensional signals.
--
fftConvolve3
  :: AFType a
  => Array a
  -- ^ is the input signal
  -> Array a
  --  ^ is the signal that shall be used for the convolution operation
  -> ConvMode
  -- ^ indicates if the convolution should be expanded or not(where output size equals input)
  -> Array a
  -- ^  is convolved array
fftConvolve3 a b (toConvMode -> c) = op2 a b (\x y z -> af_fft_convolve3 x y z c)

-- | Finite Impulse Filter.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fir.htm#ga2a850e69775eede4709e0d607bba240b)
--
-- Finite impulse filters take an input x and a co-efficient array b to generate an output y such that:
--
-- C Interface for finite impulse response filter.
--
fir
  :: AFType a
  => Array a
  -- ^ is the input signal to the filter
  -> Array a
  -- ^ is the array containing the coefficients of the filter
  -> Array a
  -- ^ is the output signal from the filter
fir a b = op2 a b af_fir

-- | Infinite Impulse Filter.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__iir.htm#ga7adcc364da0a66cdfd2bb351215456c4)
--
-- Infinite impulse filters take an input x and a feedforward array b, feedback array a to generate an output y such that:
--
-- C Interface for infinite impulse response filter.
--
-- *Note* The feedforward coefficients are currently limited to a length of 512
--
iir
  :: AFType a
  => Array a
  -- ^ the array containing the feedforward coefficient
  -> Array a
  -- ^ is the array containing the feedback coefficients
  -> Array a
  -- ^ is the input signal to the filter
  -> Array a
  -- ^ the output signal from the filter
iir a b c = op3 a b c af_iir

-- | Median Filter
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__image__func__medfilt.htm#gaaf3f62f2de0f4dc315b831e494e1b2c0)
--
-- A median filter is similar to the arbitrary filter except that instead of a weighted sum, the median value of the pixels covered by the kernel is returned.
--
-- C Interface for median filter.
--
medFilt
  :: AFType a
  => Array a
  -- ^ 'Array' is the input image
  -> Int
  -> Int
  -> BorderType
  -> Array a
  -- ^ 'Array' is the processed image
medFilt a l w (fromBorderType -> b) =
 a `op1` (\x y -> af_medfilt x y (fromIntegral l) (fromIntegral w) b)

-- | 1D Median Filter
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__image__func__medfilt.htm#gad108ea62cbbb5371bd14a17d06384359)
--
-- A median filter is similar to the arbitrary filter except that instead of a weighted sum, the median value of the pixels covered by the kernel is returned.
--
-- C Interface for 1D median filter.
--
medFilt1
  :: AFType a
  => Array a
  -- ^ 'Array' is the input signal
  -> Int
  -- ^ Is the kernel width
  -> BorderType
  -- ^ value will decide what happens to border when running filter in their neighborhood. It takes one of the values [AF_PAD_ZERO | AF_PAD_SYM]
  -> Array a
  -- ^ 'Array' is the processed signal
medFilt1 a w (fromBorderType -> b) =
 a `op1` (\x y -> af_medfilt1 x y (fromIntegral w) b)

-- | 2D Median Filter
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__image__func__medfilt.htm#ga2cb99dca5842f74f6b9cd28eb187a9cd)
--
-- A median filter is similar to the arbitrary filter except that instead of a weighted sum, the median value of the pixels covered by the kernel is returned.
--
-- C Interface for 2D median filter.
--
medFilt2
  :: AFType a
  => Array a
  -- ^ 'Array' is the input image
  -> Int
  -- ^ the kernel height
  -> Int
  -- ^ the kernel width
  -> BorderType
  -- ^ value will decide what happens to border when running filter in their neighborhood. It takes one of the values [AF_PAD_ZERO | AF_PAD_SYM]
  -> Array a
  -- ^ 'Array' is the processed image
medFilt2 a l w (fromBorderType -> b) =
 a `op1` (\x y -> af_medfilt2 x y (fromIntegral l) (fromIntegral w) b)

-- | C Interface for setting plan cache size.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__signal__func__fft.htm#ga4ddef19b43d9a50c97b1a835df60279a)
--
-- This function doesn't do anything if called when CPU backend is active. The plans associated with the most recently used array sizes are cached.
--
-- >>> setFFTPlanCacheSize 2
-- ()
--
setFFTPlanCacheSize
  :: Int
  -- ^ is the number of plans that shall be cached
  -> IO ()
setFFTPlanCacheSize =
 afCall . af_set_fft_plan_cache_size . fromIntegral


{-# LANGUAGE BangPatterns #-}

{- | Some extra facilities to work with 'Rand' monad and 'PureMT'
     random number generator.
-}

module Moo.GeneticAlgorithm.Random
    (
    -- * Random numbers from given range
      getRandomR
    , getRandom
    -- * Probability distributions
    , getNormal2
    , getNormal
    -- * Random samples and shuffles
    , randomSample
    , randomSampleIndices
    , shuffle
    -- * Building blocks
    , withProbability
    -- * Re-exports from random number generator packages
    , getBool, getInt, getWord, getInt64, getWord64, getDouble
    , runRandom, evalRandom, newPureMT
    , Rand, Random, PureMT
    ) where

import Control.Monad (liftM)
import Control.Monad.State.Strict (MonadState(..), State, runState, evalState)
import Data.Complex (Complex (..))
import System.Random (RandomGen, Random(..))
import System.Random.Mersenne.Pure64
import qualified System.Random.Shuffle as S
import qualified Data.Set as Set
import Data.Word (Word64)
import Data.Int (Int64)

-- | Monad for handling randomized computation. This uses the mersenne twister RNG.
type Rand = State PureMT

-- | Run a random computation using the generator @g@, returning the result
-- and the updated generator.
runRandom  :: Rand a -> PureMT -> (a, PureMT)
runRandom = runState

-- | Evaluate a random computation using the mersenne generator @g@.  Note that the
-- generator @g@ is not returned, so there's no way to recover the
-- updated version of @g@.
evalRandom :: Rand a -> PureMT -> a
evalRandom = evalState

-- | Yield a new boolean value from the generator.
getBool     :: Rand Bool
getBool     = fmap (<0) getInt

-- | Yield a new 'Int' value from the generator.
getInt      :: Rand Int
getInt      = state randomInt

-- | Yield a new 'Word' value from the generator.
getWord     :: Rand Word
getWord     = state randomWord

-- | Yield a new 'Int64' value from the generator.
getInt64    :: Rand Int64
getInt64    = state randomInt64

-- | Yield a new 53-bit precise 'Double' value from the generator.
getDouble :: Rand Double
getDouble = state randomDouble

-- | Yield a new 'Word64' value from the generator.
getWord64   :: Rand Word64
getWord64   = state randomWord64

-- | Yield a new randomly selected value of type @a@ in the range @(lo, hi)@.
-- See 'System.Random.randomR' for details.
getRandomR :: Random a => (a, a) -> Rand a
getRandomR range = state $ randomR range

-- | Yield a new randomly selected value of type @a@.
-- See 'System.Random.random' for details.
getRandom :: Random a => Rand a
getRandom = state random

-- | Yield two randomly selected values which follow standard
-- normal distribution.
getNormal2 :: Rand (Double, Double)
getNormal2 = do
  -- Box-Muller method
  u <- getDouble
  v <- getDouble
  let (c :+ s) = exp (0 :+ (2*pi*v))
  let r = sqrt $ (-2) * log u
  return (r*c, r*s)

-- | Yield one randomly selected value from standard normal distribution.
getNormal :: Rand Double
getNormal = fst `liftM` getNormal2

-- | Take at most n random elements from the list. Preserve order.
randomSample :: Int -> [a] -> Rand [a]
randomSample n xs =
  state $ \g -> select g n (length xs) xs []
  where
    select rng _ _ [] acc = (reverse acc, rng)
    select rng n m xs acc
        | n <= 0     = (reverse acc, rng)
        | otherwise  =
            let (k, rng') = randomR (0, m - n) rng
                (x:rest) = drop k xs
            in  select rng' (n-1) (m-k-1) rest (x:acc)

-- | Select @sampleSize@ numbers in the range from @0@ to @(populationSize-1)@.
-- The function works best when @sampleSize@ is much smaller than @populationSize@.
randomSampleIndices :: Int -> Int -> Rand [Int]
randomSampleIndices sampleSize populationSize =
    state $ \g ->
        let (sampleSet, g') = buildSampleSet g sampleSize Set.empty
        in  (Set.toList sampleSet, g')
  where
    buildSampleSet g n s
        | n <= 0 = (s, g)
        | otherwise =
            let (i, g') = randomR (0, populationSize-1) g
            in  if (i `Set.member` s)
                then buildSampleSet g' n s
                else buildSampleSet g' (n-1) (Set.insert i s)

-- | Randomly reorder the list.
shuffle :: [a] -> Rand [a]
shuffle xs = state $ randomShuffle xs (length xs)

-- | Given a sequence (e1,...en) to shuffle, its length, and a random
-- generator, compute the corresponding permutation of the input
-- sequence, return the permutation and the new state of the
-- random generator.
randomShuffle :: RandomGen gen => [a] -> Int -> gen -> ([a], gen)
randomShuffle elements len g =
    let (rs, g') = rseq len g
    in  (S.shuffle elements rs, g')
  where
  -- | The sequence (r1,...r[n-1]) of numbers such that r[i] is an
  -- independent sample from a uniform random distribution
  -- [0..n-i]
  rseq :: RandomGen gen => Int -> gen -> ([Int], gen)
  rseq n g = second lastGen . unzip $ rseq' (n - 1) g
      where
        rseq' :: RandomGen gen => Int -> gen -> [(Int, gen)]
        rseq' i gen
          | i <= 0    = []
          | otherwise = let (j, gen') = randomR (0, i) gen
                        in  (j, gen') : rseq' (i - 1) gen'
        -- apply a function on the second element of a pair
        second :: (b -> c) -> (a, b) -> (a, c)
        second f (x,y) = (x, f y)
        -- the last returned random number generator
        lastGen [] = g   -- didn't use the generator yet
        lastGen (lst:[]) = lst
        lastGen gens = lastGen (drop 1 gens)

-- |Modify value with probability @p@. Return the unchanged value with probability @1-p@.
withProbability :: Double -> (a -> Rand a) -> (a -> Rand a)
withProbability p modify x = do
  t <- getDouble
  if t < p
     then modify x
     else return x


{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveAnyClass                           #-}
{-# LANGUAGE DeriveDataTypeable                       #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE DerivingVia                              #-}
{-# LANGUAGE FlexibleInstances                        #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE MultiParamTypeClasses                    #-}
{-# LANGUAGE PatternSynonyms                          #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TemplateHaskell                          #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Backprop.Learn.Model.Neural.LSTM (
    -- * LSTM
    lstm
  , LSTMp(..), lstmForget, lstmInput, lstmUpdate, lstmOutput
  , reshapeLSTMpInput
  , reshapeLSTMpOutput
  , lstm'
    -- * GRU
  , gru
  , GRUp(..), gruMemory, gruUpdate, gruOutput
  , gru'
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Function
import           Backprop.Learn.Model.Neural
import           Backprop.Learn.Model.Regression
import           Backprop.Learn.Model.State
import           Backprop.Learn.Model.Types
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Primitive
import           Data.Type.Tuple
import           Data.Typeable
import           GHC.Generics                          (Generic)
import           GHC.TypeNats
import           Lens.Micro
import           Lens.Micro.TH
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.OneLiner
import           Numeric.Opto.Ref
import           Numeric.Opto.Update
import           Statistics.Distribution
import qualified Data.Binary                           as Bi
import qualified Numeric.LinearAlgebra.Static          as H
import qualified System.Random.MWC                     as MWC

-- TODO: allow parameterize internal activation function?
-- TODO: Peepholes

-- | 'LSTM' layer parmateters
data LSTMp (i :: Nat) (o :: Nat) =
    LSTMp { _lstmForget :: !(FCp (i + o) o)
          , _lstmInput  :: !(FCp (i + o) o)
          , _lstmUpdate :: !(FCp (i + o) o)
          , _lstmOutput :: !(FCp (i + o) o)
          }
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Bi.Binary, Regularize, Backprop)

deriving via (GNum (LSTMp i o)) instance (KnownNat i, KnownNat o) => Num (LSTMp i o)
deriving via (GNum (LSTMp i o)) instance (KnownNat i, KnownNat o) => Fractional (LSTMp i o)
deriving via (GNum (LSTMp i o)) instance (KnownNat i, KnownNat o) => Floating (LSTMp i o)

makeLenses ''LSTMp

instance (PrimMonad m, KnownNat i, KnownNat o) => Mutable m (LSTMp i o) where
    type Ref m (LSTMp i o) = GRef m (LSTMp i o)
    thawRef   = gThawRef
    freezeRef = gFreezeRef
    copyRef   = gCopyRef
instance (PrimMonad m, KnownNat i, KnownNat o) => LinearInPlace m Double (LSTMp i o)

instance (PrimMonad m, KnownNat i, KnownNat o) => Learnable m (LSTMp i o)

-- | Stateless version of 'lstm' that takes the "previous input" as a part
-- of the input vector.
lstm'
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LSTMp i o)) ('Just (R o)) (R (i + o)) (R o)
lstm' = modelD $ \(PJust p) x (PJust s) ->
    let forget = logistic $ runLRp (p ^^. lstmForget) x
        input  = logistic $ runLRp (p ^^. lstmInput ) x
        update = tanh     $ runLRp (p ^^. lstmUpdate) x
        s'     = forget * s + input * update
        o      = logistic $ runLRp (p ^^. lstmOutput) x
        h      = o * tanh s'
    in  (h, PJust s')

-- | Long-term short-term memory layer
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
lstm
    :: (KnownNat i, KnownNat o)
    => Model ('Just (LSTMp i o)) ('Just (R o :# R o)) (R i) (R o)
lstm = recurrent H.split (H.#) id lstm'

reshapeLSTMpInput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat i', KnownNat o)
    => d
    -> MWC.Gen (PrimState m)
    -> LSTMp i o
    -> m (LSTMp i' o)
reshapeLSTMpInput d g (LSTMp forget input update output) =
    LSTMp <$> reshaper forget
          <*> reshaper input
          <*> reshaper update
          <*> reshaper output
  where
    reshaper = reshapeLRpInput d g

reshapeLSTMpOutput
    :: (ContGen d, PrimMonad m, KnownNat i, KnownNat o, KnownNat o')
    => d
    -> MWC.Gen (PrimState m)
    -> LSTMp i o
    -> m (LSTMp i o')
reshapeLSTMpOutput d g (LSTMp forget input update output) =
    LSTMp <$> reshaper forget
          <*> reshaper input
          <*> reshaper update
          <*> reshaper output
  where
    reshaper = reshapeLRpInput  d g
           <=< reshapeLRpOutput d g

-- | Forget biases initialized to 1
instance (KnownNat i, KnownNat o) => Initialize (LSTMp i o) where
    initialize d g = LSTMp <$> set (mapped . fcBias) 1 (initialize d g)
                           <*> initialize d g
                           <*> initialize d g
                           <*> initialize d g

-- | 'GRU' layer parmateters
data GRUp (i :: Nat) (o :: Nat) =
    GRUp { _gruMemory :: !(FCp (i + o) o)
         , _gruUpdate :: !(FCp (i + o) o)
         , _gruOutput :: !(FCp (i + o) o)
         }
  deriving stock     (Generic, Typeable, Show)
  deriving anyclass  (NFData, Linear Double, Metric Double, Bi.Binary, Initialize, Regularize, Backprop)

deriving via (GNum (GRUp i o)) instance (KnownNat i, KnownNat o) => Num (GRUp i o)
deriving via (GNum (GRUp i o)) instance (KnownNat i, KnownNat o) => Fractional (GRUp i o)
deriving via (GNum (GRUp i o)) instance (KnownNat i, KnownNat o) => Floating (GRUp i o)

makeLenses ''GRUp

instance (PrimMonad m, KnownNat i, KnownNat o) => Mutable m (GRUp i o) where
    type Ref m (GRUp i o) = GRef m (GRUp i o)
    thawRef   = gThawRef
    freezeRef = gFreezeRef
    copyRef   = gCopyRef
instance (KnownNat i, KnownNat o, Mutable m (GRUp i o)) => LinearInPlace m Double (GRUp i o)

instance (KnownNat i, KnownNat o, PrimMonad m) => Learnable m (GRUp i o)

-- | Stateless version of 'gru' that takes the "previous input" as a part
-- of the input vector.
gru'
    :: forall i o. (KnownNat i, KnownNat o)
    => Model ('Just (GRUp i o)) 'Nothing (R (i + o)) (R o)
gru' = modelStatelessD $ \(PJust p) x ->
    let z      = logistic $ runLRp (p ^^. gruMemory) x
        r      = logistic $ runLRp (p ^^. gruUpdate) x
        r'     = 1 # r
        h'     = tanh     $ runLRp (p ^^. gruOutput) (r' * x)
    in  (1 - z) * snd (split @i x) + z * h'

-- | Gated Recurrent Unit
--
-- <http://colah.github.io/posts/2015-08-Understanding-LSTMs/>
--
gru :: (KnownNat i, KnownNat o)
    => Model ('Just (GRUp i o)) ('Just (R o)) (R i) (R o)
gru = recurrent H.split (H.#) id gru'


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

import Data.Maybe
import Data.Number.Symbolic
import qualified Data.Number.Symbolic as Sym
import Data.Proxy

import qualified Naperian as N
import qualified Data.Foldable as F
import           Control.Applicative ( liftA2 )
import qualified GHC.TypeLits as M
import           Data.Functor
import           Data.List.Split

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

kx, ky :: Floating a => a
kx = 0.5
ky = 0.75

-- spatial mesh size
nx, ny :: Int
nx = 30
ny = 60

-- x mesh spacing
-- y mesh spacing
dx :: Floating a => a
dx = 1 / (fromIntegral nx - 1)

dy :: Floating a => a
dy = 1 / (fromIntegral ny - 1)

c1, c2 :: Floating a => a
c1 = kx/dx/dx
c2 = ky/dy/dy

cc4' :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
cc4' = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f m n i j k l | i == 0               = 0
                        | j == 0               = 0
                        | i == n - 1           = 0
                        | j == m - 1           = 0
                        | k == i - 1 && l == j = 1
                        | k == i     && l == j = -2
                        | k == i + 1 && l == j = 1
                        | otherwise            = 0

cc5' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
cc5' = N.binary (*) (N.Scalar c2) cc4'

cc5Sym' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
          N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
cc5Sym' = N.binary (*) (N.Scalar $ var "c2") cc4'

yy4' :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
yy4' = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f :: Int -> Int -> Int -> Int -> Int -> Int -> b
          f m n i j k l | i == 0                   = 0
                        | j == 0                   = 0
                        | i == n - 1               = 0
                        | j == m - 1               = 0
                        | k == i     && l == j - 1 = 1
                        | k == i     && l == j     = -2
                        | k == i     && l == j + 1 = 1
                        | otherwise                = 0

yy5' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
yy5' = N.binary (*) (N.Scalar c1) yy4'

yy5Sym' :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
           N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
yy5Sym' = N.binary (*) (N.Scalar $ var "c1") yy4'

ccSym5 = cc5Sym' @Double @4 @5
yy5Sym = yy5Sym' @Double @4 @5

ccSym5

yy5Sym

fmap (N.elements . N.Prism . N.Prism . N.Scalar) $ N.elements $ N.crystal $ N.crystal $ N.binary (+) cc5Sym yy5Sym

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedLists       #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}

import Data.Maybe
import Data.Number.Symbolic
import Data.Proxy

import qualified Naperian as N
import qualified Data.Foldable as F
import           Control.Applicative ( liftA2 )
import qualified GHC.TypeLits as M

import           Numeric.Sundials.ARKode.ODE
import           Numeric.LinearAlgebra

x1, a, x2 :: Double
x1 = 0
a = 1.0
x2 = a

y1, y2 :: Double
y1 = 0.0
y2 = 1.0

bigT :: Double
bigT = 1000.0

n :: Int
n = 2

dx :: Double
dx = a / (fromIntegral n + 1)

dy :: Double
dy = a / (fromIntegral n + 1)

beta, s :: Double
beta = 1.0e-5
s = beta / (dx * dx)

kx, ky :: Double
kx = 0.5
ky = 0.75

c1, c2 :: Double
c1 = kx/dx/dx
c2 = ky/dy/dy

bigAA1 :: Matrix Double
bigAA1 = assoc (n * n, n * n) 0.0 [((i, j), f (i, j)) | i <- [0 .. n * n - 1]
                                                      , j <- [i - n, i,  i + n]
                                                      , j `elem` [0 .. n * n -1]]
  where
    f (i, j) | i     == j = (-2.0) * c1
             | i - n == j = 1.0    * c1
             | i + n == j = 1.0    * c1
             | otherwise = error $ show (i, j)

bigAA2 :: Matrix Double
bigAA2 = diagBlock (replicate n bigA)
  where
    bigA :: Matrix Double
    bigA = assoc (n, n) 0.0 [((i, j), f (i, j)) | i <- [0 .. n - 1]
                                                , j <- [i-1..i+1]
                                                , j `elem` [0..n-1]]
      where
        f (i, j) | i     == j = (-2.0) * c2
                 | i - 1 == j = 1.0    * c2
                 | i + 1 == j = 1.0    * c2

bigAA :: Matrix Double
bigAA = bigAA1 + bigAA2

bigAA1

bigAA2

n

bigZZ1 :: Matrix Double
bigZZ1 = assoc (m * m, m * m) 0.0 [((i, j), f (i, j)) | i <- [0 .. m * m - 1]
                                                      , j <- [0 .. m * m - 1]]
  where
    m = n + 2
    f (i, j) | i     == 0     = 0.0
             | j     == 0     = 0.0
             | i     == j     = (-2.0) * c1
             | i - n == j     = 1.0    * c1
             | i + n == j     = 1.0    * c1
             | i     == n + 1 = 0.0
             | j     == n + 1 = 0.0
             | otherwise      = 0.0

bigZZ1

x :: forall m n . (M.KnownNat m, M.KnownNat n) => N.Vector n (N.Vector m (Sym Int))
x = (fromJust . N.fromList) $
    map (fromJust . N.fromList) ([[var $ (\(x,y) -> "A" ++ show x ++ "," ++ show y) (x,y) | y <- [1..m]] | x <- [1..n]] :: [[Sym Int]])
    where
      m = M.natVal (undefined :: Proxy m)
      n = M.natVal (undefined :: Proxy n)

u1 :: N.Hyper '[N.Vector 3, N.Vector 2] (Sym Int)
u1 = N.Prism $ N.Prism (N.Scalar x)

u1

y :: forall n . M.KnownNat n => N.Vector n (Sym Int)
y = (fromJust . N.fromList) $
    (map (var . ("v" ++) . show) [1..n ] :: [Sym Int])
    where
    n = M.natVal (undefined :: Proxy n)

u2 :: N.Hyper '[N.Vector 3] (Sym Int)
u2 = N.Prism (N.Scalar y)

N.innerH u1 u2