{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE Strict                #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Tanh
Description : Hyperbolic tangent nonlinear layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Tanh
  ( Tanh(..)
  , SpecTanh (..)
  , specTanh1D
  , specTanh2D
  , specTanh3D
  , specTanh
  , tanhLayer
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Constraint                (Dict (..))
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build
import           Grenade.Utils.Conversion
import           Grenade.Utils.Vector

-- | A Tanh layer.
--   A layer which can act between any shape of the same dimension, performing a tanh function.
data Tanh = Tanh
  deriving (Generic,NFData,Show)

instance UpdateLayer Tanh where
  type Gradient Tanh = ()
  runUpdate _ _ _ = Tanh

instance RandomLayer Tanh where
  createRandomWith _ _ = return Tanh

instance Serialize Tanh where
  put _ = return ()
  get = return Tanh

instance (a ~ b, SingI a) => Layer Tanh a b where
  type Tape Tanh a b = S a
  -- runForwards _ (S1DV v) = (S1DV v, S1DV $ mapVectorInPlace tanh v) -- This is inplace replacement!
  -- runForwards _ (S2DV v) = (S2DV v, S2DV $ mapVectorInPlace tanh v) -- This is inplace replacement!
  runForwards _ a        = (a, tanh a)
  runBackwards _ (S1DV v) (S1DV gs) = ((), S1DV $ zipWithVector (\t g -> max 0.005 (1 - (tanh t) ^ (2 :: Int)) * g) v gs)
  runBackwards _ (S2DV v) (S2DV gs) = ((), S2DV $ zipWithVector (\t g -> max 0.005 (1 - (tanh t) ^ (2 :: Int)) * g) v gs)
  runBackwards _ (S1D a) (S1D g)= ((), S1D $ LAS.dvmap (max 0.005) (tanh' a) * g)
  runBackwards _ (S2D a) (S2D g)= ((), S2D $ LAS.dmmap (max 0.005) (tanh' a) * g)
  runBackwards _ (S3D a) (S3D g)= ((), S3D $ LAS.dmmap (max 0.005) (tanh' a) * g)
  runBackwards l x y = runBackwards l x (toLayerShape x y)
  -- runBackwards _ a g = ((), tanh' a * g)

tanh' :: (Floating a) => a -> a
tanh' t = 1 - s ^ (2 :: Int)  where s = tanh t

-------------------- DynamicNetwork instance --------------------

instance FromDynamicLayer Tanh where
  fromDynamicLayer inp _ _ = SpecNetLayer $ SpecTanh (tripleFromSomeShape inp)

instance ToDynamicLayer SpecTanh where
  toDynamicLayer _ _ (SpecTanh (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     case (rows, cols, depth) of
         (_, 1, 1)    -> return $ SpecLayer Tanh (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> return $ SpecLayer Tanh (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _    -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer Tanh (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a Tanh layer.
specTanh1D :: Integer -> SpecNet
specTanh1D i = specTanh3D (i, 1, 1)

-- | Create a specification for a Tanh layer.
specTanh2D :: (Integer, Integer) -> SpecNet
specTanh2D (i, j) = specTanh3D (i, j, 1)

-- | Create a specification for a Tanh layer.
specTanh3D :: (Integer, Integer, Integer) -> SpecNet
specTanh3D = SpecNetLayer . SpecTanh

-- | Create a specification for a Tanh layer.
specTanh :: (Integer, Integer, Integer) -> SpecNet
specTanh = SpecNetLayer . SpecTanh

-- | Add a Tanh layer to your build.
tanhLayer :: BuildM ()
tanhLayer = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecTanh

-------------------- GNum instances --------------------

instance GNum Tanh where
  _ |* Tanh = Tanh
  _ |+ Tanh = Tanh
  sumG _ = Tanh


{-# LANGUAGE FlexibleContexts #-} {- Container (Vector a) => ..., etc. -}
{-# LANGUAGE RankNTypes       #-} {- For type traces w/ ST             -}
{-# LANGUAGE Strict           #-} {- Lazy evaluation is no good here   -}
{-# LANGUAGE TemplateHaskell  #-} {- Since we're using lenses aplenty  -}

{-
   TODO LIST:
   * Make sure basic functionality works, write some tests.
   * Use Numeric.LinearAlgebra.Devel to guarantee that ALL matrix operations
     are zero-copy or low allocation overhead.
   * Leverage chunk-based parallelism (lots of embarassing parallelism!)
-}

module Optimization.NelderMead
  ( nelderMeadAt
  , nelderMeadFull
  , nelderMead
  ) where

import Control.Applicative
import Control.Lens
import Control.Monad
import Control.Monad.ST

import Data.Maybe
import Data.Ord   ( comparing )
import Data.STRef ( STRef
                  , modifySTRef
                  , newSTRef
                  , readSTRef
                  , writeSTRef )

import Data.Vector.Storable          as VS  hiding ( take
                                                   , forM_
                                                   , (!) )
import Data.Vector.Algorithms.Intro  ( sortBy )
import Data.Vector.Mutable           as VM  hiding ( take )
import Data.Vector.Storable.Mutable  as VSM hiding ( take )

import Numeric.LinearAlgebra.Data    as LA
import Numeric.LinearAlgebra.HMatrix as HM
import Numeric.LinearAlgebra.Devel   as MutM

import Optimization.Types

-- One point of the simplex and its associated score.
data Vertex s = Vertex {
  which    :: Int,
  score    :: s
  }

-- Ranked vertices for NM search.
data VertexRank a b = VertexRank {
  bestOf         :: Vertex b,
  nextBestOf     :: Vertex b,
  worstOf        :: Vertex b,
  goodCentroidOf :: Vector a
  }

-- Create a VertexRank from the available data.
pickBestAndWorstVerts ::
  (Numeric a, Fractional a, Ord b) =>
  --(Container Vector a, Element a, Ord b, Transposable (Simplex a) (Simplex a)) =>
  (Vector a -> b)  -> -- A function from a point on the simplex to some orderable type.
  Simplex a        -> -- The simplex to consider.
  VertexRank a b      -- Vertex rank.
pickBestAndWorstVerts f s = runST $ do
  let
    nDims  = rows s
    points = tr s

  scores <- VM.new nDims
  forM_ [0..(nDims - 1)] (\i -> VM.unsafeWrite scores i (i, f $ points ! i))
  sortBy (comparing snd) scores
  let
    mvertex j = liftM (uncurry Vertex) (VM.unsafeRead scores j)
  h  <- mvertex 0
  sh <- mvertex 1
  l  <- mvertex (nDims - 1)

  return $ VertexRank h sh l (center $ allButIth (which l) s)

-- Increment a termination criterion by one iteration.
increment :: TerminationCriterion a -> TerminationCriterion a
increment p
  | isJust (view maxit p) = over maxit (fmap pred) p
  | otherwise             = p

-- Compute the center of the NM Simplex.
center ::
  (Numeric a, Fractional a) =>
  Simplex a ->
  Vector a
center s = runST $ do
  centroid <- VSM.new (rows s)

  forM_ [0..((rows s) - 1)] (\i -> do
    forM_ [0..((cols s) - 1)] (\j -> do
      let
        u = fromIntegral (cols s)
      VSM.unsafeWrite centroid i ((s `atIndex` (i,j)) / u) ))

  freeze centroid

-- Create a simplex with all points but the i-th.
allButIth :: (Numeric a) => Int -> Simplex a -> Simplex a
allButIth l s
  | l == 0        = subMatrix (0,1) (nDims, nDims) s
  | l == (cols s) = subMatrix (0,0) (nDims, nDims) s
  | otherwise     = ls ||| rs where
      nDims = rows s
      ls = subMatrix (0, l - 1) (nDims, l - 1)     s
      rs = subMatrix (0, l + 1) (nDims, nDims - l) s

-- The centroid of all points except whichever is worst.
goodCentroid :: (Numeric a, Fractional a, Ord b) =>
  VertexRank a b ->
  Simplex a      ->
  Vector a
goodCentroid r splx = center $ allButIth ((which . worstOf) r) splx

-- Get the i-th point of a simplex.
nthPoint :: (Numeric a, Element a) => Int -> Simplex a -> Vector a
nthPoint i s = flatten $ subMatrix (0, i) (rows s, 1) s

-- Get an estimate of simplex size by computing the sum of all
-- elements in the simplex, translated so that its centroid
-- lies at zero.
simplexSize :: (Numeric a, Fractional a) =>
  Simplex a ->
  a
simplexSize s = fastSimplexL1 (moveToOrigin s) where
  fastSimplexL1 m = runST $ do
    acc <- newSTRef 0
    mut <- unsafeThawMatrix m
    forM_ [0..((rows m) - 1)] $ \i -> do
      forM_ [0..((cols m) - 1)] $ \j -> do
        elem <- unsafeReadMatrix mut i j
        modifySTRef acc (+ elem)
    readSTRef acc

moveToOrigin :: (Numeric a, Fractional a) =>
  Simplex a ->
  Simplex a
moveToOrigin s = runSTMatrix $ do
  sm  <- thawMatrix s
  ctr <- unsafeThawVector (center s)

  forM_ [0..((rows s) - 1)] $ \i -> do
    c_entry <- unsafeReadVector ctr i
    forM_ [0..((cols s) - 1)] $ \j -> do
      m_entry <- unsafeReadMatrix sm i j
      unsafeWriteMatrix sm i j (m_entry - c_entry)

  return sm

-- Check if current simplex state and termination criterion indicate that we should
-- be finished.
isDone tc simplex = noMoreIters || closeEnough || neitherValid where
  noMoreIters  = (fromMaybe 1 (view maxit tc)) == 0
  thresh       = view threshold tc
  closeEnough  = fromMaybe False (fmap ((simplexSize simplex) <) thresh)
  neitherValid = (isNothing (view maxit tc)) && (isNothing (view threshold tc))

-- Update the i-th point of simplex sx with vector p.
updateSimplex :: (Numeric a) => Int -> Vector a -> Simplex a -> Simplex a
updateSimplex i v sx = runSTMatrix $ do
  sxm <- thawMatrix sx
  vm  <- unsafeThawVector v
  forM_ [0..((cols sx) - 1)] $ \d -> do
    updated <- unsafeReadVector vm d
    unsafeWriteMatrix sxm d i updated
  return sxm

-- Reflect the simplex's "worst" point around the hyperplane defined by the
-- other points. If this constitutes a significant improvement, then consider
-- "expanding" even farther in the same direction. Otherwise, give up.
reflectAndExpand :: forall a b s. (Num a, Numeric a, Fractional a, Ord b) =>
  NMStepParameters a              ->
  (Vector a -> b)                 ->
  STRef s (Simplex a)             ->
  STRef s (VertexRank a b)        ->
  ST    s Bool
reflectAndExpand params f s vrank = do
  r  <- readSTRef vrank
  sx <- readSTRef s

  let
    x0      = goodCentroid r sx
    x_best  = nthPoint ((which . bestOf) r) sx
    sx_best = (score . bestOf)     r -- Score of best x
    sx_next = (score . nextBestOf) r -- Score of next best x

    a   = view alpha params
    xr  = x0 + ((x0 - x_best) `scaleBy` a) -- x reflected
    sxr = f xr                             -- score of x reflected

  -- How does the reflected point fare?
  if (sxr < sx_next && sxr > sx_best)
    then do
    modifySTRef s (updateSimplex ((which . worstOf) r) xr)
    return True
    else do

    -- Consider expansion.
    if (sxr < sx_best)
      then do
      let
        g   = view gamma params
        xe  = x0 + ((xr - x0) `scaleBy` g)
        sxe = f xe

      -- Is expansion better?
      if (sxe < sxr)
        then do
        modifySTRef s (updateSimplex ((which . worstOf) r) xe)
        return True
        else do
        modifySTRef s (updateSimplex ((which . worstOf) r) xr)
        return True
      else do
      return False

-- Move the worst point of the simplex inward toward the centroid of all the
-- better points. If this doesn't result in an improvement, don't do anything.
contract :: forall a b s. (Numeric a, Fractional a, Ord b) =>
  NMStepParameters a       ->
  (Vector a -> b)          ->
  STRef s (Simplex a)      ->
  STRef s (VertexRank a b) ->
  ST s Bool
contract params f s vrank = do
  r  <- readSTRef vrank
  sx <- readSTRef s
  let
    x0       = goodCentroid r sx
    x_worst  = nthPoint ((which . worstOf) r) sx
    sx_worst = (score . worstOf) r

    rh  = view rho params
    xc  = x0 + ((x_worst - x0) `scaleBy` rh)
    sxc = f xc

  if (sxc < sx_worst)
    then do
    modifySTRef s (updateSimplex ((which . worstOf) r) xc)
    return True
    else do
    return False

-- Move all points of a simplex slightly towards the best point.
shrink :: forall a b s. (Num a, Fractional a, Storable a, Ord b) =>
  NMStepParameters a       ->
  (Vector a -> b)          ->
  STRef s (Simplex a)      ->
  STRef s (VertexRank a b) ->
  ST s Bool
shrink params f s vrank = do
  sx <- readSTRef s
  r  <- readSTRef vrank
  let
    sig    = view sigma r
    x_best = nthPoint ((which . bestOf) r) sx

  forM_ (\i -> do
            modifySTRef s (\simplex ->
                             if (i == ((which . bestOf) r))
                             then
                               simplex
                             else
                               let
                                 xi  = nthPoint i simplex
                                 xi' = x_best + (sig * (xi - x_best)) in
                                 updateSimplex i xi' simplex
                          )
        ) [0..((cols sx) - 1)]
  return True

-- |Predict the location of a minimum of an objective function using the Nelder-Mead simplex
--  search. This function permits tuning of all parameters, including the starting simplex.
nelderMeadFull :: (Num a, Fractional a, Storable a, Ord b) =>
  NMStepParameters a     -> -- ^Nelder-Mead step parameters
  TerminationCriterion a -> -- ^Reason we would terminate
  (Vector a -> b)        -> -- ^Objective function
  Simplex a              -> -- ^Starting simplex
  Vector a                  -- ^Predicted minimum
nelderMeadFull params term_conditions func initial = (center . runST) $ do

  simplex  <- newSTRef initial
  vtx_rank <- newSTRef (pickBestAndWorstVerts simplex)

  nelderMeadStep params term_conditions func simplex vtx_rank where

    nelderMeadStep p tc f s vrank = do
      -- Update rank and see if we're done.
      modifySTRef vrank (pickBestAndWorstVerts func)
      finished <- isDone <$> readSTRef s

      if (finished)
        then do
        rnk <- readSTRef vrank
        sx  <- readSTRef s
        return $ nthPoint ((which . bestOf) rnk) sx

        -- Attempt to improve simplex.
        else do
        reflected <- reflectAndExpand p f s vrank
        unless reflected $ do
          contracted <- contract p f s vrank
          unless contracted $ do
            shrink p f s vrank

        nelderMeadStep p (iterate tc) f s vrank

-- |Predict the location of a minimum of an objective function using the
--  Nelder-Mead simplex search. This version is like nelderMeadFull, but accepts
--  a starting point rather than an entire simplex.
nelderMeadAt :: (Num a, Fractional a, Storable a, Ord b) =>
  NMStepParameters a     -> -- ^Nelder-Mead step parameters
  TerminationCriterion a -> -- ^Reason we would terminate
  (Vector a -> b)        -> -- ^Objective function
  Vector a               -> -- ^Starting point.
  Vector a                  -- ^Predicted minimum.
nelderMeadAt p t f i = nelderMeadFull p t f (makeSimplexAround i e) where
  e = view epsilon p

-- |Minimize an objective function using the Nelder-Mead simplex search. This
--  invokation uses "default" step parameters.
nelderMead :: (Num a, Fractional a, Storable a, Ord b) =>
  TerminationCriterion a -> -- ^Reason we would terminate
  (Vector a -> b)        -> -- ^Objective function
  Vector a                -- ^Initial guess
nelderMead tc f x = nelderMeadAt defaultStep tc f x

{-
   Construct a simplex around a point by creating another vertex epsilon away from the initial
   guess along each vector of the canonical basis.
-}
makeSimplexAround :: (Element a) =>
  Vector a -> -- One vertex of the simplex
  a        -> -- Epsilon used to expand the simplex along the canonical basis
  Simplex a   -- Resultant simplex.
makeSimplexAround initial = fromBlocks [[asColumn initial, cloned + identity]]
  where
    dimension = size initial
    cloned    = asColumns (take dimension $ repeat initial)
    identity  = ident dimension

scaleBy :: (Num a) => Vector a -> a -> Vector a
scaleBy v a = runST $ do
  vm <- thawVector
  forM_ [0..((VS.length v) - 1)] $ \i -> do
    e <- unsafeReadVector vm i
    unsafeWriteVector vm i (e * a)
  freezeVector vm


{- |
Copyright: (c) 2020 Thomas Tuegel
SPDX-License-Identifier: BSD-3-Clause
Maintainer: Thomas Tuegel <ttuegel@mailbox.org>

-}

module Injection
    ( Injection (..)
    , Retraction (..)
    ) where

import Data.Complex (Complex ((:+)))
import Data.Dynamic (Dynamic, Typeable, fromDynamic, toDyn)
import Data.Fixed (Fixed, HasResolution)
import Data.Functor.Const (Const (..))
import Data.Functor.Identity (Identity (..))
import Data.List.NonEmpty (NonEmpty (..))
import Data.Maybe (maybeToList)
import Data.Monoid (Dual (..))
import Data.Monoid (Product (..))
import Data.Monoid (Sum (..))
import Data.Monoid (Any (..))
import Data.Monoid (All (..))
import qualified Data.Monoid as Monoid (First (..), Last (..))
import Data.Ord (Down (..))
import Data.Ratio (Ratio)
import qualified Data.Ratio as Ratio
import Data.Semigroup (Max (..), Min (..))
import qualified Data.Semigroup as Semigroup (First (..), Last (..))
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.Lazy as Lazy (Text)
import qualified Data.Text.Lazy as Text.Lazy
import Data.Void (Void)
import Numeric.Natural (Natural)

{- | @Injection@ describes a lossless conversion that includes one type in another.

The sole method of this class,

> inject :: from -> into

takes a value @input :: from@ and returns a value @output :: into@ which preserves all the information contained in the input.
Specifically, each @input@ is mapped to a /unique/ @output@.
In mathematical terminology, @inject@ is /injective/:

> inject a ≡ inject b → a ≡ b

The name of the class is derived from the mathematical term.

@Injection@ models the "is-a" relationship used in languages with subtypes (such as in object-oriented programming),
but an explicit cast with @inject@ is required in Haskell.

Although it is often possible to infer the type parameters of this class,
it is advisable to specify one or both of the parameters to @inject@
using a type signature or the @TypeApplications@ language extension.
Specifying the type parameters will give clearer error messages from the type checker in any case.

-}
class Injection from into where
    inject :: from -> into

{- | @Retraction@ undoes an 'Injection'.

Because 'Injection' is a lossless conversion, we can define a @Retraction@ which undoes it.
The method

> retract :: into -> Maybe from

is the (left) inverse of 'inject':

> retract (inject x) = Just x

'retract' is partial (returns 'Maybe') because the type @into@ may be larger than the type @from@;
that is, there may be values in @into@ which are not 'inject'-ed from @from@,
and in that case @retract@ may return 'Nothing'.

Although it is often possible to infer the type parameters of this class,
it is advisable to specify one or both of the parameters to @retract@
using a type signature or the @TypeApplications@ language extension.
Specifying the type parameters will give clearer error messages from the type checker in any case.

-}
class Injection from into => Retraction from into where
    retract :: into -> Maybe from

instance Injection a a where
    inject = id
    {-# INLINE inject #-}

instance Retraction a a where
    retract = Just
    {-# INLINE retract #-}

instance Typeable a => Injection a Dynamic where
    inject = toDyn
    {-# INLINE inject #-}

instance Typeable a => Retraction a Dynamic where
    retract = fromDynamic
    {-# INLINE retract #-}

instance Injection a b => Injection a (Maybe b) where
    inject = Just . inject
    {-# INLINE inject #-}

instance Retraction a b => Retraction a (Maybe b) where
    retract = \x -> x >>= retract @a @b
    {-# INLINE retract #-}

instance Injection a b => Injection (Maybe a) [b] where
    inject = maybeToList . fmap (inject @a @b)
    {-# INLINE inject #-}

instance Retraction a b => Retraction (Maybe a) [b] where
    retract [] = Just Nothing
    retract [b] = Just <$> retract @a @b b
    retract _ = Nothing

instance Injection Natural Integer where
    inject = toInteger
    {-# INLINE inject #-}

instance Retraction Natural Integer where
    retract x
        | x < 0 = Nothing
        | otherwise = Just (fromInteger x)
    {-# INLINE retract #-}

instance Injection Void any where
    inject = \case {}
    {-# INLINE inject #-}

-- | 'Text.unpack' is the canonical injection @'Text' -> 'String'@.
-- There is no injection in the other direction because 'String' can represent
-- invalid surrogate code points, but 'Text' cannot; for details, see
-- "Data.Text".
instance Injection Text String where
    inject = Text.unpack
    {-# INLINE inject #-}

-- | 'Text.Lazy.unpack' is the canonical injection @'Lazy.Text' -> 'String'@.
-- There is no injection in the other direction because 'String' can represent
-- invalid surrogate code points, but 'Lazy.Text' cannot; for details, see
-- "Data.Text.Lazy".
instance Injection Lazy.Text String where
    inject = Text.Lazy.unpack
    {-# INLINE inject #-}

instance Injection Text Lazy.Text where
    inject = Text.Lazy.fromStrict
    {-# INLINE inject #-}

instance Injection Lazy.Text Text where
    inject = Text.Lazy.toStrict
    {-# INLINE inject #-}

instance HasResolution a => Injection Integer (Fixed a) where
    inject = fromInteger
    {-# INLINE inject #-}

instance HasResolution a => Retraction Integer (Fixed a) where
    retract x = retract @Integer (toRational x)
    {-# INLINE retract #-}

instance Injection a (Const a b) where
    inject = Const
    {-# INLINE inject #-}

instance Injection (Const a b) a where
    inject = getConst
    {-# INLINE inject #-}

instance Injection Integer (Ratio Integer) where
    inject = fromInteger
    {-# INLINE inject #-}

instance Retraction Integer (Ratio Integer) where
    retract x
        | Ratio.denominator x == 1 = Just (Ratio.numerator x)
        | otherwise = Nothing
    {-# INLINE retract #-}

instance Num a => Injection a (Complex a) where
    inject = (:+ 0)
    {-# INLINE inject #-}

instance (Eq a, Num a) => Retraction a (Complex a) where
    retract (x :+ y)
      | y == 0 = Just x
      | otherwise = Nothing
    {-# INLINE retract #-}

instance Injection a (Identity a) where
    inject = Identity
    {-# INLINE inject #-}

instance Injection (Identity a) a where
    inject = runIdentity
    {-# INLINE inject #-}

instance Injection (NonEmpty a) [a] where
    inject (x :| xs) = x : xs
    {-# INLINE inject #-}

instance Retraction (NonEmpty a) [a] where
    retract (x : xs) = Just (x :| xs)
    retract [] = Nothing
    {-# INLINE retract #-}

instance Injection a (Down a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Down a) a where
    inject = \(Down a) -> a
    {-# INLINE inject #-}

instance Injection a (Product a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Product a) a where
    inject = getProduct
    {-# INLINE inject #-}

instance Injection a (Sum a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Sum a) a where
    inject = getSum
    {-# INLINE inject #-}

instance Injection a (Dual a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Dual a) a where
    inject = getDual
    {-# INLINE inject #-}

instance Injection a (Monoid.Last a) where
    inject = pure
    {-# INLINE inject #-}

instance Retraction a (Monoid.Last a) where
    retract = Monoid.getLast
    {-# INLINE retract #-}

instance Injection a (Monoid.First a) where
    inject = pure
    {-# INLINE inject #-}

instance Retraction a (Monoid.First a) where
    retract = Monoid.getFirst
    {-# INLINE retract #-}

instance Injection a (Semigroup.First a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Semigroup.First a) a where
    inject = Semigroup.getFirst
    {-# INLINE inject #-}

instance Injection a (Semigroup.Last a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Semigroup.Last a) a where
    inject = Semigroup.getLast
    {-# INLINE inject #-}

instance Injection a (Max a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Max a) a where
    inject = getMax
    {-# INLINE inject #-}

instance Injection a (Min a) where
    inject = pure
    {-# INLINE inject #-}

instance Injection (Min a) a where
    inject = getMin
    {-# INLINE inject #-}

instance Injection a (r -> a) where
    inject = const
    {-# INLINE inject #-}

instance Injection Bool Any where
    inject = Any
    {-# INLINE inject #-}

instance Injection Any Bool where
    inject = getAny
    {-# INLINE inject #-}

instance Injection Bool All where
    inject = All
    {-# INLINE inject #-}

instance Injection All Bool where
    inject = getAll
    {-# INLINE inject #-}


{-# LANGUAGE TypeFamilies #-}
-- | Tools for implementing (and debugging the use of) gradient descent schemes.
module HordeAd.Core.OptimizerTools
  ( updateWithGradient
  , gradientIsNil, minimumGradient, maximumGradient
  , ArgsAdam(..), defaultArgsAdam
  , StateAdam(..), initialStateAdam
  , liftMatrix43, updateWithGradientAdam
  ) where

import Prelude

import           Control.Monad.ST.Strict (runST)
import qualified Data.Array.DynamicS as OT
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as VM
import           HordeAd.Internal.OrthotopeOrphanInstances (liftVT2)
import           Numeric.LinearAlgebra (Element, Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM
import           Numeric.LinearAlgebra.Data (flatten)
import           Numeric.LinearAlgebra.Devel
  (MatrixOrder (..), liftMatrix, liftMatrix2, matrixFromVector, orderOf)

import HordeAd.Internal.Delta (Domains, isTensorDummy)

{-
60% of heap allocation in matrix- and vector-based MNIST
with simple gradient descent (no mini-batches) is performed
by @updateWithGradient.updateVector@ below

  let updateVector i r = i - HM.scale gamma r

due to allocating once in @scale@ and again in @-@
(and there seems to be one more allocation judging by the numbers).
Something like the following code would be needed to eliminate
one allocation, but it would requre the hmatrix maintainer to expose
internal modules.

import           Internal.Vectorized
  ( FunCodeSV (Scale), FunCodeVV (Sub), applyRaw, c_vectorMapValR
  , c_vectorZipR, createVector )

minusTimesGamma :: Storable r => r -> Vector r -> Vector r -> Vector r
minusTimesGamma gamma u v = unsafePerformIO $ do
  r <- createVector (dim0 u)
  pval <- newArray [gamma]
  (v `applyRaw` (r `applyRaw` id))
    (c_vectorMapValR (fromei Scale) pval)
    #| "minusTimesGamma1"
  free pval
  (u `applyRaw` (v `applyRaw` (r `applyRaw` id)))
    (c_vectorZipR (fromei Sub))
    #| "minusTimesGamma2"
  return r

BTW, a version with HM.Devel.zipVectorWith makes
the test twice slower and allocate twice more

  let updateVector = zipVectorWith (\i r -> i - gamma * r)

and a version with Vector.Storable makes the test thrice slower
and allocate thrice more

  let updateVector = V.zipWith (\i r -> i - gamma * r)

which is probably a bug in stream fusion which, additionally in this case,
can't fuse with anything and so can't pay for its overhead.

-}

updateWithGradient :: (Numeric r, Floating (Vector r))
                   => r -> Domains r -> Domains r -> Domains r
updateWithGradient gamma (params0, params1, params2, paramsX)
                         (gradient0, gradient1, gradient2, gradientX) =
  let updateVector i r = i - HM.scale gamma r
      !params0New = updateVector params0 gradient0
      update1 i r = if V.null r  -- eval didn't update it, would crash
                    then i
                    else updateVector i r
      !params1New = V.zipWith update1 params1 gradient1
      update2 i r = if HM.rows r <= 0  -- eval didn't update it, would crash
                    then i
                    else liftMatrix2 updateVector i r
      !params2New = V.zipWith update2 params2 gradient2
      updateX i r = if isTensorDummy r  -- eval didn't update it, would crash
                    then i
                    else liftVT2 updateVector i r
                      -- TODO: this is slow; add @liftArray2@ and use HM,
                      -- unless we move away from HM; similarly other OT calls
      !paramsXNew = V.zipWith updateX paramsX gradientX
  in (params0New, params1New, params2New, paramsXNew)

gradientIsNil :: (Eq r, Numeric r)
              => Domains r -> Bool
gradientIsNil (gradient0, gradient1, gradient2, gradientX) =
  V.all (== 0) gradient0
  && V.all V.null gradient1
  && V.all (\r -> HM.rows r <= 0) gradient2
  && V.all isTensorDummy gradientX

minimumGradient :: (Ord r, Numeric r)
                => Domains r -> r
minimumGradient (gradient0, gradient1, gradient2, gradientX) =
  min (if V.null gradient0 then 0 else HM.minElement gradient0)
      (min (if V.null gradient1 then 0
            else V.minimum (V.map HM.minElement gradient1))
           (min (if V.null gradient2 then 0
                 else V.minimum (V.map HM.minElement gradient2))
                (if V.null gradientX then 0
                 else V.minimum (V.map OT.minimumA gradientX))))

maximumGradient :: (Ord r, Numeric r)
                => Domains r -> r
maximumGradient (gradient0, gradient1, gradient2, gradientX) =
  max (if V.null gradient0 then 0 else HM.maxElement gradient0)
      (max (if V.null gradient1 then 0
            else V.maximum (V.map HM.maxElement gradient1))
           (max (if V.null gradient2 then 0
                 else V.maximum (V.map HM.maxElement gradient2))
                (if V.null gradientX then 0
                 else V.maximum (V.map OT.maximumA gradientX))))

data ArgsAdam r = ArgsAdam
  { alpha   :: r
  , betaOne :: r
  , betaTwo :: r
  , epsilon :: r
  }

-- The defaults taken from
-- https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
defaultArgsAdam :: Fractional r => ArgsAdam r
defaultArgsAdam = ArgsAdam
  { alpha = 0.001
  , betaOne = 0.9
  , betaTwo = 0.999
  , epsilon = 1e-7
  }

data StateAdam r = StateAdam
  { tAdam :: Int  -- iteration count
  , mAdam :: Domains r
  , vAdam :: Domains r
  }

-- The arguments are just sample params0, for dimensions.
zeroParameters :: Numeric r
               => Domains r -> Domains r
zeroParameters (params0, params1, params2, paramsX) =
  let zeroVector v = runST $ do
        vThawed <- V.thaw v
        VM.set vThawed 0
        V.unsafeFreeze vThawed
  in ( zeroVector params0
     , V.map zeroVector params1
     , V.map (liftMatrix zeroVector) params2
     , V.map (\a -> OT.constant (OT.shapeL a) 0) paramsX )  -- fast allright

initialStateAdam :: Numeric r
                 => Domains r -> StateAdam r
initialStateAdam parameters0 =
  let zeroP = zeroParameters parameters0
  in StateAdam
       { tAdam = 0
       , mAdam = zeroP
       , vAdam = zeroP
       }

-- | Application of a vector function on the flattened matrices elements.
liftMatrix43 :: ( Numeric a, Numeric b, Numeric c, Numeric d
                , Element x, Element y, Element z )
             => (Vector a -> Vector b -> Vector c -> Vector d
                 -> (Vector x, Vector y, Vector z))
             -> Matrix a -> Matrix b -> Matrix c -> Matrix d
             -> (Matrix x, Matrix y, Matrix z)
liftMatrix43 f m1 m2 m3 m4 =
  let sz@(r, c) = HM.size m1
      rowOrder = orderOf m1
        -- checking @m4@ (gradient) makes RNN LL test much faster and BB slower
        -- so this needs much more benchmarking and understading to tweak
  in if sz == HM.size m2 && sz == HM.size m3 && sz == HM.size m4
     then
       let (vx, vy, vz) = case rowOrder of
             RowMajor -> f (flatten m1) (flatten m2) (flatten m3) (flatten m4)
             ColumnMajor -> f (flatten (HM.tr' m1)) (flatten (HM.tr' m2))
                              (flatten (HM.tr' m3)) (flatten (HM.tr' m4))
       in ( matrixFromVector rowOrder r c vx
          , matrixFromVector rowOrder r c vy
          , matrixFromVector rowOrder r c vz
          )
     else error $ "nonconformant matrices in liftMatrix43: "
                  ++ show (HM.size m1, HM.size m2, HM.size m3, HM.size m4)

-- TOOD: make sure this is not worse that OT.zipWith3A when transposing
-- between each application or that we never encounter such situations
--
-- | Application of a vector function on the flattened arrays elements.
liftArray43 :: ( Numeric a, Numeric b, Numeric c, Numeric d
               , Numeric x, Numeric y, Numeric z )
            => (Vector a -> Vector b -> Vector c -> Vector d
                -> (Vector x, Vector y, Vector z))
            -> OT.Array a -> OT.Array b -> OT.Array c -> OT.Array d
            -> (OT.Array x, OT.Array y, OT.Array z)
liftArray43 f m1 m2 m3 m4 =
  let sz = OT.shapeL m1
  in if sz == OT.shapeL m2 && sz == OT.shapeL m3 && sz == OT.shapeL m4
     then
       let (vx, vy, vz) = f (OT.toVector m1) (OT.toVector m2)
                            (OT.toVector m3) (OT.toVector m4)
       in ( OT.fromVector sz vx
          , OT.fromVector sz vy
          , OT.fromVector sz vz
          )
     else error
          $ "nonconformant arrays in liftArray43: "
            ++ show (OT.shapeL m1, OT.shapeL m2, OT.shapeL m3, OT.shapeL m4)

updateWithGradientAdam
  :: forall r. (Numeric r, Floating r, Floating (Vector r))
  => ArgsAdam r -> StateAdam r -> Domains r -> Domains r
  -> (Domains r, StateAdam r)
updateWithGradientAdam ArgsAdam{..}
                       StateAdam{ tAdam
                                , mAdam = (mAdam0, mAdam1, mAdam2, mAdamX)
                                , vAdam = (vAdam0, vAdam1, vAdam2, vAdamX)
                                }
                       (params0, params1, params2, paramsX)
                       (gradient0, gradient1, gradient2, gradientX) =
  let tAdamNew = tAdam + 1
      oneMinusBeta1 = 1 - betaOne
      oneMinusBeta2 = 1 - betaTwo
      updateVector :: Vector r -> Vector r
                   -> Vector r -> Vector r
                   -> (Vector r, Vector r, Vector r)
      updateVector mA vA p g =
        let mANew = HM.scale betaOne mA + HM.scale oneMinusBeta1 g
            vANew = HM.scale betaTwo vA + HM.scale oneMinusBeta2 (g * g)
            alphat = alpha * sqrt (1 - betaTwo ^ tAdamNew)
                             / (1 - betaOne ^ tAdamNew)
        in ( mANew
           , vANew
           , p - HM.scale alphat mANew
                 / (sqrt vANew + HM.scalar epsilon) )  -- the @scalar@ is safe
                      -- @addConstant@ would be better, but it's not exposed
      (!mAdam0New, !vAdam0New, !params0New) =
        updateVector mAdam0 vAdam0 params0 gradient0
      update1 mA vA p g = if V.null g  -- eval didn't update it, would crash
                          then (mA, vA, p)
                          else updateVector mA vA p g
      (!mAdam1New, !vAdam1New, !params1New) =
        V.unzip3 $ V.zipWith4 update1 mAdam1 vAdam1 params1 gradient1
      update2 mA vA p g = if HM.rows g <= 0  -- eval didn't update it; crash
                          then (mA, vA, p)
                          else liftMatrix43 updateVector mA vA p g
      (!mAdam2New, !vAdam2New, !params2New) =
        V.unzip3 $ V.zipWith4 update2 mAdam2 vAdam2 params2 gradient2
      updateX mA vA p g = if isTensorDummy g  -- eval didn't update it
                          then (mA, vA, p)
                          else liftArray43 updateVector mA vA p g
      (!mAdamXNew, !vAdamXNew, !paramsXNew) =
        V.unzip3 $ V.zipWith4 updateX mAdamX vAdamX paramsX gradientX
  in ( (params0New, params1New, params2New, paramsXNew)
     , StateAdam { tAdam = tAdamNew
                 , mAdam = (mAdam0New, mAdam1New, mAdam2New, mAdamXNew)
                 , vAdam = (vAdam0New, vAdam1New, vAdam2New, vAdamXNew)
                 }
     )


{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# HLINT ignore "Use camelCase"      #-}
{-|
Module      : Grenade.Layers.Internal.Add
Description : Fast functions for performing convolutions and helper functions
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Internal.Convolution (
  -- * im2col functions
    im2col
  , col2im
  , col2vid
  , vid2col

  -- * convolution functions
  , forwardConv2d
  , forwardBiasConv2d
  , backwardConv2d
  , backwardBiasConv2d
  ) where

import qualified Data.Vector.Storable        as U (unsafeFromForeignPtr0,
                                                   unsafeToForeignPtr0)

import           Foreign                     (mallocForeignPtrArray,
                                              withForeignPtr)
import           Foreign.Ptr                 (Ptr)
import           Grenade.Types
import           Numeric.LinearAlgebra       (Matrix, Vector, cols, flatten,
                                              rows)
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as U
import           System.IO.Unsafe            (unsafePerformIO)

import Grenade.Layers.Internal.Hmatrix

-- | Rearrange the matrix columns into blocks, calculates the number of channels automatically.
col2vid :: Int            -- ^ kernel rows
        -> Int            -- ^ kernel columns
        -> Int            -- ^ stride rows
        -> Int            -- ^ stride columns
        -> Int            -- ^ input height
        -> Int            -- ^ input width
        -> Matrix RealNum -- ^ input matrix
        -> Matrix RealNum -- ^ output matrix
col2vid kernelRows kernelColumns strideRows strideColumns height width dataCol =
  let channels = cols dataCol `div` (kernelRows * kernelColumns)
      outRows  = (height - kernelRows) `div` strideRows + 1
      outCols  = (width - kernelColumns) `div` strideColumns + 1
  in  col2im_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 outRows outCols dataCol

-- | Rearrange the matrix columns into blocks, assumes there is only one channel
col2im :: Int            -- ^ kernel rows
       -> Int            -- ^ kernel columns
       -> Int            -- ^ stride rows
       -> Int            -- ^ stride columns
       -> Int            -- ^ input height
       -> Int            -- ^ input width
       -> Matrix RealNum -- ^ input matrix
       -> Matrix RealNum -- ^ output matrix
col2im kernelRows kernelColumns strideRows strideColumns height width dataCol =
  let channels = 1
      outRows  = (height - kernelRows) `div` strideRows + 1
      outCols  = (width - kernelColumns) `div` strideColumns + 1
  in  col2im_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 outRows outCols dataCol

col2im_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im_c channels height width kernelRows kernelColumns strideRows strideColumns padl padt outRows outCols dataCol =
  let vec = flatten dataCol
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (height * width * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        col2im_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns padt padl outRows outCols outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (height * width * channels)
    return $ U.matrixFromVector U.RowMajor (height * channels) width matVec
{-# INLINE col2im_c #-}

foreign import ccall unsafe
    col2im_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()

-- | Rearrange the matrix blocks into columns, calculates the number of channels automatically.
vid2col :: Int            -- ^ kernel rows
        -> Int            -- ^ kernel columns
        -> Int            -- ^ stride rows
        -> Int            -- ^ stride columns
        -> Int            -- ^ input height
        -> Int            -- ^ input width
        -> Matrix RealNum -- ^ input matrix
        -> Matrix RealNum -- ^ output matrix
vid2col kernelRows kernelColumns strideRows strideColumns height width dataVid =
  let channels = rows dataVid `div` height
      rowOut          = (height - kernelRows) `div` strideRows + 1
      colOut          = (width - kernelColumns) `div` strideColumns + 1
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 dataVid rowOut colOut

-- | Rearrange the matrix blocks into columns, assumes there is only one channel
im2col :: Int            -- ^ kernel rows
       -> Int            -- ^ kernel columns
       -> Int            -- ^ stride rows
       -> Int            -- ^ stride columns
       -> Matrix RealNum -- ^ input matrix
       -> Matrix RealNum -- ^ output matrix
im2col kernelRows kernelColumns strideRows strideColumns dataIm =
  let channels = 1
      height = rows dataIm
      width  = cols dataIm
      rowOut          = (height - kernelRows) `div` strideRows + 1
      colOut          = (width - kernelColumns) `div` strideColumns + 1
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns 0 0 dataIm rowOut colOut

im2col_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Int -> Int -> Matrix RealNum
im2col_c channels height width kernelRows kernelColumns strideRows strideColumns padl padt dataIm outRows outCols =
  let vec             = flatten dataIm
      kernelSize      = kernelRows * kernelColumns
      numberOfPatches = outRows * outCols
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (numberOfPatches * kernelSize * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        im2col_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns padt padl outRows outCols outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (numberOfPatches * kernelSize * channels)
    return $ U.matrixFromVector U.RowMajor (kernelSize * channels) numberOfPatches matVec
{-# INLINE im2col_c #-}

foreign import ccall unsafe
    im2col_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()

foreign import ccall unsafe
    in_place_add_per_channel_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Ptr RealNum -> IO ()

-- | Efficient implementation of convolutions using the im2col trick as popularised by Caffe
forwardConv2d :: Matrix RealNum -- ^ input
              -> Int            -- ^ input channels
              -> Int            -- ^ input rows
              -> Int            -- ^ input columns
              -> Matrix RealNum -- ^ kernel
              -> Int            -- ^ kernel filters
              -> Int            -- ^ kernel rows
              -> Int            -- ^ kernel columns
              -> Int            -- ^ stride height
              -> Int            -- ^ stride width
              -> Int            -- ^ output rows
              -> Int            -- ^ output columns
              -> Int            -- ^ pad left
              -> Int            -- ^ pad top
              -> Matrix RealNum -- ^ output of convolution
forwardConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols outRows outCols padLeft padTop
  = let dataCol  = im2col_c channels rows cols kernelRows kernelCols strideRows strideCols padLeft padTop input outRows outCols
        gemmM    = LA.tr kernel LA.<> dataCol
        dataVid  = reshapeMatrix (filters * outRows) outCols gemmM
    in  dataVid

-- | Efficient implementation of convolutions using the im2col trick as popularised by Caffe
--   also adds some bias to the functions, note that it does it in place to prevent unnecessary allocation.
forwardBiasConv2d :: Matrix RealNum -- ^ input
                  -> Int            -- ^ input channels
                  -> Int            -- ^ input rows
                  -> Int            -- ^ input columns
                  -> Vector RealNum -- ^ bias
                  -> Matrix RealNum -- ^ kernel
                  -> Int            -- ^ kernel filters
                  -> Int            -- ^ kernel rows
                  -> Int            -- ^ kernel columns
                  -> Int            -- ^ stride height
                  -> Int            -- ^ stride width
                  -> Int            -- ^ output rows
                  -> Int            -- ^ output columns
                  -> Int            -- ^ pad left
                  -> Int            -- ^ pad top
                  -> Matrix RealNum -- ^ output of convolution
forwardBiasConv2d input channels rows cols bias kernel filters kernelRows kernelCols strideRows strideCols outRows outCols padLeft padTop =
  let outSize         = outRows * outCols * filters
  in  unsafePerformIO $ do
    let dataCol  = im2col_c channels rows cols kernelRows kernelCols strideRows strideCols padLeft padTop input outRows outCols
        gemmM    = LA.tr kernel LA.<> dataCol
        dataVid  = flatten gemmM
        (xPtr, _) = U.unsafeToForeignPtr0 dataVid
        (bPtr, _) = U.unsafeToForeignPtr0 bias

    withForeignPtr xPtr $ \xPtr' ->
      withForeignPtr bPtr $ \bPtr' ->
        in_place_add_per_channel_cpu xPtr' filters outRows outCols bPtr'

    let matVec = U.unsafeFromForeignPtr0 xPtr outSize
    return $ U.matrixFromVector U.RowMajor (outRows * filters) outCols matVec

-- | Efficient implementation of the backward pass for convolutions using the im2col trick as popularised by Caffe
backwardConv2d :: Matrix RealNum -- ^ input
               -> Int            -- ^ input channels
               -> Int            -- ^ input rows
               -> Int            -- ^ input columns
               -> Matrix RealNum -- ^ kernel
               -> Int            -- ^ kernel filters
               -> Int            -- ^ kernel rows
               -> Int            -- ^ kernel columns
               -> Int            -- ^ stride height
               -> Int            -- ^ stride width
               -> Int            -- ^ pad left
               -> Int            -- ^ pad top
               -> Int            -- ^ pad right
               -> Int            -- ^ pad bottom
               -> Matrix RealNum -- ^ derivative wrt out of layer
               -> Int            -- ^ output rows
               -> Int            -- ^ output columns
               -> (Matrix RealNum, Matrix RealNum)    -- ^ (derivated wrt input, derivative wrt kernel)
backwardConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols padl padt padr padb dout outRows outCols =
  let dout'  = LA.reshape (outRows * outCols) $ LA.flatten dout
      dX_col = kernel LA.<> dout'
      rowCol = (div (rows + padt + padb - kernelRows) strideRows) + 1
      colCol = (div (cols + padl + padr - kernelCols) strideCols) + 1
      dX     = col2im_c channels rows cols kernelRows kernelCols strideRows strideCols padl padt rowCol colCol dX_col

      x_col  = im2col_c channels rows cols kernelRows kernelCols strideRows strideCols padl padt input outRows outCols
      dw_col = dout' LA.<> (LA.tr x_col)
      dw     = LA.tr $ reshapeMatrix filters (kernelRows * kernelCols * channels) dw_col
  in (dX, dw)

-- | Efficient implementation of the backward pass for convolutions using the im2col trick as popularised by Caffe
backwardBiasConv2d :: Matrix RealNum    -- ^ input
                   -> Int -- ^ input channels
                   -> Int -- ^ input rows
                   -> Int -- ^ input columns
                   -> Matrix RealNum -- ^ kernel
                   -> Int -- ^ kernel filters
                   -> Int -- ^ kernel rows
                   -> Int -- ^ kernel columns
                   -> Int -- ^ stride height
                   -> Int -- ^ stride width
                   -> Int -- ^ pad left
                   -> Int -- ^ pad top
                   -> Int -- ^ pad right
                   -> Int -- ^ pad bottom
                   -> Matrix RealNum -- ^ derivative wrt out of layer
                   -> Int -- ^ output rows
                   -> Int -- ^ output columns
                   -> (Matrix RealNum, Matrix RealNum, Vector RealNum)  -- ^ (derivated wrt input, derivative wrt kernel, derivative wrt bias)
backwardBiasConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols padl padt padr padb dout outRows outCols =
  let (dX, dw) = backwardConv2d input channels rows cols kernel filters kernelRows kernelCols strideRows strideCols padl padt padr padb dout outRows outCols
      dB       = sum_over_channels_c dout filters outRows outCols
  in (dX, dw, dB)

sum_over_channels_c :: Matrix RealNum -> Int -> Int -> Int -> Vector RealNum
sum_over_channels_c mat channels rows cols = unsafePerformIO $ do
  let (xPtr, _) = U.unsafeToForeignPtr0 . flatten $ mat
  outPtr <- mallocForeignPtrArray channels

  withForeignPtr xPtr $ \xPtr' ->
    withForeignPtr outPtr $ \outPtr' ->
      sum_over_channels_cpu xPtr' channels rows cols outPtr'

  return $ U.unsafeFromForeignPtr0 outPtr channels

foreign import ccall unsafe
    sum_over_channels_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Ptr RealNum -> IO ()


{-# LANGUAGE CPP, OverloadedStrings, TypeOperators #-}

module Main where

import Control.Exception
import qualified Data.Array.Accelerate as A
import qualified Data.Array.Accelerate.Interpreter as ALI
import Data.Array.Accelerate.Math.Hilbert
import Data.Array.Accelerate.Math.Qtfd
import Data.Array.Accelerate.Math.WindowFunc
import Data.Attoparsec.Text
import ParseArgs
import ParseFile
import System.Directory
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
import qualified Data.Array.Accelerate.LLVM.Native as ALN
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
import qualified Data.Array.Accelerate.LLVM.PTX as ALP
#endif
import Data.Array.Accelerate.Array.Sugar as S
import qualified Data.ByteString.Builder as BB
import Data.Complex
import qualified Data.Double.Conversion.Convertable as DC
import Data.List
import Data.Monoid
import qualified Data.Text as T
import qualified Data.Text.IO as TI
import qualified Data.Vector as V
import GHC.Float
import Math.Gamma
import System.IO

main :: IO ()
main = do
  args <- opts
  makeQTFDAll args

-- | Transform data to text(using Float conversion. It is much faster) and write it to file.

writeData ::
  Handle            -- ^ output file
  -> [Float]       -- ^ Data to write
  -> Int            -- ^ number of element
  -> Int            -- ^ elements per row
  -> IO ()
writeData _ [] _ _  = return ()
writeData file (x:xs) w n = do
  TI.hPutStr file $ DC.toPrecision 5 x <> " "
  if (n `mod` w) == 0
  then do
    TI.hPutStr file "\n"
    writeData file xs w (n + 1)
  else writeData file xs w (n + 1)

makeString ::
     Handle
  -> [Float]       -- ^ Data to write
  -> Int            -- ^ number of element
  -> Int            -- ^ elements per row
  -> BB.Builder
  -> IO ()
makeString _ [] _ _ _  = return ()
makeString file (x:xs) w n acc =
  if (n `mod` w) == 0
  then BB.hPutBuilder file (acc <> DC.toPrecision 5 x <> " " <> "\n") >> makeString file xs w (n+1) ""
  else makeString file xs w (n+1) (acc <> DC.toPrecision 5 x <> " ")

-- | make Pseudo Wigner-Ville transform to all files in a directory.

makeQTFDAll :: Opts -> IO ()
makeQTFDAll opts = do
  let path = getPath opts
      nosAVGflag = subAVGflag opts
  files <- listDirectory path
  setCurrentDirectory path
  let fFiles = filter (isSuffixOf ".txt") files
      sAVGflag = A.constant $ not nosAVGflag
  case opts of
    (OptsPWV _ dev wlen wfun _) ->
      do
        if (even wlen)
        then error "Window length maust be odd !"
        else
          do
            let window = makeWindow wfun (A.unit $ A.constant wlen) :: A.Acc (A.Array A.DIM1 Float)
            mapM_ (makePWV dev window sAVGflag) fFiles
    (OptsWV path dev nosAVGflag) -> mapM_ (makeWV dev sAVGflag) fFiles
    (OptsCW path dev sigma uWindow uWindowFunc nWindow nWindowFunc normalise nosAVGflag) ->
      do
        if (even uWindow || even nWindow)
        then error "Window length maust be odd !"
        else
          let uWindowArray = makeWindow uWindowFunc (A.unit $ A.constant uWindow) :: A.Acc (A.Array A.DIM1 Float)
              nWindowArray = makeWindow nWindowFunc (A.unit $ A.constant nWindow) :: A.Acc (A.Array A.DIM1 Float)
              bothRectangular = uWindowFunc == Rect && nWindowFunc == Rect
          in mapM_ (makeCW dev sAVGflag sigma uWindowArray nWindowArray normalise bothRectangular) fFiles
    (OptsBJ path dev alpha uWindow uWindowFunc nWindow nWindowFunc nosAVGflag) ->
      do
        if (even uWindow || even nWindow)
        then error "Window length maust be odd !"
        else
          let uWindowArray = makeWindow uWindowFunc (A.unit $ A.constant uWindow) :: A.Acc (A.Array A.DIM1 Float)
              nWindowArray = makeWindow nWindowFunc (A.unit $ A.constant nWindow) :: A.Acc (A.Array A.DIM1 Float)
              bothRectangular = uWindowFunc == Rect && nWindowFunc == Rect
          in mapM_ (makeBJ dev sAVGflag alpha uWindowArray nWindowArray bothRectangular) fFiles
    (OptsEMBD path dev alpha beta uWindow uWindowFunc nWindow nWindowFunc normalise nosAVGflag) ->
      do
        if (even uWindow || even nWindow)
        then error "Window length maust be odd !"
        else let gammas = makeGammaArray beta uWindow
                 uWindowArray = makeWindow uWindowFunc (A.unit $ A.constant uWindow) :: A.Acc (A.Array A.DIM1 Float)
                 nWindowArray = makeWindow nWindowFunc (A.unit $ A.constant nWindow) :: A.Acc (A.Array A.DIM1 Float)
                 bothRectangular = uWindowFunc == Rect && nWindowFunc == Rect
             in mapM_ (makeEMBD dev sAVGflag alpha gammas uWindowArray nWindowArray bothRectangular normalise) fFiles
  return ()

makeCW :: CalcDev                   -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool                     -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                          -- ^ sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> FilePath                       -- ^ Name of file with data.
  -> IO ()
makeCW dev sAVGflag sigma uWindowArray nWindowArray normalise bothRect file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startCW dev file sAVGflag sigma uWindowArray nWindowArray normalise bothRect) dataF

makeBJ :: CalcDev                   -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool                     -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                          -- ^ sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if Both windows are rectangular
  -> FilePath                       -- ^ Name of file with data.
  -> IO ()
makeBJ dev sAVGflag alpha uWindowArray nWindowArray bothRect file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startBJ dev file sAVGflag alpha uWindowArray nWindowArray bothRect) dataF

-- | Make Pseudo Wigner-Ville transform to all columns in the given file

makePWV ::
  CalcDev                            -- ^ Calculation device - CPU or GPU
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> A.Exp Bool                      -- ^ Apply subtraction of the mean value from all elements in a column.
  -> FilePath                        -- ^ Name of file with data.
  -> IO ()
makePWV dev window sAVGflag file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startPWV dev window file sAVGflag) dataF

--makeEMBD dev alpha gammas uWindowArray nWindowArray bothRectangular normalise

makeEMBD :: CalcDev                   -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool                     -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                          -- ^ alpha
  -> [Float]                        -- ^ gammas
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> FilePath                       -- ^ Name of file with data.
  -> IO ()
makeEMBD dev sAVGflag alpha gammas uWindowArray nWindowArray normalise bothRect file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startEMBD dev file sAVGflag alpha gammas uWindowArray nWindowArray normalise bothRect) dataF

-- | Make Pseudo Wigner-Ville transform to the given column.
-- At first it makes subtraction of average value (if supAVGflag) and applies hilbert transform to make analytic signal. After that it applies Pseudo-Wigner transftorm
-- and save data to file with

startPWV ::
  CalcDev                          -- ^ Calculation device - CPU or GPU
  -> A.Acc (A.Array A.DIM1 Float) -- ^ Smoothing window in time-domain. Length must be odd.
  -> FilePath                      -- ^ Name of file with data.
  -> A.Exp Bool                    -- ^ Apply subtraction of the mean value from all elements in a column.
  -> (T.Text,[Float])             -- ^ Name of column and parsed data from column
  -> IO ()
startPWV dev window oldName sAVGflag (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      pData = A.fromList (A.Z A.:. leng) dataF
      appPWV = (pWignerVille window) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appPWV pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appPWV pData
#endif
                    CPU -> ALI.run1 appPWV pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $! processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

-- | Make Wigner-Ville transform to all columns in the given file

makeWV ::
  CalcDev        -- ^ Calculation device - CPU or GPU
  -> A.Exp Bool  -- ^ Apply subtraction of the mean value from all elements in a column.
  -> FilePath    -- ^ Name of file with data.
  -> IO ()
makeWV dev sAVGflag file = do
  putStrLn $ "processing " ++ file ++ "..."
  text <- TI.readFile file
  let parseRes = parseOnly parseFile text
  case parseRes of
    (Left errStr) -> error $ errStr
    (Right dataF) -> mapM_ (startWV dev file sAVGflag) dataF

-- | Make Wigner-Ville transform to the given column.
-- At first it makes subtraction of average value (if supAVGflag) and applies hilbert transform to make analytic signal. After that it applies Pseudo-Wigner transftorm
-- and save data to file with

startWV ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startWV dev oldName sAVGflag (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      pData = A.fromList (A.Z A.:. leng) dataF
      appWV = wignerVilleNew . hilbert . supAVG sAVGflag
      appWVGPU = wignerVille . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appWV pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appWV pData
#endif
                    CPU -> ALI.run1 appWV pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  hSetBinaryMode file True >> hSetBuffering file LineBuffering
  -- let resultBsBuilder = makeString pList leng 1 ""
  onException (makeString file pList leng 1 "") (removeFile newFName) >> hFlush file
  hClose file
  putStrLn "Done !"

startCW ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float               -- ^ Sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startCW dev oldName sAVGflag sigma uWindowArray nWindowArray normalise bothRect (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      uWindow = A.length uWindowArray
      nWindow = A.length nWindowArray
      mWindowArrays = if bothRect then Nothing else Just (uWindowArray,nWindowArray)
      pData = A.fromList (A.Z A.:. leng) dataF
      appCW = (choiWilliams (A.constant sigma) mWindowArrays uWindow nWindow normalise) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appCW pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appCW pData
#endif
                    CPU -> ALI.run1 appCW pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

startBJ ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float               -- ^ Sigma
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if Both windows are rectangular
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startBJ dev oldName sAVGflag alpha uWindowArray nWindowArray bothRect (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      uWindow = A.length uWindowArray
      nWindow = A.length nWindowArray
      mWindowArrays = if bothRect then Nothing else Just (uWindowArray,nWindowArray)
      pData = A.fromList (A.Z A.:. leng) dataF
      appCW = (bornJordan (A.constant alpha) mWindowArrays uWindow nWindow) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appCW pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appCW pData
#endif
                    CPU -> ALI.run1 appCW pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

startEMBD ::
  CalcDev                 -- ^ Calculation device - CPU or GPU
  -> FilePath             -- ^ Name of file with data.
  -> A.Exp Bool           -- ^ Apply subtraction of the mean value from all elements in a column.
  -> Float                -- ^ Alpha
  -> [Float]              -- ^ Gammas
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in frequensy-domain. Length must be odd.
  -> A.Acc (A.Array A.DIM1 Float)   -- ^ Smoothing window in time-domain. Length must be odd.
  -> Bool                           -- ^ if the core array should be normalised
  -> Bool                           -- ^ if Both windows are rectangular
  -> (T.Text,[Float])    -- ^ Name of column and parsed data from column
  -> IO ()
startEMBD dev oldName sAVGflag alpha gammas uWindowArray nWindowArray normalise bothRect (column_name,dataF) = do
  let newFName = oldName ++ "-" ++ T.unpack column_name ++ ".txt"
  putStr $ "   Creating file " ++ newFName ++ " ..."
  let leng = length dataF
      uWindow = A.length uWindowArray
      nWindow = A.length nWindowArray
      gammasArr = A.use $ A.fromList (A.Z A.:. (length gammas)) gammas
      mWindowArrays = if bothRect then Nothing else Just (uWindowArray,nWindowArray)
      pData = A.fromList (A.Z A.:. leng) dataF
      appEMBD = (eModifiedB (A.constant alpha) gammasArr mWindowArrays uWindow nWindow normalise) . hilbert . supAVG sAVGflag
      processed = case dev of
#ifdef ACCELERATE_LLVM_NATIVE_BACKEND
                    CPU -> ALN.run1 appEMBD pData
#endif
#ifdef ACCELERATE_LLVM_PTX_BACKEND
                    GPU -> ALP.run1 appEMBD pData
#endif
                    CPU -> ALI.run1 appEMBD pData
                    GPU -> error "Compiled without GPU support"
      pList = S.toList $  processed
  file <- openFile newFName WriteMode
  onException (writeData file pList leng 1) (removeFile newFName)
  hClose file
  putStrLn "Done !"

-- | Subtraction of average value from all elements of column

supAVG :: A.Exp Bool -> A.Acc (Array DIM1 Float) -> A.Acc (Array DIM1 Float)
supAVG flag arr =
  A.acond flag (A.map (\x -> x - avg) arr) arr
  where leng = A.length arr
        avg = (A.the $ A.sum arr)/(A.fromIntegral leng)

makeGammaArray :: Float -> Int -> [Float]
makeGammaArray beta wLength =
  let dwLength = fromIntegral wLength :: Float
      v = [(-0.5), ((-0.5) + 1.0/dwLength)..0.5]
      s = map (\x -> beta :+ pi*x) v
      f = map (\x -> abs $ ((magnitude $ gamma x) ^ 2)/((gamma beta) ^ 2)) s
      n = floor $ (fromIntegral $ length f)/2.0
      frst = Data.List.take (n-1) f
      secnd  = Data.List.drop n f
  in secnd ++ frst


{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, OverloadedStrings,
    RecordWildCards #-}

-- |
-- Module    : Statistics.Resampling.Bootstrap
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The bootstrap method for statistical inference.

module Statistics.Resampling.Bootstrap
    (
      Estimate(..)
    , bootstrapBCA
    , scale
    -- * References
    -- $references
    ) where

import Control.Applicative ((<$>), (<*>))
import Control.DeepSeq (NFData)
import Control.Exception (assert)
import Control.Monad.Par (parMap, runPar)
import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data)
import Data.Typeable (Typeable)
import Data.Vector.Unboxed ((!))
import GHC.Generics (Generic)
import Statistics.Distribution (cumulative, quantile)
import Statistics.Distribution.Normal
import Statistics.Resampling (Resample(..), jackknife)
import Statistics.Sample (mean)
import Statistics.Types (Estimator, Sample)
import qualified Data.Vector.Unboxed as U
import qualified Statistics.Resampling as R

-- | A point and interval estimate computed via an 'Estimator'.
data Estimate = Estimate {
      estPoint           :: {-# UNPACK #-} !Double
    -- ^ Point estimate.
    , estLowerBound      :: {-# UNPACK #-} !Double
    -- ^ Lower bound of the estimate interval (i.e. the lower bound of
    -- the confidence interval).
    , estUpperBound      :: {-# UNPACK #-} !Double
    -- ^ Upper bound of the estimate interval (i.e. the upper bound of
    -- the confidence interval).
    , estConfidenceLevel :: {-# UNPACK #-} !Double
    -- ^ Confidence level of the confidence intervals.
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON Estimate
instance ToJSON Estimate

instance Binary Estimate where
    put (Estimate w x y z) = put w >> put x >> put y >> put z
    get = Estimate <$> get <*> get <*> get <*> get
instance NFData Estimate

-- | Multiply the point, lower bound, and upper bound in an 'Estimate'
-- by the given value.
scale :: Double                 -- ^ Value to multiply by.
      -> Estimate -> Estimate
scale f e@Estimate{..} = e {
                           estPoint = f * estPoint
                         , estLowerBound = f * estLowerBound
                         , estUpperBound = f * estUpperBound
                         }

estimate :: Double -> Double -> Double -> Double -> Estimate
estimate pt lb ub cl =
    assert (lb <= ub) .
    assert (cl > 0 && cl < 1) $
    Estimate { estPoint = pt
             , estLowerBound = lb
             , estUpperBound = ub
             , estConfidenceLevel = cl
             }

data T = {-# UNPACK #-} !Double :< {-# UNPACK #-} !Double
infixl 2 :<

-- | Bias-corrected accelerated (BCA) bootstrap. This adjusts for both
-- bias and skewness in the resampled distribution.
bootstrapBCA :: Double          -- ^ Confidence level
             -> Sample          -- ^ Sample data
             -> [Estimator]     -- ^ Estimators
             -> [Resample]      -- ^ Resampled data
             -> [Estimate]
bootstrapBCA confidenceLevel sample estimators resamples
  | confidenceLevel > 0 && confidenceLevel < 1
      = runPar $ parMap (uncurry e) (zip estimators resamples)
  | otherwise = error "Statistics.Resampling.Bootstrap.bootstrapBCA: confidence level outside (0,1) range"
  where
    e est (Resample resample)
      | U.length sample == 1 || isInfinite bias =
          estimate pt pt pt confidenceLevel
      | otherwise =
          estimate pt (resample ! lo) (resample ! hi) confidenceLevel
      where
        pt    = R.estimate est sample
        lo    = max (cumn a1) 0
          where a1 = bias + b1 / (1 - accel * b1)
                b1 = bias + z1
        hi    = min (cumn a2) (ni - 1)
          where a2 = bias + b2 / (1 - accel * b2)
                b2 = bias - z1
        z1    = quantile standard ((1 - confidenceLevel) / 2)
        cumn  = round . (*n) . cumulative standard
        bias  = quantile standard (probN / n)
          where probN = fromIntegral . U.length . U.filter (<pt) $ resample
        ni    = U.length resample
        n     = fromIntegral ni
        accel = sumCubes / (6 * (sumSquares ** 1.5))
          where (sumSquares :< sumCubes) = U.foldl' f (0 :< 0) jack
                f (s :< c) j = s + d2 :< c + d2 * d
                    where d  = jackMean - j
                          d2 = d * d
                jackMean     = mean jack
        jack  = jackknife est sample

-- $references
--
-- * Davison, A.C; Hinkley, D.V. (1997) Bootstrap methods and their
--   application. <http://statwww.epfl.ch/davison/BMA/>


-- This file is auto-generated.  Do not edit directly.
-- |
--   Stability: Experimental
--
--   Generic interface to Blas using unsafe foreign calls.  Refer to the GHC
--   documentation for more information regarding appropriate use of safe and unsafe
--   foreign calls.
--
--   The functions here are named in a similar fashion to the original Blas interface, with
--   the type-dependent letter(s) removed.  Some functions have been merged with
--   others to allow the interface to work on both real and complex numbers.  If you can't a
--   particular function, try looking for its corresponding complex equivalent (e.g.
--   @symv@ is a special case of 'hemv' applied to real numbers).
--
--   Note: although complex versions of @rot@ and @rotg@ exist in many implementations,
--   they are not part of the official Blas standard and therefore not included here.  If
--   you /really/ need them, submit a ticket so we can try to come up with a solution.
--
--   The documentation here is still incomplete.  Consult the
--   <http://netlib.org/blas/#_blas_routines official documentation> for more
--   information.
--
--   Notation:
--
--     * @⋅@ denotes dot product (without any conjugation).
--     * @*@ denotes complex conjugation.
--     * @⊤@ denotes transpose.
--     * @†@ denotes conjugate transpose (Hermitian conjugate).
--
--   Conventions:
--
--     * All scalars are denoted with lowercase Greek letters
--     * All vectors are denoted with lowercase Latin letters and are
--       assumed to be column vectors (unless transposed).
--     * All matrices are denoted with uppercase Latin letters.
--
--   /Since: 0.1.1/
module Blas.Specialized.ComplexDouble.Unsafe (
  -- * Level 1: vector-vector operations
  -- ** Basic operations
    swap
  , scal
  , copy
  , axpy
  , dotu
  , dotc
  -- ** Norm operations
  , nrm2
  , asum
  , iamax
  -- * Level 2: matrix-vector operations
  -- ** Multiplication
  , gemv
  , gbmv
  , hemv
  , hbmv
  , hpmv
  -- ** Triangular operations
  , trmv
  , tbmv
  , tpmv
  , trsv
  , tbsv
  , tpsv
  -- ** Rank updates
  , geru
  , gerc
  , her
  , hpr
  , her2
  , hpr2
  -- * Level 3: matrix-matrix operations
  -- ** Multiplication
  , gemm
  , symm
  , hemm
  -- ** Rank updates
  , syrk
  , herk
  , syr2k
  , her2k
  -- ** Triangular operations
  , trmm
  , trsm
  ) where
import Prelude (Double, Int, IO)
import Foreign (Ptr)
import Data.Complex (Complex((:+)))
import FFI (getReturnValue)
import Blas.Primitive.Types (Order, Transpose, Uplo, Diag, Side)
import qualified Blas.Primitive.Unsafe as C

-- | Swap two vectors:
--
--   > (x, y) ← (y, x)
swap :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
swap = C.zswap

-- | Multiply a vector by a scalar.
--
--   > x ← α x
scal :: Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
scal n (alpha :+ 0) = C.zdscal n alpha
scal n alpha = C.zscal n alpha

-- | Copy a vector into another vector:
--
--   > y ← x
copy :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
copy = C.zcopy

-- | Add a scalar-vector product to a vector.
--
--   > y ← α x + y
axpy :: Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
axpy = C.zaxpy

-- | Calculate the bilinear dot product of two vectors:
--
--   > x ⋅ y ≡ ∑[i] x[i] y[i]
dotu :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO (Complex Double)
dotu a b c d e = getReturnValue (C.zdotu_sub a b c d e)

-- | Calculate the sesquilinear dot product of two vectors.
--
--   > x* ⋅ y ≡ ∑[i] x[i]* y[i]
dotc :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO (Complex Double)
dotc a b c d e = getReturnValue (C.zdotc_sub a b c d e)

-- | Calculate the Euclidean (L²) norm of a vector:
--
--   > ‖x‖₂ ≡ √(∑[i] x[i]²)
nrm2 :: Int
     -> Ptr (Complex Double)
     -> Int
     -> IO Double
nrm2 = C.dznrm2

-- | Calculate the Manhattan (L¹) norm, equal to the sum of the magnitudes of the elements:
--
--   > ‖x‖₁ = ∑[i] |x[i]|
asum :: Int
     -> Ptr (Complex Double)
     -> Int
     -> IO Double
asum = C.dzasum

-- | Calculate the index of the element with the maximum magnitude (absolute value).
iamax :: Int
      -> Ptr (Complex Double)
      -> Int
      -> IO Int
iamax = C.izamax

-- | Perform a general matrix-vector update.
--
--   > y ← α T(A) x + β y
gemv :: Order
     -> Transpose
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
gemv = C.zgemv

-- | Perform a general banded matrix-vector update.
--
--   > y ← α T(A) x + β y
gbmv :: Order
     -> Transpose
     -> Int
     -> Int
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
gbmv = C.zgbmv

-- | Perform a hermitian matrix-vector update.
--
--   > y ← α A x + β y
hemv :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
hemv = C.zhemv

-- | Perform a hermitian banded matrix-vector update.
--
--   > y ← α A x + β y
hbmv :: Order
     -> Uplo
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
hbmv = C.zhbmv

-- | Perform a hermitian packed matrix-vector update.
--
--   > y ← α A x + β y
hpmv :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
hpmv = C.zhpmv

-- | Multiply a triangular matrix by a vector.
--
--   > x ← T(A) x
trmv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trmv = C.ztrmv

-- | Multiply a triangular banded matrix by a vector.
--
--   > x ← T(A) x
tbmv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tbmv = C.ztbmv

-- | Multiply a triangular packed matrix by a vector.
--
--   > x ← T(A) x
tpmv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tpmv = C.ztpmv

-- | Multiply an inverse triangular matrix by a vector.
--
--   > x ← T(A⁻¹) x
trsv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trsv = C.ztrsv

-- | Multiply an inverse triangular banded matrix by a vector.
--
--   > x ← T(A⁻¹) x
tbsv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tbsv = C.ztbsv

-- | Multiply an inverse triangular packed matrix by a vector.
--
--   > x ← T(A⁻¹) x
tpsv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tpsv = C.ztpsv

-- | Perform an unconjugated rank-1 update of a general matrix.
--
--   > A ← α x y⊤ + A
geru :: Order
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
geru = C.zgeru

-- | Perform a conjugated rank-1 update of a general matrix.
--
--   > A ← α x y† + A
gerc :: Order
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
gerc = C.zgerc

-- | Perform a rank-1 update of a Hermitian matrix.
--
--   > A ← α x y† + A
her :: Order
    -> Uplo
    -> Int
    -> Double
    -> Ptr (Complex Double)
    -> Int
    -> Ptr (Complex Double)
    -> Int
    -> IO ()
her = C.zher

-- | Perform a rank-1 update of a Hermitian packed matrix.
--
--   > A ← α x y† + A
hpr :: Order
    -> Uplo
    -> Int
    -> Double
    -> Ptr (Complex Double)
    -> Int
    -> Ptr (Complex Double)
    -> IO ()
hpr = C.zhpr

-- | Perform a rank-2 update of a Hermitian matrix.
--
--   > A ← α x y† + y (α x)† + A
her2 :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
her2 = C.zher2

-- | Perform a rank-2 update of a Hermitian packed matrix.
--
--   > A ← α x y† + y (α x)† + A
hpr2 :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> IO ()
hpr2 = C.zhpr2

-- | Perform a general matrix-matrix update.
--
--   > C ← α T(A) U(B) + β C
gemm :: Order -- ^ Layout of all the matrices.
     -> Transpose -- ^ The operation @T@ to be applied to @A@.
     -> Transpose -- ^ The operation @U@ to be applied to @B@.
     -> Int -- ^ Number of rows of @T(A)@ and @C@.
     -> Int -- ^ Number of columns of @U(B)@ and @C@.
     -> Int -- ^ Number of columns of @T(A)@ and number of rows of @U(B)@.
     -> Complex Double -- ^ Scaling factor @α@ of the product.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @A@.
     -> Int -- ^ Stride of the major dimension of @A@.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @B@.
     -> Int -- ^ Stride of the major dimension of @B@.
     -> Complex Double -- ^ Scaling factor @β@ of the original @C@.
     -> Ptr (Complex Double) -- ^ Pointer to a mutable matrix @C@.
     -> Int -- ^ Stride of the major dimension of @C@.
     -> IO ()
gemm = C.zgemm

-- | Perform a symmetric matrix-matrix update.
--
--   > C ← α A B + β C    or    C ← α B A + β C
--
--   where @A@ is symmetric.  The matrix @A@ must be in an unpacked format, although the
--   routine will only access half of it as specified by the @'Uplo'@ argument.
symm :: Order -- ^ Layout of all the matrices.
     -> Side -- ^ Side that @A@ appears in the product.
     -> Uplo -- ^ The part of @A@ that is used.
     -> Int -- ^ Number of rows of @C@.
     -> Int -- ^ Number of columns of @C@.
     -> Complex Double -- ^ Scaling factor @α@ of the product.
     -> Ptr (Complex Double) -- ^ Pointer to a symmetric matrix @A@.
     -> Int -- ^ Stride of the major dimension of @A@.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @B@.
     -> Int -- ^ Stride of the major dimension of @B@.
     -> Complex Double -- ^ Scaling factor @α@ of the original @C@.
     -> Ptr (Complex Double) -- ^ Pointer to a mutable matrix @C@.
     -> Int -- ^ Stride of the major dimension of @C@.
     -> IO ()
symm = C.zsymm

-- | Perform a Hermitian matrix-matrix update.
--
--   > C ← α A B + β C    or    C ← α B A + β C
--
--   where @A@ is Hermitian.  The matrix @A@ must be in an unpacked format, although the
--   routine will only access half of it as specified by the @'Uplo'@ argument.
hemm :: Order -- ^ Layout of all the matrices.
     -> Side -- ^ Side that @A@ appears in the product.
     -> Uplo -- ^ The part of @A@ that is used.
     -> Int -- ^ Number of rows of @C@.
     -> Int -- ^ Number of columns of @C@.
     -> Complex Double -- ^ Scaling factor @α@ of the product.
     -> Ptr (Complex Double) -- ^ Pointer to a Hermitian matrix @A@.
     -> Int -- ^ Stride of the major dimension of @A@.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @B@.
     -> Int -- ^ Stride of the major dimension of @B@.
     -> Complex Double -- ^ Scaling factor @α@ of the original @C@.
     -> Ptr (Complex Double) -- ^ Pointer to a mutable matrix @C@.
     -> Int -- ^ Stride of the major dimension of @C@.
     -> IO ()
hemm = C.zhemm

-- | Perform a symmetric rank-k update.
--
--   > C ← α A A⊤ + β C    or    C ← α A⊤ A + β C
syrk :: Order
     -> Uplo
     -> Transpose
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
syrk = C.zsyrk

-- | Perform a Hermitian rank-k update.
--
--   > C ← α A A† + β C    or    C ← α A† A + β C
herk :: Order
     -> Uplo
     -> Transpose
     -> Int
     -> Int
     -> Double
     -> Ptr (Complex Double)
     -> Int
     -> Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
herk = C.zherk

-- | Perform a symmetric rank-2k update.
--
--   > C ← α A B⊤ + α* B A⊤ + β C    or    C ← α A⊤ B + α* B⊤ A + β C
syr2k :: Order
      -> Uplo
      -> Transpose
      -> Int
      -> Int
      -> Complex Double
      -> Ptr (Complex Double)
      -> Int
      -> Ptr (Complex Double)
      -> Int
      -> Complex Double
      -> Ptr (Complex Double)
      -> Int
      -> IO ()
syr2k = C.zsyr2k

-- | Perform a Hermitian rank-2k update.
--
--   > C ← α A B† + α* B A† + β C    or    C ← α A† B + α* B† A + β C
her2k :: Order
      -> Uplo
      -> Transpose
      -> Int
      -> Int
      -> Complex Double
      -> Ptr (Complex Double)
      -> Int
      -> Ptr (Complex Double)
      -> Int
      -> Double
      -> Ptr (Complex Double)
      -> Int
      -> IO ()
her2k = C.zher2k

-- | Perform a triangular matrix-matrix multiplication.
--
--   > B ← α T(A) B    or    B ← α B T(A)
--
--   where @A@ is triangular.
trmm :: Order
     -> Side
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trmm = C.ztrmm

-- | Perform an inverse triangular matrix-matrix multiplication.
--
--   > B ← α T(A⁻¹) B    or    B ← α B T(A⁻¹)
--
--   where @A@ is triangular.
trsm :: Order
     -> Side
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trsm = C.ztrsm