{-# 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 BangPatterns           #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilyDependencies #-}

module Backprop.Learn.Model (
    module M, Backprop(..)
  , runLearn_, runLearnStoch_, runLearnStateless_, runLearnStochStateless_
  , gradLearn, gradLearnStoch
  -- * Work with parameters
  , initParam, initParamMaybe, initParamNormal, initParamNormalMaybe
  , encodeParam, decodeParam, decodeParamOrFail, saveParam, loadParam, loadParamOrFail
  -- * Iterated runners
  , iterateLearn, iterateLearnM, iterateLearnStoch
  , scanLearn, scanLearnStoch
  -- * No final state
  , iterateLearn_, iterateLearnM_, iterateLearnStoch_
  , scanLearn_, scanLearnStoch_
  -- * "Prime" runners
  , primeLearn, primeLearnStoch, selfPrime, selfPrimeM
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Class       as M
import           Backprop.Learn.Model.Combinator  as M
import           Backprop.Learn.Model.Function    as M
import           Backprop.Learn.Model.Neural      as M
import           Backprop.Learn.Model.Neural.LSTM as M
import           Backprop.Learn.Model.Parameter   as M
import           Backprop.Learn.Model.Regression  as M
import           Backprop.Learn.Model.State       as M
import           Backprop.Learn.Model.Stochastic  as M
import           Control.Monad.Primitive
import           Control.Monad.ST
import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Foldable
import           Data.Functor.Identity
import           Data.Type.Mayb
import           Data.Word
import           Numeric.Backprop
import           Statistics.Distribution
import qualified Data.Binary                      as Bi
import qualified Data.ByteString.Lazy             as BSL
import qualified Data.Vector.Unboxed              as VU
import qualified System.Random.MWC                as MWC

-- TODO: this can be more efficient by breaking out into separate functions
runLearn_
    :: (Learn a b l, MaybeC Backprop (LStateMaybe l), Backprop b)
    => l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> (b, LState_ I l)
runLearn_ l mp x ms = case mp of
    N_ -> case ms of
      N_       -> (evalBP (runLearnStateless l N_) x, N_)
      J_ (I s) -> second (J_ . I)
                . evalBP2 (\x' s' -> uncurry T2
                                       . second fromJ_
                                       $ runLearn l N_ x' (J_ s')
                          ) x
                $ s
    J_ (I p) -> case ms of
      N_       -> (evalBP2 (runLearnStateless l . J_) p x, N_)
      J_ (I s) -> second (J_ . I)
                . evalBPN (\(p' :< x' :< s' :< Ø) ->
                                uncurry T2
                              . second fromJ_
                              $ runLearn l (J_ p') x' (J_ s')
                          )
                $ (p ::< x ::< s ::< Ø)

runLearnStoch_
    :: (Learn a b l, MaybeC Backprop (LStateMaybe l), Backprop b, PrimMonad m)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> m (b, LState_ I l)
runLearnStoch_ l g mp x ms = do
    -- TODO: is this the best way to handle this?
    seed <- MWC.uniformVector @_ @Word32 @VU.Vector g 2
    pure $ case mp of
      N_ -> case ms of
        N_       -> (, N_) $ evalBP (\x' -> runST $ do
            g' <- MWC.initialize seed
            runLearnStochStateless l g' N_ x'
          ) x
        J_ (I s) -> second (J_ . I) $ evalBP2 (\x' s' -> runST $ do
            g' <- MWC.initialize seed
            uncurry T2 . second fromJ_
              <$> runLearnStoch l g' N_ x' (J_ s')
          ) x s
      J_ (I p) -> case ms of
        N_       -> (, N_) $ evalBP2 (\p' x' -> runST $ do
            g' <- MWC.initialize seed
            runLearnStochStateless l g' (J_ p') x'
          ) p x
        J_ (I s) -> second (J_ . I) $ evalBPN (\(p' :< x' :< s' :< Ø) -> runST $ do
            g' <- MWC.initialize seed
            uncurry T2 . second fromJ_
              <$> runLearnStoch l g' (J_ p') x' (J_ s')
          ) (p ::< x ::< s ::< Ø)

runLearnStateless_
    :: (Learn a b l, NoState l)
    => l
    -> LParam_ I l
    -> a
    -> b
runLearnStateless_ l = \case
    N_       -> evalBP  (runLearnStateless l N_  )
    J_ (I p) -> evalBP2 (runLearnStateless l . J_) p

runLearnStochStateless_
    :: (Learn a b l, NoState l, PrimMonad m)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> a
    -> m b
runLearnStochStateless_ l g mp x = do
    seed <- MWC.uniformVector @_ @Word32 @VU.Vector g 2
    pure $ case mp of
      N_       -> evalBP  (\x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless l g' N_ x'
        ) x
      J_ (I p) -> evalBP2 (\p' x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless l g' (J_ p') x'
        ) p x

gradLearn
    :: (Learn a b l, NoState l, Backprop a, Backprop b, MaybeC Backprop (LParamMaybe l))
    => l
    -> LParam_ I l
    -> a
    -> (LParam_ I l, a)
gradLearn l = \case
    N_       ->       (N_,)    . gradBP  (runLearnStateless l   N_)
    J_ (I p) -> first (J_ . I) . gradBP2 (runLearnStateless l . J_) p

gradLearnStoch
    :: (Learn a b l, NoState l, Backprop a, Backprop b, MaybeC Backprop (LParamMaybe l), PrimMonad m)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> a
    -> m (LParam_ I l, a)
gradLearnStoch l g mp x = do
    seed <- MWC.uniformVector @_ @Word32 @VU.Vector g 2
    pure $ case mp of
      N_       ->       (N_,)    $ gradBP  (\x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless l g' N_ x'
        ) x
      J_ (I p) -> first (J_ . I) $ gradBP2 (\p' x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless l g' (J_ p') x'
        ) p x

iterateLearn
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l))
    => (b -> a)         -- ^ loop
    -> Int              -- ^ num times
    -> l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> ([b], LState_ I l)
iterateLearn f n l p x = runIdentity . iterateLearnM (Identity . f) n l p x

iterateLearn_
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l))
    => (b -> a)         -- ^ loop
    -> l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> [b]
iterateLearn_ f l p = go
  where
    go !x !s = y : go (f y) s'
      where
        (y, s') = runLearn_ l p x s

selfPrime
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l))
    => (b -> a)         -- ^ loop
    -> l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> [LState_ I l]
selfPrime f l p = go
  where
    go !x !s = s' : go (f y) s'
      where
        (y, s') = runLearn_ l p x s

iterateLearnM
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), Monad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> m ([b], LState_ I l)
iterateLearnM f n l p = go 0
  where
    go !i !x !s
      | i <= n    = do
          let (y, s') = runLearn_ l p x s
          (ys, s'') <- flip (go (i + 1)) s' =<< f y
          pure (y : ys, s'')
      | otherwise = pure ([], s)

iterateLearnM_
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), Monad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> m [b]
iterateLearnM_ f n l p x = fmap fst . iterateLearnM f n l p x

selfPrimeM
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), Monad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> l
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> m (LState_ I l)
selfPrimeM f n l p x = fmap snd . iterateLearnM f n l p x

iterateLearnStoch
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), PrimMonad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> m ([b], LState_ I l)
iterateLearnStoch f n l g p = go 0
  where
    go !i !x !s
      | i <= n    = do
          (y , s' ) <- runLearnStoch_ l g p x s
          (ys, s'') <- flip (go (i + 1)) s' =<< f y
          pure (y : ys, s'')
      | otherwise = pure ([], s)

iterateLearnStoch_
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), PrimMonad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> a
    -> LState_ I l
    -> m [b]
iterateLearnStoch_ f n l g p x = fmap fst . iterateLearnStoch f n l g p x

scanLearn
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), Traversable t)
    => l
    -> LParam_ I l
    -> t a
    -> LState_ I l
    -> (t b, LState_ I l)
scanLearn l p = runState . traverse (state . runLearn_ l p)

scanLearn_
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), Traversable t)
    => l
    -> LParam_ I l
    -> t a
    -> LState_ I l
    -> t b
scanLearn_ l p xs = fst . scanLearn l p xs

primeLearn
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), Foldable t)
    => l
    -> LParam_ I l
    -> t a
    -> LState_ I l
    -> LState_ I l
primeLearn l p = execState . traverse_ (state . runLearn_ l p)

scanLearnStoch
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), PrimMonad m, Traversable t)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> t a
    -> LState_ I l
    -> m (t b, LState_ I l)
scanLearnStoch l g p = runStateT . traverse (StateT . runLearnStoch_ l g p)

scanLearnStoch_
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), PrimMonad m, Traversable t)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> t a
    -> LState_ I l
    -> m (t b)
scanLearnStoch_ l g p xs = fmap fst . scanLearnStoch l g p xs

primeLearnStoch
    :: (Learn a b l, Backprop b, MaybeC Backprop (LStateMaybe l), PrimMonad m, Foldable t)
    => l
    -> MWC.Gen (PrimState m)
    -> LParam_ I l
    -> t a
    -> LState_ I l
    -> m (LState_ I l)
primeLearnStoch l g p = execStateT . traverse_ (StateT . runLearnStoch_ l g p)

initParamMaybe
    :: forall l m d proxy.
     ( MaybeC Initialize (LParamMaybe l)
     , ContGen d
     , PrimMonad m
     , KnownMayb (LParamMaybe l)
     )
    => proxy l                            -- ^ ignored
    -> d
    -> MWC.Gen (PrimState m)
    -> LParam_ m l
initParamMaybe _ d g = case knownMayb @(LParamMaybe l) of
                    N_            -> N_
                    J_ (_ :: P p) -> J_ $ initialize @p d g

initParam
    :: forall l m d proxy.
     ( Initialize (LParam l)
     , ContGen d
     , PrimMonad m
     )
    => proxy l                            -- ^ ignored
    -> d
    -> MWC.Gen (PrimState m)
    -> m (LParam l)
initParam _ = initialize

initParamNormalMaybe
    :: forall l m proxy.
     ( MaybeC Initialize (LParamMaybe l)
     , PrimMonad m
     , KnownMayb (LParamMaybe l)
     )
    => proxy l                            -- ^ ignored
    -> Double
    -> MWC.Gen (PrimState m)
    -> LParam_ m l
initParamNormalMaybe _ d g = case knownMayb @(LParamMaybe l) of
    N_            -> N_
    J_ (_ :: P p) -> J_ $ initializeNormal @p d g

initParamNormal
    :: forall l m proxy.
     ( Initialize (LParam l)
     , PrimMonad m
     )
    => proxy l                            -- ^ ignored
    -> Double
    -> MWC.Gen (PrimState m)
    -> m (LParam l)
initParamNormal _ = initializeNormal

encodeParam
    :: Bi.Binary (LParam l)
    => proxy l                              -- ^ ignored
    -> LParam l
    -> BSL.ByteString
encodeParam _ = Bi.encode

decodeParam
    :: Bi.Binary (LParam l)
    => proxy l                              -- ^ ignored
    -> BSL.ByteString
    -> LParam l
decodeParam _ = Bi.decode

decodeParamOrFail
    :: Bi.Binary (LParam l)
    => proxy l                              -- ^ ignored
    -> BSL.ByteString
    -> Either String (LParam l)
decodeParamOrFail _ = bimap thrd thrd . Bi.decodeOrFail

saveParam
    :: Bi.Binary (LParam l)
    => proxy l
    -> FilePath
    -> LParam l
    -> IO ()
saveParam p fp = BSL.writeFile fp . encodeParam p

loadParam
    :: Bi.Binary (LParam l)
    => proxy l
    -> FilePath
    -> IO (LParam l)
loadParam p fp = decodeParam p <$> BSL.readFile fp

loadParamOrFail
    :: Bi.Binary (LParam l)
    => proxy l
    -> FilePath
    -> IO (Either String (LParam l))
loadParamOrFail p fp = decodeParamOrFail p <$> BSL.readFile fp

thrd :: (a,b,c) -> c
thrd (_,_,z) = z


{-# LANGUAGE ForeignFunctionInterface #-}
{-|
Module      : Grenade.Layers.Internal.Activations
Description : Fast activation functions that call efficient C function
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Internal.Activations
  ( relu
  , relu1d
  , leakyRelu
  , leakyRelu1d
  ) where

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

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

import           Grenade.Types

-- | Apply ReLU on a matrix
relu :: Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
relu channels rows cols = leakyRelu channels rows cols 0
{-# INLINE relu #-}

-- | Apply ReLU on a vector
relu1d :: Int -> Vector RealNum -> Vector RealNum
relu1d size = leakyRelu1d size 0
{-# INLINE relu1d #-}

-- | Apply LeakyReLU on a matrix
leakyRelu :: Int -> Int -> Int -> RealNum -> Matrix RealNum -> Matrix RealNum
leakyRelu channels rows cols alpha m
  = let vec = flatten m
        out = leakyRelu1d (channels * rows * cols) alpha vec
    in  U.matrixFromVector U.RowMajor (rows * channels) cols out
{-# INLINE leakyRelu #-}

-- | Apply LeakyReLU on a vector
leakyRelu1d :: Int -> RealNum -> Vector RealNum -> Vector RealNum
leakyRelu1d size alpha vec
  = unsafePerformIO $ do
      outPtr        <- mallocForeignPtrArray size
      let (inPtr, _) = U.unsafeToForeignPtr0 vec

      withForeignPtr inPtr $ \inPtr' ->
        withForeignPtr outPtr $ \outPtr' ->
          leaky_relu_forward inPtr' size alpha outPtr'

      return $ U.unsafeFromForeignPtr0 outPtr size
{-# INLINE leakyRelu1d #-}

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


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

module FinancialTimeseries.Report.Standard where

import Data.Time (UTCTime)

import Data.Bifunctor (bimap)
import Data.Distributive (Distributive, distribute)

import qualified Data.Vector as Vec
import Data.Vector (Vector)

import qualified System.Random as R

import qualified Text.Blaze.Html5 as H5

import qualified Graphics.Rendering.Chart.Easy as E

import qualified Statistics.Sample.Histogram as Histo

import qualified FinancialTimeseries.Algorithm.MonteCarlo as AMC
import FinancialTimeseries.Algorithm.Evaluate (Profit, long, evaluate, evaluateFraction)
import FinancialTimeseries.Render.Chart (chartWithEquity)
import qualified FinancialTimeseries.Render.HtmlReader as HtmlReader -- (Config, runHtmlReader)
import FinancialTimeseries.Render.Render (display)
import FinancialTimeseries.Render.Statement (statement, currentTime)
import FinancialTimeseries.Statistics.Statistics (yield, tradeStatistics, stats2cdfChart, stats2list)
import FinancialTimeseries.Test (check_timeseries_prop)
import FinancialTimeseries.Type.Fraction (Fraction(..))
import FinancialTimeseries.Type.Histogram (Histogram(..))
import FinancialTimeseries.Type.Labeled (Labeled(..))
import FinancialTimeseries.Type.MonteCarlo (Broom(..))
import FinancialTimeseries.Type.Strategy (Strategy(..))
import FinancialTimeseries.Type.Timeseries (TimeseriesRaw, first, slice)
import qualified FinancialTimeseries.Type.Timeseries as TS
import FinancialTimeseries.Type.Types (StripPrice, stripPrice, Equity(..), partitionInvested)

import FinancialTimeseries.Util.DistributivePair (distributePair)
import FinancialTimeseries.Util.Pretty (Pretty, pretty)

data Config params gen price a = Config {
  now :: UTCTime
  , reportConfig :: HtmlReader.Config
  , monteCarloConfig :: AMC.Config gen a
  , strategy :: Strategy params price a
  }

report ::
  (Profit price, Distributive price, StripPrice price, Pretty params, R.RandomGen gen, Num a, Fractional a, Real a, E.PlotValue a, Show a, Pretty a, TS.Length (TimeseriesRaw price a)) =>
  Config params gen price a -> TimeseriesRaw price a -> H5.Html
report cfg ts =
  let t = (freeStrategy (strategy cfg)) (parameters (strategy cfg)) ts
      lg = long (partitionInvested (slice t))

      mteCrlo = AMC.mc (monteCarloConfig cfg) lg

      yields =
        let k = fmap (Vec.fromList . map (yield . Vec.map snd))
        in fmap (fmap (bimap k k)) lg

      tradeYields =
        let  j zs = (fmap stats2cdfChart zs, fmap stats2list zs)
             k = j . fmap ((:[]) . Labeled "Trade Yields" . tradeStatistics)
        in fmap (fmap (snd . fmap k)) yields

      tsCharts =
        let convert = Equity . snd . stripPrice
            k = fmap (flip chartWithEquity [t]) . distribute
        in fmap (bimap k k . distributePair) (evaluate (convert (first ts)) lg)

      histogram =
        let histo zs =
              let n = round (fromIntegral (Vec.length zs) / fromIntegral (5 :: Integer) :: Double)
              in Histogram (Histo.histogram n (Vec.map realToFrac zs) :: (Vector Double, Vector Int))
        in fmap (fmap (fmap histo . snd)) yields

      ms = map (\f -> mteCrlo f (evaluateFraction f)) (AMC.fractions (monteCarloConfig cfg))
      ys = fmap (fmap (fmap snd)) (AMC.timeseriesYields ms)
      rdds = fmap (fmap (fmap snd)) (AMC.relativeDrawdowns ms)

      broom = fmap (fmap (fmap (bimap (fmap Broom) (fmap Broom)))) (mteCrlo (Fraction 1.0) evaluate)

      html = HtmlReader.runHtmlReader (reportConfig cfg) $ mconcat $
        currentTime (now cfg)
        : statement ("Timeseries is well formed: " ++ show (check_timeseries_prop t))
        : statement ("Parameters: " ++ pretty (parameters (strategy cfg)))
        : display tsCharts
        : display tradeYields
        : display histogram
        : display broom
        : display ys
        : display rdds
        : []

  in html


{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}

module Plotting where

import Data.List (sort)
import qualified Data.Vector as Vector
import Graphics.Rendering.Chart.Easy
import Statistics.Sample (meanVariance)

{-# ANN module "HLint: ignore Reduce duplication" #-}

errBars :: Double -> Vector.Vector Double -> (Double, Double)
errBars k xs = (mean, k * stdDev)
  where
    (mean, var) = meanVariance xs
    stdDev = sqrt var

anytimeEval :: ([(a, Double)] -> Double) -> [(a, Double)] -> [Int] -> [Double]
anytimeEval f xs ns = map eval (filter (<= length xs) ns)
  where
    eval n = f (take n xs)

anytimePlot :: String -> String -> [Int] -> [(String, [Double])] -> EC (Layout LogValue LogValue) ()
anytimePlot x_title y_title ns inputs = do
  let xLabelShow = map (show . ceiling)
  let yLabelShow = map show
  let generatePlot (algName, ys) = plot $ line algName [zip (map fromIntegral ns) (map LogValue ys)]
  mapM_ generatePlot inputs
  layout_x_axis . laxis_generate .= autoScaledLogAxis (loga_labelf .~ xLabelShow $ def)
  layout_y_axis . laxis_generate .= autoScaledLogAxis (loga_labelf .~ yLabelShow $ def)
  layout_x_axis . laxis_title .= x_title
  layout_y_axis . laxis_title .= y_title

oneShotPlot :: String -> String -> [(String, [(Double, (Double, Double))])] -> EC (Layout LogValue LogValue) ()
oneShotPlot x_title y_title inputs = do
  let xLabelShow = map (show . ceiling)
  let yLabelShow = map show
  let processPoint (x, (y, dy)) = symErrPoint (LogValue x) (LogValue y) 0 (LogValue dy)
  let generatePlot (algName, ps) = do
        plot $ points algName $ map (\(x, (y, _)) -> (LogValue x, LogValue y)) ps
        plot $ return (def & plot_errbars_values .~ map processPoint ps)
  mapM_ generatePlot inputs
  layout_x_axis . laxis_generate .= autoScaledLogAxis (loga_labelf .~ xLabelShow $ def)
  layout_y_axis . laxis_generate .= autoScaledLogAxis (loga_labelf .~ yLabelShow $ def)
  layout_x_axis . laxis_title .= x_title
  layout_y_axis . laxis_title .= y_title

errorbarPlot ::
  String ->
  String ->
  [(String, [(Double, [Double])])] ->
  EC (Layout LogValue LogValue) ()
errorbarPlot x_title y_title inputs = do
  let xLabelShow = map (show . ceiling)
  let yLabelShow = map show
  let makeBars (x, ys) = ErrPoint (ErrValue x' x' x') (ErrValue ymin y ymax)
        where
          x' = LogValue x
          y = LogValue $ sum ys / fromIntegral (length ys)
          sorted = sort ys
          n = length ys
          ymin = LogValue $ sorted !! (n `div` 4)
          ymax = LogValue $ sorted !! (n * 3 `div` 4)
  let generatePlot (algName, ps) = do
        plot $ liftEC $ do
          -- ensure error bars are the same color as the points
          (color : _) <- liftCState $ use colors
          plot_errbars_values .= map makeBars ps
          plot_errbars_line_style . line_color .= color
        plot $ points algName $
          map (\(x, ys) -> (LogValue x, LogValue (sum ys / fromIntegral (length ys)))) ps
  mapM_ generatePlot inputs
  layout_x_axis . laxis_generate .= autoScaledLogAxis (loga_labelf .~ xLabelShow $ def)
  layout_y_axis . laxis_generate .= autoScaledLogAxis (loga_labelf .~ yLabelShow $ def)
  layout_x_axis . laxis_title .= x_title
  layout_y_axis . laxis_title .= y_title


{-# 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


{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module Marvin.Test.TestUtils (
  isAround
  , nestedFromList
  , trainMatrix
  , equals
  , (+-)
) where

import Marvin.Test.Metric
import Marvin.API as Marvin

import Test.QuickCheck hiding (vector)
import qualified Test.QuickCheck as QC (vector)
import Numeric.LinearAlgebra.Data as LA
import Foreign.Storable

import qualified Data.Vector as Vec

nestedFromList = Vec.fromList . Prelude.map Vec.fromList

genVectorOfSize :: (Storable a, Arbitrary a) => Int -> Gen (Vector a)
genVectorOfSize n = do
  xs <- vectorOf n arbitrary
  return $ LA.fromList xs

instance (Storable a, Arbitrary a) => Arbitrary (Vector a) where
  arbitrary = sized genVectorOfSize

instance (Arbitrary a) => Arbitrary (Vec.Vector a) where
  arbitrary = Vec.fromList <$> arbitrary

instance Arbitrary NumericColumn where
  arbitrary = do
      t <- arbitrary
      h <- arbitrary
      let Right col = Marvin.fromList (h:t)
      return col

instance Arbitrary Error where
  arbitrary = do
    msg <- arbitrary
    return $ Marvin.Failure msg

isAround :: Metric a => a -> a -> Bool
isAround a x = dist a x <= threshold

infix 7 `equals`
equals :: Metric a => a -> a -> FloatPrecision -> Bool
equals x y (P e) = dist x y < e

newtype FloatPrecision = P Double

infix 8 +-
(+-) :: Double -> FloatPrecision
(+-) = P

threshold = 0.00001

trainMatrix = [
  [2,6,2,3,2,2,1,3,2,2,2,15.0],
  [1,6,1,3,3,1,1,2,2,4,2,14.5],
  [2,6,3,1,4,2,3,3,1,3,1,14.0],
  [2,6,1,4,4,2,2,3,2,4,2,16.0],
  [1,1,2,2,2,1,1,4,1,3,1,14.0],
  [1,3,2,1,1,1,2,2,1,2,1,16.5],
  [1,4,1,1,3,1,3,1,1,3,1,14.0],
  [1,3,2,3,2,1,2,3,1,3,1,16.5],
  [2,4,1,2,1,2,3,2,1,4,1,15.0],
  [2,4,1,1,2,2,3,3,1,3,1,16.0],
  [2,2,1,2,3,2,1,2,1,3,1,16.5],
  [2,5,2,2,1,2,3,1,2,1,2,15.0],
  [2,5,2,2,3,2,3,1,2,1,2,15.5],
  [2,6,1,4,3,2,2,3,2,4,2,15.5],
  [2,3,2,3,1,2,1,2,1,3,1,15.0],
  [1,1,2,4,3,1,1,1,1,2,1,15.0],
  [1,1,3,2,2,1,2,2,2,2,2,14.5],
  [1,2,1,1,2,1,2,1,1,4,1,16.0],
  [2,1,2,4,4,2,2,2,1,1,1,15.5],
  [2,2,2,2,3,2,3,4,1,2,1,14.5],
  [1,5,2,3,3,1,3,2,1,1,1,13.5],
  [2,5,3,4,4,2,2,4,2,3,2,16.0],
  [1,2,2,3,4,1,3,3,1,2,1,15.0],
  [2,3,1,2,2,2,2,2,1,1,1,14.5],
  [2,2,1,3,1,2,1,1,2,4,2,15.0],
  [1,1,2,3,4,1,1,2,1,4,1,14.5],
  [2,1,3,4,2,2,1,4,1,1,1,14.5],
  [1,3,3,4,3,1,3,4,2,4,2,15.5],
  [2,2,3,1,3,2,2,1,2,2,2,15.5],
  [1,5,1,2,4,1,2,1,2,4,2,15.5],
  [2,2,2,3,1,2,3,2,2,4,2,14.0],
  [1,2,3,4,2,1,1,4,1,3,1,17.0],
  [1,1,3,4,2,1,2,3,2,3,2,15.0],
  [1,2,1,2,1,1,2,3,1,3,1,16.0],
  [2,3,2,3,3,2,1,2,2,3,2,15.0],
  [2,5,1,1,4,2,1,4,2,1,2,16.0],
  [1,4,3,3,3,1,2,3,1,3,1,15.5],
  [1,2,3,3,3,1,1,3,2,1,2,16.5],
  [2,6,1,3,1,2,2,2,1,1,1,15.5],
  [1,6,2,1,1,1,2,3,2,3,2,15.0],
  [2,6,1,2,3,2,2,1,1,2,1,15.0],
  [2,2,1,4,2,2,1,4,2,1,2,15.0],
  [2,5,1,4,4,2,1,3,2,2,2,14.5],
  [1,5,1,3,1,1,2,2,2,2,2,14.5],
  [2,4,1,4,1,2,3,4,1,2,1,16.0],
  [1,6,1,4,2,1,1,4,2,3,2,14.5],
  [1,6,1,2,4,1,1,1,2,2,2,13.5],
  [2,6,1,1,2,2,2,4,1,3,1,15.0],
  [1,4,1,1,1,1,3,1,2,3,2,14.5],
  [2,1,1,4,3,2,3,3,2,3,2,15.0],
  [1,3,2,1,2,1,2,2,2,2,2,14.5],
  [2,5,1,2,2,2,1,1,1,4,1,15.5],
  [1,3,3,1,3,1,3,2,1,1,1,15.0],
  [1,2,3,4,3,1,1,4,2,3,2,16.5],
  [1,1,2,3,1,1,1,2,1,4,1,15.0],
  [1,5,3,1,2,1,1,1,1,2,1,14.0],
  [1,6,3,2,3,1,3,2,1,1,1,15.0],
  [2,5,3,1,3,2,2,3,1,4,1,14.0],
  [1,2,2,2,4,1,3,1,1,3,1,16.5],
  [2,1,2,3,4,2,2,3,1,4,1,16.0],
  [2,6,2,3,3,2,1,3,1,2,1,14.0],
  [1,3,3,1,1,1,3,2,2,1,2,15.5],
  [2,3,2,4,2,2,1,3,2,2,2,14.0],
  [1,5,2,2,1,1,3,4,1,4,1,14.5],
  [2,6,2,2,1,2,1,1,1,4,1,15.0],
  [1,2,3,1,4,1,1,2,2,4,2,16.0],
  [1,4,3,1,3,1,2,2,1,2,1,14.5],
  [2,5,3,4,3,2,2,4,2,3,2,15.0],
  [1,6,2,3,1,1,2,2,2,2,2,13.5],
  [2,4,1,3,4,2,3,1,1,1,1,15.0],
  [2,4,3,4,1,2,1,2,2,3,2,14.0],
  [1,4,3,4,1,1,2,4,2,1,2,15.5],
  [1,3,2,2,3,1,2,1,1,4,1,16.0],
  [1,6,3,3,2,1,3,4,1,4,1,15.0],
  [1,3,3,2,2,1,3,1,1,3,1,15.5],
  [1,6,2,4,3,1,2,1,2,4,2,15.5],
  [2,2,3,1,1,2,2,1,2,2,2,14.5],
  [2,2,2,3,2,2,3,2,1,4,1,14.5],
  [2,2,1,2,2,2,1,2,1,3,1,15.0],
  [2,5,2,3,2,2,3,3,1,3,1,15.0]]

testMatrix = [
  [1,1,1,4,4,1,3,4,1,4,1,16.0],
  [1,4,3,3,1,1,2,3,1,3,1,15.5],
  [2,2,1,4,4,2,1,4,2,1,2,15.0],
  [1,5,2,4,3,1,3,1,1,3,1,15.0],
  [1,4,2,1,2,1,1,2,2,4,2,14.5],
  [1,1,2,1,4,1,1,3,2,1,2,15.0],
  [1,1,3,1,3,1,2,1,2,4,2,15.0],
  [1,4,2,1,1,1,1,2,1,4,1,14.5],
  [2,6,1,4,2,2,2,3,1,4,1,15.0],
  [2,5,1,3,2,2,1,2,2,3,2,15.5],
  [1,6,1,3,1,1,1,2,1,4,1,15.0],
  [2,3,2,1,3,2,1,1,1,4,1,15.5],
  [1,2,2,3,3,1,3,3,2,2,2,14.0],
  [1,4,1,1,2,1,3,1,2,3,2,14.0],
  [1,2,1,2,4,1,2,3,2,3,2,15.0],
  [1,2,2,2,1,1,3,1,2,3,2,16.0],
  [1,5,1,1,1,1,2,3,1,3,1,15.5],
  [1,3,2,4,3,1,2,4,2,1,2,17.0],
  [2,1,3,4,3,2,1,4,2,1,2,15.5],
  [2,5,2,3,1,2,3,3,2,3,2,15.0]]


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.LeakyRelu
Description : Rectifying linear unit layer
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.LeakyRelu (
    LeakyRelu (..)
  , SpecLeakyRelu (..)
  , specLeakyRelu1D
  , specLeakyRelu2D
  , specLeakyRelu3D
  , leakyRelu
  ) where

import           Control.DeepSeq                (NFData (..), force)
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.Types
import           Grenade.Utils.Conversion       (toLayerShape)
import           Grenade.Utils.Vector

-- | A rectifying linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
data LeakyRelu = LeakyRelu
  deriving (Generic, NFData, Show)

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

instance RandomLayer LeakyRelu where
  createRandomWith _ _ = return LeakyRelu

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

instance (KnownNat i) => Layer LeakyRelu ('D1 i) ('D1 i) where
  type Tape LeakyRelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (relu y))
    where
      relu = LAS.dvmap (\a -> if a < 0 then alpha * a else a)
  runForwards _ (S1DV y) = (S1DV y, S1DV (mapVectorInPlace relu y))
    where
      relu = (\a -> if a < 0 then alpha * a else a) -- we do not change the sign, so we can use inPlace here
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (relu' y * dEdy))
    where
      relu' = LAS.dvmap (\a -> if a < 0 then alpha else 1)
  runBackwards _ (S1DV y) (S1DV dEdy) = ((), zipWithVectorInPlaceSnd reluDif y dEdy `seq` S1DV dEdy)
  runBackwards x y dEdy = runBackwards x y (toLayerShape y dEdy)

instance (KnownNat i, KnownNat j) => Layer LeakyRelu ('D2 i j) ('D2 i j) where
  type Tape LeakyRelu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (relu y))
    where
      relu = LAS.dmmap (\a -> if a < 0 then alpha * a else a)
  runForwards _ (S2DV y) = (S2DV y, S2DV (mapVectorInPlace relu y))
    where
      relu = (\a -> if a < 0 then alpha * a else a)
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (relu' y * dEdy))
    where
      relu' = LAS.dmmap (\a -> if a < 0 then alpha else 1)
  runBackwards _ (S2DV y) (S2DV dEdy) = ((), zipWithVectorInPlaceSnd reluDif y dEdy `seq` S2DV dEdy)
  runBackwards x y dEdy = runBackwards x y (toLayerShape y dEdy)

instance (KnownNat i, KnownNat j, KnownNat k) => Layer LeakyRelu ('D3 i j k) ('D3 i j k) where

  type Tape LeakyRelu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (relu y))
    where
      relu = LAS.dmmap (\a -> if a < 0 then alpha * a else a)
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (relu' y * dEdy))
    where
      relu' = LAS.dmmap (\a -> if a < 0 then alpha else 1)

alpha :: RealNum
alpha = 0.02

reluDif :: RealNum -> RealNum -> RealNum
reluDif a g
  | a < 0 = alpha * g
  | otherwise = g

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

instance FromDynamicLayer LeakyRelu where
  fromDynamicLayer inp _ _ = SpecNetLayer $ SpecLeakyRelu (tripleFromSomeShape inp)

instance ToDynamicLayer SpecLeakyRelu where
  toDynamicLayer _ _ (SpecLeakyRelu (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 LeakyRelu (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> return $ SpecLayer LeakyRelu (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _    -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer LeakyRelu (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a elu layer.
specLeakyRelu1D :: Integer -> SpecNet
specLeakyRelu1D i = specLeakyRelu3D (i, 1, 1)

-- | Create a specification for a elu layer.
specLeakyRelu2D :: (Integer, Integer) -> SpecNet
specLeakyRelu2D (i, j) = specLeakyRelu3D (i, j, 1)

-- | Create a specification for a elu layer.
specLeakyRelu3D :: (Integer, Integer, Integer) -> SpecNet
specLeakyRelu3D = SpecNetLayer . SpecLeakyRelu

-- | Add a LeakyRelu layer to your build.
leakyRelu :: BuildM ()
leakyRelu = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecLeakyRelu

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

instance GNum LeakyRelu where
  _ |* LeakyRelu = LeakyRelu
  _ |+ LeakyRelu = LeakyRelu


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

module Numeric.LinearAlgebra.Arnoldi
       ( Arpack, eig
         -- * Options
       , Options(..), Which(..)
         -- * Exceptions
       , MaxIterations(..)
       , NoShifts(..)
       , Reallocate(..)
       , XYAUPD(..)
       , XYEUPD(..)
       , Unimplemented(..)
       ) where

import Data.Vector.Storable (Vector)
import Data.Vector.Storable.Mutable (IOVector)
import Numeric.LinearAlgebra (Matrix)
import System.IO.Unsafe (unsafePerformIO)

import Arpack.Exceptions
import Arpack.Foreign
import Arpack.Options

eig :: Arpack t => Options t -> Int -> (IOVector t -> IOVector t -> IO ())
    -> (Vector t, Matrix t)
eig !options !dim !multiply = unsafePerformIO (arpack options dim multiply)


{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE ConstrainedClassMethods  #-}
{-# LANGUAGE DeriveGeneric            #-}
{-# LANGUAGE FlexibleContexts         #-}
{-# LANGUAGE FlexibleInstances        #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE TypeFamilies             #-}
{-# LANGUAGE TypeOperators            #-}
{-# LANGUAGE ViewPatterns             #-}

-- |
-- Module      :  Internal.Matrix
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--
-- Internal matrix representation
--

module Internal.Matrix where

import           Control.DeepSeq       (NFData (..))
import           Data.Complex          (Complex)
import           Foreign.C.String      (CString, newCString)
import           Foreign.C.Types       (CInt (..))
import           Foreign.Marshal.Alloc (free)
import           Foreign.Marshal.Array (newArray)
import           Foreign.Ptr           (Ptr)
import           Foreign.Storable      (Storable)
import           Internal.Devel
import           Internal.Vector
import           Internal.Vectorized   hiding (( # ), ( #! ))
import           System.IO.Unsafe      (unsafePerformIO)
import           Text.Printf

-----------------------------------------------------------------

data MatrixOrder = RowMajor | ColumnMajor deriving (Show,Eq)

-- | Matrix representation suitable for BLAS\/LAPACK computations.

data Matrix t = Matrix
    { irows :: {-# UNPACK #-} !Int
    , icols :: {-# UNPACK #-} !Int
    , xRow  :: {-# UNPACK #-} !Int
    , xCol  :: {-# UNPACK #-} !Int
    , xdat  :: {-# UNPACK #-} !(Vector t)
    }

rows :: Matrix t -> Int
rows = irows
{-# INLINE rows #-}

cols :: Matrix t -> Int
cols = icols
{-# INLINE cols #-}

size :: Matrix t -> (Int, Int)
size m = (irows m, icols m)
{-# INLINE size #-}

rowOrder :: Matrix t -> Bool
rowOrder m = xCol m == 1 || cols m == 1
{-# INLINE rowOrder #-}

colOrder :: Matrix t -> Bool
colOrder m = xRow m == 1 || rows m == 1
{-# INLINE colOrder #-}

is1d :: Matrix t -> Bool
is1d (size->(r,c)) = r==1 || c==1
{-# INLINE is1d #-}

-- data is not contiguous
isSlice :: Storable t => Matrix t -> Bool
isSlice m@(size->(r,c)) = r*c < dim (xdat m)
{-# INLINE isSlice #-}

orderOf :: Matrix t -> MatrixOrder
orderOf m = if rowOrder m then RowMajor else ColumnMajor

showInternal :: Storable t => Matrix t -> IO ()
showInternal m = printf "%dx%d %s %s %d:%d (%d)\n" r c slc ord xr xc dv
  where
    r  = rows m
    c  = cols m
    xr = xRow m
    xc = xCol m
    slc = if isSlice m then "slice" else "full"
    ord = if is1d m then "1d" else if rowOrder m then "rows" else "cols"
    dv = dim (xdat m)

--------------------------------------------------------------------------------

-- | Matrix transpose.
trans :: Matrix t -> Matrix t
trans m@Matrix { irows = r, icols = c, xRow = xr, xCol = xc } =
             m { irows = c, icols = r, xRow = xc, xCol = xr }

cmat :: (Element t) => Matrix t -> Matrix t
cmat m
    | rowOrder m = m
    | otherwise  = extractAll RowMajor m

fmat :: (Element t) => Matrix t -> Matrix t
fmat m
    | colOrder m = m
    | otherwise  = extractAll ColumnMajor m

-- C-Haskell matrix adapters
{-# INLINE amatr #-}
amatr :: Storable a => Matrix a -> (f -> IO r) -> (CInt -> CInt -> Ptr a -> f) -> IO r
amatr x f g = unsafeWith (xdat x) (f . g r c)
  where
    r  = fi (rows x)
    c  = fi (cols x)

{-# INLINE amat #-}
amat :: Storable a => Matrix a -> (f -> IO r) -> (CInt -> CInt -> CInt -> CInt -> Ptr a -> f) -> IO r
amat x f g = unsafeWith (xdat x) (f . g r c sr sc)
  where
    r  = fi (rows x)
    c  = fi (cols x)
    sr = fi (xRow x)
    sc = fi (xCol x)

instance Storable t => TransArray (Matrix t)
  where
    type TransRaw (Matrix t) b = CInt -> CInt -> Ptr t -> b
    type Trans (Matrix t) b    = CInt -> CInt -> CInt -> CInt -> Ptr t -> b
    apply = amat
    {-# INLINE apply #-}
    applyRaw = amatr
    {-# INLINE applyRaw #-}

infixr 1 #
(#) :: TransArray c => c -> (b -> IO r) -> Trans c b -> IO r
a # b = apply a b
{-# INLINE (#) #-}

(#!) :: (TransArray c, TransArray c1) => c1 -> c -> Trans c1 (Trans c (IO r)) -> IO r
a #! b = a # b # id
{-# INLINE (#!) #-}

--------------------------------------------------------------------------------

copy :: Element t => MatrixOrder -> Matrix t -> IO (Matrix t)
copy ord m = extractR ord m 0 (idxs[0,rows m-1]) 0 (idxs[0,cols m-1])

extractAll :: Element t => MatrixOrder -> Matrix t -> Matrix t
extractAll ord m = unsafePerformIO (copy ord m)

{- | Creates a vector by concatenation of rows. If the matrix is ColumnMajor, this operation requires a transpose.

>>> flatten (ident 3)
[1.0,0.0,0.0,0.0,1.0,0.0,0.0,0.0,1.0]
it :: (Num t, Element t) => Vector t

-}
flatten :: Element t => Matrix t -> Vector t
flatten m
    | isSlice m || not (rowOrder m) = xdat (extractAll RowMajor m)
    | otherwise                     = xdat m

-- | the inverse of 'Data.Packed.Matrix.fromLists'
toLists :: (Element t) => Matrix t -> [[t]]
toLists = map toList . toRows

-- | common value with \"adaptable\" 1
compatdim :: [Int] -> Maybe Int
compatdim [] = Nothing
compatdim [a] = Just a
compatdim (a:b:xs)
    | a==b = compatdim (b:xs)
    | a==1 = compatdim (b:xs)
    | b==1 = compatdim (a:xs)
    | otherwise = Nothing

-- | Create a matrix from a list of vectors.
-- All vectors must have the same dimension,
-- or dimension 1, which is are automatically expanded.
fromRows :: Element t => [Vector t] -> Matrix t
fromRows [] = emptyM 0 0
fromRows vs = case compatdim (map dim vs) of
    Nothing -> error $ "fromRows expects vectors with equal sizes (or singletons), given: " ++ show (map dim vs)
    Just 0  -> emptyM r 0
    Just c  -> matrixFromVector RowMajor r c . vjoin . map (adapt c) $ vs
  where
    r = length vs
    adapt c v
        | c == 0 = fromList[]
        | dim v == c = v
        | otherwise = constantD (v@>0) c

-- | extracts the rows of a matrix as a list of vectors
toRows :: Element t => Matrix t -> [Vector t]
toRows m
    | rowOrder m = map sub rowRange
    | otherwise  = map ext rowRange
  where
    rowRange = [0..rows m-1]
    sub k = subVector (k*xRow m) (cols m) (xdat m)
    ext k = xdat $ unsafePerformIO $ extractR RowMajor m 1 (idxs[k]) 0 (idxs[0,cols m-1])

-- | Creates a matrix from a list of vectors, as columns
fromColumns :: Element t => [Vector t] -> Matrix t
fromColumns m = trans . fromRows $ m

-- | Creates a list of vectors from the columns of a matrix
toColumns :: Element t => Matrix t -> [Vector t]
toColumns m = toRows . trans $ m

-- | Reads a matrix position.
(@@>) :: Storable t => Matrix t -> (Int,Int) -> t
infixl 9 @@>
m@Matrix {irows = r, icols = c} @@> (i,j)
    | i<0 || i>=r || j<0 || j>=c = error "matrix indexing out of range"
    | otherwise = atM' m i j
{-# INLINE (@@>) #-}

--  Unsafe matrix access without range checking
atM' :: Storable t => Matrix t -> Int -> Int -> t
atM' m i j = xdat m `at'` (i * (xRow m) + j * (xCol m))
{-# INLINE atM' #-}

------------------------------------------------------------------

matrixFromVector :: Storable t => MatrixOrder -> Int -> Int -> Vector t -> Matrix t
matrixFromVector _ 1 _ v@(dim->d) = Matrix { irows = 1, icols = d, xdat = v, xRow = d, xCol = 1 }
matrixFromVector _ _ 1 v@(dim->d) = Matrix { irows = d, icols = 1, xdat = v, xRow = 1, xCol = d }
matrixFromVector o r c v
    | r * c == dim v = m
    | otherwise = error $ "can't reshape vector dim = "++ show (dim v)++" to matrix " ++ shSize m
  where
    m | o == RowMajor = Matrix { irows = r, icols = c, xdat = v, xRow = c, xCol = 1 }
      | otherwise     = Matrix { irows = r, icols = c, xdat = v, xRow = 1, xCol = r }

-- allocates memory for a new matrix
createMatrix :: (Storable a) => MatrixOrder -> Int -> Int -> IO (Matrix a)
createMatrix ord r c = do
    p <- createVector (r*c)
    return (matrixFromVector ord r c p)

{- | Creates a matrix from a vector by grouping the elements in rows with the desired number of columns. (GNU-Octave groups by columns. To do it you can define @reshapeF r = tr' . reshape r@
where r is the desired number of rows.)

>>> reshape 4 (fromList [1..12])
(3><4)
 [ 1.0,  2.0,  3.0,  4.0
 , 5.0,  6.0,  7.0,  8.0
 , 9.0, 10.0, 11.0, 12.0 ]

-}
reshape :: Storable t => Int -> Vector t -> Matrix t
reshape 0 v = matrixFromVector RowMajor 0 0 v
reshape c v = matrixFromVector RowMajor (dim v `div` c) c v

-- | application of a vector function on the flattened matrix elements
liftMatrix :: (Element a, Element b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
liftMatrix f m@Matrix { irows = r, icols = c, xdat = d}
    | isSlice m = matrixFromVector RowMajor r c (f (flatten m))
    | otherwise = matrixFromVector (orderOf m) r c (f d)

-- | application of a vector function on the flattened matrices elements
liftMatrix2 :: (Element t, Element a, Element b) => (Vector a -> Vector b -> Vector t) -> Matrix a -> Matrix b -> Matrix t
liftMatrix2 f m1@(size->(r,c)) m2
    | (r,c)/=size m2 = error "nonconformant matrices in liftMatrix2"
    | rowOrder m1 = matrixFromVector RowMajor    r c (f (flatten m1) (flatten m2))
    | otherwise   = matrixFromVector ColumnMajor r c (f (flatten (trans m1)) (flatten (trans m2)))

------------------------------------------------------------------

-- | Supported matrix elements.
class (Storable a) => Element a where
    constantD  :: a -> Int -> Vector a
    extractR :: MatrixOrder -> Matrix a -> CInt -> Vector CInt -> CInt -> Vector CInt -> IO (Matrix a)
    setRect  :: Int -> Int -> Matrix a -> Matrix a -> IO ()
    sortI    :: Ord a => Vector a -> Vector CInt
    sortV    :: Ord a => Vector a -> Vector a
    compareV :: Ord a => Vector a -> Vector a -> Vector CInt
    selectV  :: Vector CInt -> Vector a -> Vector a -> Vector a -> Vector a
    remapM   :: Matrix CInt -> Matrix CInt -> Matrix a -> Matrix a
    rowOp    :: Int -> a -> Int -> Int -> Int -> Int -> Matrix a -> IO ()
    gemm     :: Vector a -> Matrix a -> Matrix a -> Matrix a -> IO ()
    reorderV :: Vector CInt-> Vector CInt-> Vector a -> Vector a -- see reorderVector for documentation

instance Element Float where
    constantD  = constantAux cconstantF
    extractR   = extractAux c_extractF
    setRect    = setRectAux c_setRectF
    sortI      = sortIdxF
    sortV      = sortValF
    compareV   = compareF
    selectV    = selectF
    remapM     = remapF
    rowOp      = rowOpAux c_rowOpF
    gemm       = gemmg c_gemmF
    reorderV   = reorderAux c_reorderF

-- instance Element Double where
--     constantD  = constantAux cconstantR
--     extractR   = extractAux c_extractD
--     setRect    = setRectAux c_setRectD
--     sortI      = sortIdxD
--     sortV      = sortValD
--     compareV   = compareD
--     selectV    = selectD
--     remapM     = remapD
--     rowOp      = rowOpAux c_rowOpD
--     gemm       = gemmg c_gemmD
--     reorderV   = reorderAux c_reorderD

instance Element (Complex Float) where
    constantD  = constantAux cconstantQ
    extractR   = extractAux c_extractQ
    setRect    = setRectAux c_setRectQ
    sortI      = undefined
    sortV      = undefined
    compareV   = undefined
    selectV    = selectQ
    remapM     = remapQ
    rowOp      = rowOpAux c_rowOpQ
    gemm       = gemmg c_gemmQ
    reorderV   = reorderAux c_reorderQ

-- instance Element (Complex Double) where
--     constantD  = constantAux cconstantC
--     extractR   = extractAux c_extractC
--     setRect    = setRectAux c_setRectC
--     sortI      = undefined
--     sortV      = undefined
--     compareV   = undefined
--     selectV    = selectC
--     remapM     = remapC
--     rowOp      = rowOpAux c_rowOpC
--     gemm       = gemmg c_gemmC
--     reorderV   = reorderAux c_reorderC

instance Element (CInt) where
    constantD  = constantAux cconstantI
    extractR   = extractAux c_extractI
    setRect    = setRectAux c_setRectI
    sortI      = sortIdxI
    sortV      = sortValI
    compareV   = compareI
    selectV    = selectI
    remapM     = remapI
    rowOp      = rowOpAux c_rowOpI
    gemm       = gemmg c_gemmI
    reorderV   = reorderAux c_reorderI

instance Element Z where
    constantD  = constantAux cconstantL
    extractR   = extractAux c_extractL
    setRect    = setRectAux c_setRectL
    sortI      = sortIdxL
    sortV      = sortValL
    compareV   = compareL
    selectV    = selectL
    remapM     = remapL
    rowOp      = rowOpAux c_rowOpL
    gemm       = gemmg c_gemmL
    reorderV   = reorderAux c_reorderL

-------------------------------------------------------------------

-- | reference to a rectangular slice of a matrix (no data copy)
subMatrix :: Element a
            => (Int,Int) -- ^ (r0,c0) starting position
            -> (Int,Int) -- ^ (rt,ct) dimensions of submatrix
            -> Matrix a -- ^ input matrix
            -> Matrix a -- ^ result
subMatrix (r0,c0) (rt,ct) m
    | rt <= 0 || ct <= 0 = matrixFromVector RowMajor (max 0 rt) (max 0 ct) (fromList [])
    | 0 <= r0 && 0 <= rt && r0+rt <= rows m &&
      0 <= c0 && 0 <= ct && c0+ct <= cols m = res
    | otherwise = error $ "wrong subMatrix "++show ((r0,c0),(rt,ct))++" of "++shSize m
  where
    p = r0 * xRow m + c0 * xCol m
    tot | rowOrder m = ct + (rt-1) * xRow m
        | otherwise  = rt + (ct-1) * xCol m
    res = m { irows = rt, icols = ct, xdat = subVector p tot (xdat m) }

--------------------------------------------------------------------------

maxZ :: (Num t1, Ord t1, Foldable t) => t t1 -> t1
maxZ xs = if minimum xs == 0 then 0 else maximum xs

conformMs :: Element t => [Matrix t] -> [Matrix t]
conformMs ms = map (conformMTo (r,c)) ms
  where
    r = maxZ (map rows ms)
    c = maxZ (map cols ms)

conformVs :: Element t => [Vector t] -> [Vector t]
conformVs vs = map (conformVTo n) vs
  where
    n = maxZ (map dim vs)

conformMTo :: Element t => (Int, Int) -> Matrix t -> Matrix t
conformMTo (r,c) m
    | size m == (r,c) = m
    | size m == (1,1) = matrixFromVector RowMajor r c (constantD (m@@>(0,0)) (r*c))
    | size m == (r,1) = repCols c m
    | size m == (1,c) = repRows r m
    | otherwise = error $ "matrix " ++ shSize m ++ " cannot be expanded to " ++ shDim (r,c)

conformVTo :: Element t => Int -> Vector t -> Vector t
conformVTo n v
    | dim v == n = v
    | dim v == 1 = constantD (v@>0) n
    | otherwise = error $ "vector of dim=" ++ show (dim v) ++ " cannot be expanded to dim=" ++ show n

repRows :: Element t => Int -> Matrix t -> Matrix t
repRows n x = fromRows (replicate n (flatten x))
repCols :: Element t => Int -> Matrix t -> Matrix t
repCols n x = fromColumns (replicate n (flatten x))

shSize :: Matrix t -> [Char]
shSize = shDim . size

shDim :: (Show a, Show a1) => (a1, a) -> [Char]
shDim (r,c) = "(" ++ show r ++"x"++ show c ++")"

emptyM :: Storable t => Int -> Int -> Matrix t
emptyM r c = matrixFromVector RowMajor r c (fromList[])

----------------------------------------------------------------------

instance (Storable t, NFData t) => NFData (Matrix t)
  where
    rnf m | d > 0     = rnf (v @> 0)
          | otherwise = ()
      where
        d = dim v
        v = xdat m

---------------------------------------------------------------

extractAux :: (Eq t3, Eq t2, TransArray c, Storable a, Storable t1,
                Storable t, Num t3, Num t2, Integral t1, Integral t)
           => (t3 -> t2 -> CInt -> Ptr t1 -> CInt -> Ptr t
                  -> Trans c (CInt -> CInt -> CInt -> CInt -> Ptr a -> IO CInt))
           -> MatrixOrder -> c -> t3 -> Vector t1 -> t2 -> Vector t -> IO (Matrix a)
extractAux f ord m moder vr modec vc = do
    let nr = if moder == 0 then fromIntegral $ vr@>1 - vr@>0 + 1 else dim vr
        nc = if modec == 0 then fromIntegral $ vc@>1 - vc@>0 + 1 else dim vc
    r <- createMatrix ord nr nc
    (vr # vc # m #! r) (f moder modec)  #|"extract"

    return r

type Extr x = CInt -> CInt -> CIdxs (CIdxs (OM x (OM x (IO CInt))))

-- foreign import ccall unsafe "extractD" c_extractD :: Extr Double
foreign import ccall unsafe "extractF" c_extractF :: Extr Float
-- foreign import ccall unsafe "extractC" c_extractC :: Extr (Complex Double)
foreign import ccall unsafe "extractQ" c_extractQ :: Extr (Complex Float)
foreign import ccall unsafe "extractI" c_extractI :: Extr CInt
foreign import ccall unsafe "extractL" c_extractL :: Extr Z

---------------------------------------------------------------

setRectAux :: (TransArray c1, TransArray c)
           => (CInt -> CInt -> Trans c1 (Trans c (IO CInt)))
           -> Int -> Int -> c1 -> c -> IO ()
setRectAux f i j m r = (m #! r) (f (fi i) (fi j)) #|"setRect"

type SetRect x = I -> I -> x ::> x::> Ok

-- foreign import ccall unsafe "setRectD" c_setRectD :: SetRect Double
foreign import ccall unsafe "setRectF" c_setRectF :: SetRect Float
-- foreign import ccall unsafe "setRectC" c_setRectC :: SetRect (Complex Double)
foreign import ccall unsafe "setRectQ" c_setRectQ :: SetRect (Complex Float)
foreign import ccall unsafe "setRectI" c_setRectI :: SetRect I
foreign import ccall unsafe "setRectL" c_setRectL :: SetRect Z

--------------------------------------------------------------------------------

sortG :: (Storable t, Storable a)
      => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
sortG f v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) f #|"sortG"
    return r

-- sortIdxD :: Vector Double -> Vector CInt
-- sortIdxD = sortG c_sort_indexD
sortIdxF :: Vector Float -> Vector CInt
sortIdxF = sortG c_sort_indexF
sortIdxI :: Vector CInt -> Vector CInt
sortIdxI = sortG c_sort_indexI
sortIdxL :: Vector Z -> Vector I
sortIdxL = sortG c_sort_indexL

-- sortValD :: Vector Double -> Vector Double
-- sortValD = sortG c_sort_valD
sortValF :: Vector Float -> Vector Float
sortValF = sortG c_sort_valF
sortValI :: Vector CInt -> Vector CInt
sortValI = sortG c_sort_valI
sortValL :: Vector Z -> Vector Z
sortValL = sortG c_sort_valL

-- foreign import ccall unsafe "sort_indexD" c_sort_indexD :: CV Double (CV CInt (IO CInt))
foreign import ccall unsafe "sort_indexF" c_sort_indexF :: CV Float  (CV CInt (IO CInt))
foreign import ccall unsafe "sort_indexI" c_sort_indexI :: CV CInt   (CV CInt (IO CInt))
foreign import ccall unsafe "sort_indexL" c_sort_indexL :: Z :> I :> Ok

-- foreign import ccall unsafe "sort_valuesD" c_sort_valD :: CV Double (CV Double (IO CInt))
foreign import ccall unsafe "sort_valuesF" c_sort_valF :: CV Float  (CV Float (IO CInt))
foreign import ccall unsafe "sort_valuesI" c_sort_valI :: CV CInt   (CV CInt (IO CInt))
foreign import ccall unsafe "sort_valuesL" c_sort_valL :: Z :> Z :> Ok

--------------------------------------------------------------------------------

compareG :: (TransArray c, Storable t, Storable a)
         => Trans c (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt)
         -> c -> Vector t -> Vector a
compareG f u v = unsafePerformIO $ do
    r <- createVector (dim v)
    (u # v #! r) f #|"compareG"
    return r

-- compareD :: Vector Double -> Vector Double -> Vector CInt
-- compareD = compareG c_compareD
compareF :: Vector Float -> Vector Float -> Vector CInt
compareF = compareG c_compareF
compareI :: Vector CInt -> Vector CInt -> Vector CInt
compareI = compareG c_compareI
compareL :: Vector Z -> Vector Z -> Vector CInt
compareL = compareG c_compareL

-- foreign import ccall unsafe "compareD" c_compareD :: CV Double (CV Double (CV CInt (IO CInt)))
foreign import ccall unsafe "compareF" c_compareF :: CV Float (CV Float  (CV CInt (IO CInt)))
foreign import ccall unsafe "compareI" c_compareI :: CV CInt (CV CInt   (CV CInt (IO CInt)))
foreign import ccall unsafe "compareL" c_compareL :: Z :> Z :> I :> Ok

--------------------------------------------------------------------------------

selectG :: (TransArray c, TransArray c1, TransArray c2, Storable t, Storable a)
        => Trans c2 (Trans c1 (CInt -> Ptr t -> Trans c (CInt -> Ptr a -> IO CInt)))
        -> c2 -> c1 -> Vector t -> c -> Vector a
selectG f c u v w = unsafePerformIO $ do
    r <- createVector (dim v)
    (c # u # v # w #! r) f #|"selectG"
    return r

-- selectD :: Vector CInt -> Vector Double -> Vector Double -> Vector Double -> Vector Double
-- selectD = selectG c_selectD
selectF :: Vector CInt -> Vector Float -> Vector Float -> Vector Float -> Vector Float
selectF = selectG c_selectF
selectI :: Vector CInt -> Vector CInt -> Vector CInt -> Vector CInt -> Vector CInt
selectI = selectG c_selectI
selectL :: Vector CInt -> Vector Z -> Vector Z -> Vector Z -> Vector Z
selectL = selectG c_selectL
-- selectC :: Vector CInt
--         -> Vector (Complex Double)
--         -> Vector (Complex Double)
--         -> Vector (Complex Double)
--         -> Vector (Complex Double)
-- selectC = selectG c_selectC
selectQ :: Vector CInt
        -> Vector (Complex Float)
        -> Vector (Complex Float)
        -> Vector (Complex Float)
        -> Vector (Complex Float)
selectQ = selectG c_selectQ

type Sel x = CV CInt (CV x (CV x (CV x (CV x (IO CInt)))))

-- foreign import ccall unsafe "chooseD" c_selectD :: Sel Double
foreign import ccall unsafe "chooseF" c_selectF :: Sel Float
foreign import ccall unsafe "chooseI" c_selectI :: Sel CInt
-- foreign import ccall unsafe "chooseC" c_selectC :: Sel (Complex Double)
foreign import ccall unsafe "chooseQ" c_selectQ :: Sel (Complex Float)
foreign import ccall unsafe "chooseL" c_selectL :: Sel Z

---------------------------------------------------------------------------

remapG :: (TransArray c, TransArray c1, Storable t, Storable a)
       => (CInt -> CInt -> CInt -> CInt -> Ptr t
                -> Trans c1 (Trans c (CInt -> CInt -> CInt -> CInt -> Ptr a -> IO CInt)))
       -> Matrix t -> c1 -> c -> Matrix a
remapG f i j m = unsafePerformIO $ do
    r <- createMatrix RowMajor (rows i) (cols i)
    (i # j # m #! r) f #|"remapG"
    return r

-- remapD :: Matrix CInt -> Matrix CInt -> Matrix Double -> Matrix Double
-- remapD = remapG c_remapD
remapF :: Matrix CInt -> Matrix CInt -> Matrix Float -> Matrix Float
remapF = remapG c_remapF
remapI :: Matrix CInt -> Matrix CInt -> Matrix CInt -> Matrix CInt
remapI = remapG c_remapI
remapL :: Matrix CInt -> Matrix CInt -> Matrix Z -> Matrix Z
remapL = remapG c_remapL
-- remapC :: Matrix CInt
--        -> Matrix CInt
--        -> Matrix (Complex Double)
--        -> Matrix (Complex Double)
-- remapC = remapG c_remapC
remapQ :: Matrix CInt -> Matrix CInt -> Matrix (Complex Float) -> Matrix (Complex Float)
remapQ = remapG c_remapQ

type Rem x = OM CInt (OM CInt (OM x (OM x (IO CInt))))

-- foreign import ccall unsafe "remapD" c_remapD :: Rem Double
foreign import ccall unsafe "remapF" c_remapF :: Rem Float
foreign import ccall unsafe "remapI" c_remapI :: Rem CInt
-- foreign import ccall unsafe "remapC" c_remapC :: Rem (Complex Double)
foreign import ccall unsafe "remapQ" c_remapQ :: Rem (Complex Float)
foreign import ccall unsafe "remapL" c_remapL :: Rem Z

--------------------------------------------------------------------------------

rowOpAux :: (TransArray c, Storable a) =>
            (CInt -> Ptr a -> CInt -> CInt -> CInt -> CInt -> Trans c (IO CInt))
         -> Int -> a -> Int -> Int -> Int -> Int -> c -> IO ()
rowOpAux f c x i1 i2 j1 j2 m = do
    px <- newArray [x]
    (m # id) (f (fi c) px (fi i1) (fi i2) (fi j1) (fi j2)) #|"rowOp"
    free px

type RowOp x = CInt -> Ptr x -> CInt -> CInt -> CInt -> CInt -> x ::> Ok

-- foreign import ccall unsafe "rowop_double"  c_rowOpD :: RowOp R
foreign import ccall unsafe "rowop_float"   c_rowOpF :: RowOp Float
-- foreign import ccall unsafe "rowop_TCD"     c_rowOpC :: RowOp C
foreign import ccall unsafe "rowop_TCS"     c_rowOpQ :: RowOp (Complex Float)
foreign import ccall unsafe "rowop_int32_t" c_rowOpI :: RowOp I
foreign import ccall unsafe "rowop_int64_t" c_rowOpL :: RowOp Z
foreign import ccall unsafe "rowop_mod_int32_t" c_rowOpMI :: I -> RowOp I
foreign import ccall unsafe "rowop_mod_int64_t" c_rowOpML :: Z -> RowOp Z

--------------------------------------------------------------------------------

gemmg :: (TransArray c1, TransArray c, TransArray c2, TransArray c3)
      => Trans c3 (Trans c2 (Trans c1 (Trans c (IO CInt))))
      -> c3 -> c2 -> c1 -> c -> IO ()
gemmg f v m1 m2 m3 = (v # m1 # m2 #! m3) f #|"gemmg"

type Tgemm x = x :> x ::> x ::> x ::> Ok

-- foreign import ccall unsafe "gemm_double"  c_gemmD :: Tgemm R
foreign import ccall unsafe "gemm_float"   c_gemmF :: Tgemm Float
-- foreign import ccall unsafe "gemm_TCD"     c_gemmC :: Tgemm C
foreign import ccall unsafe "gemm_TCS"     c_gemmQ :: Tgemm (Complex Float)
foreign import ccall unsafe "gemm_int32_t" c_gemmI :: Tgemm I
foreign import ccall unsafe "gemm_int64_t" c_gemmL :: Tgemm Z
foreign import ccall unsafe "gemm_mod_int32_t" c_gemmMI :: I -> Tgemm I
foreign import ccall unsafe "gemm_mod_int64_t" c_gemmML :: Z -> Tgemm Z

--------------------------------------------------------------------------------

reorderAux :: (TransArray c, Storable t, Storable a1, Storable t1, Storable a) =>
              (CInt -> Ptr a -> CInt -> Ptr t1
                    -> Trans c (CInt -> Ptr t -> CInt -> Ptr a1 -> IO CInt))
           -> Vector t1 -> c -> Vector t -> Vector a1
reorderAux f s d v = unsafePerformIO $ do
    k <- createVector (dim s)
    r <- createVector (dim v)
    (k # s # d # v #! r) f #| "reorderV"
    return r

type Reorder x = CV CInt (CV CInt (CV CInt (CV x (CV x (IO CInt)))))

-- foreign import ccall unsafe "reorderD" c_reorderD :: Reorder Double
foreign import ccall unsafe "reorderF" c_reorderF :: Reorder Float
foreign import ccall unsafe "reorderI" c_reorderI :: Reorder CInt
-- foreign import ccall unsafe "reorderC" c_reorderC :: Reorder (Complex Double)
foreign import ccall unsafe "reorderQ" c_reorderQ :: Reorder (Complex Float)
foreign import ccall unsafe "reorderL" c_reorderL :: Reorder Z

-- | Transpose an array with dimensions @dims@ by making a copy using @strides@. For example, for an array with 3 indices,
--   @(reorderVector strides dims v) ! ((i * dims ! 1 + j) * dims ! 2 + k) == v ! (i * strides ! 0 + j * strides ! 1 + k * strides ! 2)@
--   This function is intended to be used internally by tensor libraries.
reorderVector :: Element a
                    => Vector CInt -- ^ @strides@: array strides
                    -> Vector CInt -- ^ @dims@: array dimensions of new array @v@
                    -> Vector a    -- ^ @v@: flattened input array
                    -> Vector a    -- ^ @v'@: flattened output array
reorderVector = reorderV

--------------------------------------------------------------------------------

foreign import ccall unsafe "saveMatrix" c_saveMatrix
    :: CString -> CString -> Double ::> Ok

{- | save a matrix as a 2D ASCII table
-}
saveMatrix
    :: FilePath
    -> String        -- ^ \"printf\" format (e.g. \"%.2f\", \"%g\", etc.)
    -> Matrix Double
    -> IO ()
saveMatrix name format m = do
    cname   <- newCString name
    cformat <- newCString format
    (m # id) (c_saveMatrix cname cformat) #|"saveMatrix"
    free cname
    free cformat
    return ()

--------------------------------------------------------------------------------


{-# LANGUAGE OverloadedStrings #-}
module Lisp.PrintSpec (spec) where

import Data.Complex
import Data.Ratio
import Data.Text
import Test.Hspec
import Test.QuickCheck

import Lisp

spec :: Spec

spec = do
  let foo = Symbol "foo"
  let bar = Symbol "bar"

  describe "LispVal.pp" $ do
    it "prints symbols as its own value" $
      pp foo `shouldBe` "foo"

    it "QuickCheck symbols print as their own value" $
      property $ \s -> let t = pack s in pp (Symbol t) == t

    it "prints integers as their own value" $
      pp (Int 5) `shouldBe` "5"

    it "prints decimals as thier own value" $
      pp (Real 3.14) `shouldBe` "3.14"

    it "prints rationals as their own value, lisp-style" $
      pp (Rational (3%4)) `shouldBe` "3/4"

    it "prints rationals reduced" $
      pp (Rational (4%8)) `shouldBe` "1/2"

    it "prints complex number as their own value, lisp-style" $
      pp (Complex (3:+4)) `shouldBe` "3.0+4.0i"

    it "prints strings in quotes" $
      pp (String "foo") `shouldBe` "\"foo\""

    it "prints true as #t" $
      pp (Bool True) `shouldBe` "#t"

    it "prints false as #f" $
      pp (Bool False) `shouldBe` "#f"

    it "prints nils/empty lists as ()" $
      pp Nil `shouldBe` "()"

    it "prints pairs as dotted pairs" $
      pp (Pair foo bar) `shouldBe` "(foo . bar)"

    it "prints pairs with Nil cdr as singleton lists" $
      pp (Pair foo Nil) `shouldBe` "(foo)"

    it "prints lists as lists" $
      pp (Pair foo
           (Pair bar
            (Pair (Int 5)
             (Pair (Bool False) Nil)))) `shouldBe` "(foo bar 5 #f)"

    it "prints dotted lists with dots" $
      pp (Pair foo
           (Pair bar
            (Pair (Int 5)
             (Pair (Bool False) (Bool True))))) `shouldBe` "(foo bar 5 #f . #t)"


module STCPinwheelCorner where

import           Control.Monad                  as M
import           Control.Monad.Parallel         as MP
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           Filter.Utils
import           FokkerPlanck
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           FourierPinwheel
import           Image.IO
import           Pinwheel.FourierSeries2D
import           STC                            hiding (convolve)
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.Time

main = do
  args@(deviceIDsStr:numPointsStr:deltaStr:thresholdStr:numPointsReconStr:deltaReconStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:tauStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initDistStr:initScaleStr:histFilePath:stdR2Str:stdStr:numBatchR2Str:numBatchR2FreqsStr:numBatchOriStr:batchSizeStr:sStr:radiusStr:deltaTStr:weightStr:posStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      delta = read deltaStr :: Double
      threshold = read thresholdStr :: Double
      numPointsRecon = read numPointsReconStr :: Int
      deltaRecon = read deltaReconStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tau = read tauStr :: Double
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      initScale = read initScaleStr :: Double
      initDist = read initDistStr :: [(Double, Double, Double, Double)]
      initPoints = L.map (\(x, y, t, s) -> Point x y t s) initDist
      initSource = [L.head initPoints]
      initSink = [L.last initPoints]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCPinwheelCorner"
      stdR2 = read stdR2Str :: Double
      std = read stdStr :: Double
      numBatchR2 = read numBatchR2Str :: Int
      numBatchR2Freqs = read numBatchR2FreqsStr :: Int
      numBatchOri = read numBatchOriStr :: Int
      batchSize = read batchSizeStr :: Int
      s = read sStr :: Double
      radius = read radiusStr :: Double
      deltaT = read deltaTStr :: Double
      weight = read weightStr :: Double
      pos = read posStr :: (Int, Int)
      periodEnv = periodR2 ^ 2 / 4
  -- removePathForcibly folderPath
  createDirectoryIfMissing True folderPath
  -- let deltaTheta = 2 * pi / fromIntegral numOrientation
  --     cornerDist =
  --       cornerDistribution
  --         delta
  --         initScale
  --         thetaSigma
  --         tau
  --         threshold
  --         [fromIntegral i * deltaTheta | i <- [0 .. numOrientation - 1]]
  --         weight
  --         pos
  --     idx = [(i, j) | i <- [-10 .. 10], j <- [-10 .. 10]]
  --     xs =
  --       L.map
  --         (\(i, j) ->
  --            let x = i * delta
  --                y = j * delta
  --             in ( (x, y)
  --                , L.sum $
  --                  cornerDistribution
  --                    delta
  --                    initScale
  --                    thetaSigma
  --                    tau
  --                    threshold
  --                    [ fromIntegral o * deltaTheta
  --                    | o <- [0 .. numOrientation - 1]
  --                    ]
  --                    weight
  --                    (x, y)))
  --         idx
  -- print . L.sortOn snd $ xs
  -- printf "%f\n" (L.sum cornerDist)
  -- print .
  --   L.sortOn snd .
  --   L.zip
  --     [fromIntegral i * deltaTheta / pi * 180 | i <- [0 .. numOrientation - 1]] $
  --   cornerDist
  -- plotImageRepa (folderPath </> "test.png") .
  --   ImageRepa 8 .
  --   fromListUnboxed (Z :. (1 :: Int) :. (21 :: Int) :. (21 :: Int)) .
  --   snd . L.unzip $
  --   xs
  -- let maxIdx =
  --       fst .
  --       L.maximumBy (\(_, a) (_, b) -> compare a b) .
  --       L.zip
  --         [ fromIntegral i * deltaTheta / pi * 180
  --         | i <- [0 .. numOrientation - 1]
  --         ] $
  --       cornerDist
  -- maxDist <- cornerDistribution' delta initScale thetaSigma tau threshold [fromIntegral i * deltaTheta | i <- [0 .. numOrientation - 1]] weight pos (maxIdx / 180 * pi)
  -- print maxIdx
  -- print maxDist
  flag <- doesFileExist histFilePath
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  hist <-
    if flag
      then do
        printCurrentTime "read from files..."
        decodeFile histFilePath
      else sampleCartesianCorner
             histFilePath
             folderPath
             ptxs
             numPoints
             periodEnv
             delta
             numOrientation
             initScale
             thetaSigma
             tau
             threshold
             s
             phiFreq
             rhoFreq
             thetaFreq
             scaleFreq
             stdR2
             weight
  printCurrentTime "Done"
  printCurrentTime "Start Convloution.."
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPointsRecon
      numPointsRecon
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
      (2 * phiFreq + 1)
      (2 * rhoFreq + 1)
  printCurrentTime "Compute DFT Plan done."
  let coefficients = getNormalizedHistogramArr hist
  harmonicsArray <-
    createHarmonics
      numR2Freq
      phiFreq
      rhoFreq
      thetaFreq
      scaleFreq
      (-s)
      periodR2
      periodEnv
      coefficients
  let initDistSource =
        computeInitialDistributionFourierPinwheel
          numR2Freq
          periodR2
          periodEnv
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          initSource
      initDistSink =
        computeInitialDistributionFourierPinwheel
          numR2Freq
          periodR2
          periodEnv
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          initSink
  source <- convolve harmonicsArray initDistSource
  sink <- convolve harmonicsArray initDistSink
  sourceR2 <- plotFPArray plan (folderPath </> "Source.png") source
  sinkR2 <- plotFPArray plan (folderPath </> "Sink.png") sink
  plotRThetaDist
    (folderPath </> "Source_Theta.png")
    (folderPath </> "Source_R.png")
    numPointsRecon
    360
    90
    periodEnv
    pos
    sourceR2
  plotRThetaDist
    (folderPath </> "Sink_Theta.png")
    (folderPath </> "Sink_R.png")
    numPointsRecon
    360
    90
    periodEnv
    pos
    sinkR2
  sinkR2TR <-
    computeUnboxedP .
    timeReversalRepa [fromIntegral (-thetaFreq) .. fromIntegral thetaFreq] $
    sinkR2
  completionR2 <- completionFieldRepa plan sourceR2 sinkR2TR
  (sumP . sumS . R.map (\x -> magnitude x ** 2) . rotate4D2 $ completionR2) >>=
    plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    computeS .
    extend (Z :. (1 :: Int) :. All :. All) -- . R.map (\x -> log (x + 1))


{-# LANGUAGE FlexibleContexts #-}
-- |
-- Module    : System.Random.MWC.CondensedTable
-- Copyright : (c) 2012 Aleksey Khudyakov
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Table-driven generation of random variates.  This approach can
-- generate random variates in /O(1)/ time for the supported
-- distributions, at a modest cost in initialization time.
module System.Random.MWC.CondensedTable (
    -- * Condensed tables
    CondensedTable
  , CondensedTableV
  , CondensedTableU
  , genFromTable
    -- * Constructors for tables
  , tableFromProbabilities
  , tableFromWeights
  , tableFromIntWeights
    -- ** Disrete distributions
  , tablePoisson
  , tableBinomial
    -- * References
    -- $references
  ) where

import Control.Arrow           (second,(***))

import Data.Word
import Data.Int
import Data.Bits
import qualified Data.Vector.Generic         as G
import           Data.Vector.Generic           ((++))
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed         as U
import qualified Data.Vector                 as V
import Data.Vector.Generic (Vector)
import Numeric.SpecFunctions (logFactorial)
import System.Random.Stateful

import Prelude hiding ((++))

-- | A lookup table for arbitrary discrete distributions. It allows
-- the generation of random variates in /O(1)/. Note that probability
-- is quantized in units of @1/2^32@, and all distributions with
-- infinite support (e.g. Poisson) should be truncated.
data CondensedTable v a =
  CondensedTable
  {-# UNPACK #-} !Word64 !(v a) -- Lookup limit and first table
  {-# UNPACK #-} !Word64 !(v a) -- Second table
  {-# UNPACK #-} !Word64 !(v a) -- Third table
  !(v a)                        -- Last table

-- Implementation note. We have to store lookup limit in Word64 since
-- we need to accomodate two cases. First is when we have no values in
-- lookup table, second is when all elements are there
--
-- Both are pretty easy to realize. For first one probability of every
-- outcome should be less then 1/256, latter arise when probabilities
-- of two outcomes are [0.5,0.5]

-- | A 'CondensedTable' that uses unboxed vectors.
type CondensedTableU = CondensedTable U.Vector

-- | A 'CondensedTable' that uses boxed vectors, and is able to hold
-- any type of element.
type CondensedTableV = CondensedTable V.Vector

-- | Generate a random value using a condensed table.
genFromTable :: (StatefulGen g m, Vector v a) => CondensedTable v a -> g -> m a
{-# INLINE genFromTable #-}
genFromTable table gen = do
  w <- uniformM gen
  return $ lookupTable table $ fromIntegral (w :: Word32)

lookupTable :: Vector v a => CondensedTable v a -> Word64 -> a
{-# INLINE lookupTable #-}
lookupTable (CondensedTable na aa nb bb nc cc dd) i
  | i < na    = aa `at` ( i       `shiftR` 24)
  | i < nb    = bb `at` ((i - na) `shiftR` 16)
  | i < nc    = cc `at` ((i - nb) `shiftR` 8 )
  | otherwise = dd `at` ( i - nc)
  where
    at arr j = G.unsafeIndex arr (fromIntegral j)

----------------------------------------------------------------
-- Table generation
----------------------------------------------------------------

-- | Generate a condensed lookup table from a list of outcomes with
-- given probabilities. The vector should be non-empty and the
-- probabilites should be non-negative and sum to 1. If this is not
-- the case, this algorithm will construct a table for some
-- distribution that may bear no resemblance to what you intended.
tableFromProbabilities
    :: (Vector v (a,Word32), Vector v (a,Double), Vector v a, Vector v Word32)
       => v (a, Double) -> CondensedTable v a
{-# INLINE tableFromProbabilities #-}
tableFromProbabilities v
  | G.null tbl = pkgError "tableFromProbabilities" "empty vector of outcomes"
  | otherwise  = tableFromIntWeights $ G.map (second $ toWeight . (* mlt)) tbl
  where
    -- 2^32. N.B. This number is exatly representable.
    mlt = 4.294967296e9
    -- Drop non-positive probabilities
    tbl = G.filter ((> 0) . snd) v
    -- Convert Double weight to Word32 and avoid overflow at the same
    -- time. It's especially dangerous if one probability is
    -- approximately 1 and others are 0.
    toWeight w | w > mlt - 1 = 2^(32::Int) - 1
               | otherwise   = round w

-- | Same as 'tableFromProbabilities' but treats number as weights not
-- probilities. Non-positive weights are discarded, and those
-- remaining are normalized to 1.
tableFromWeights
    :: (Vector v (a,Word32), Vector v (a,Double), Vector v a, Vector v Word32)
       => v (a, Double) -> CondensedTable v a
{-# INLINE tableFromWeights #-}
tableFromWeights = tableFromProbabilities . normalize . G.filter ((> 0) . snd)
  where
    normalize v
      | G.null v  = pkgError "tableFromWeights" "no positive weights"
      | otherwise = G.map (second (/ s)) v
      where
        -- Explicit fold is to avoid 'Vector v Double' constraint
        s = G.foldl' (flip $ (+) . snd) 0 v

-- | Generate a condensed lookup table from integer weights. Weights
-- should sum to @2^32@ at least approximately. This function will
-- correct small deviations from @2^32@ such as arising from rounding
-- errors. But for large deviations it's likely to product incorrect
-- result with terrible performance.
tableFromIntWeights :: (Vector v (a,Word32), Vector v a, Vector v Word32)
                    => v (a, Word32)
                    -> CondensedTable v a
{-# INLINE tableFromIntWeights #-}
tableFromIntWeights v
  | n == 0    = pkgError "tableFromIntWeights" "empty table"
    -- Single element tables should be treated sepately. Otherwise
    -- they will confuse correctWeights
  | n == 1    = let m = 2^(32::Int) - 1 -- Works for both Word32 & Word64
                in CondensedTable
                   m (G.replicate 256 $ fst $ G.head tbl)
                   m  G.empty
                   m  G.empty
                      G.empty
  | otherwise = CondensedTable
                na aa
                nb bb
                nc cc
                   dd
  where
    -- We must filter out zero-probability outcomes because they may
    -- confuse weight correction algorithm
    tbl   = G.filter ((/=0) . snd) v
    n     = G.length tbl
    -- Corrected table
    table = uncurry G.zip $ id *** correctWeights $ G.unzip tbl
    -- Make condensed table
    mkTable  d =
      G.concatMap (\(x,w) -> G.replicate (fromIntegral $ digit d w) x) table
    len = fromIntegral . G.length
    -- Tables
    aa = mkTable 0
    bb = mkTable 1
    cc = mkTable 2
    dd = mkTable 3
    -- Offsets
    na =       len aa `shiftL` 24
    nb = na + (len bb `shiftL` 16)
    nc = nb + (len cc `shiftL` 8)

-- Calculate N'th digit base 256
digit :: Int -> Word32 -> Word32
digit 0 x =  x `shiftR` 24
digit 1 x = (x `shiftR` 16) .&. 0xff
digit 2 x = (x `shiftR` 8 ) .&. 0xff
digit 3 x =  x .&. 0xff
digit _ _ = pkgError "digit" "the impossible happened!?"
{-# INLINE digit #-}

-- Correct integer weights so they sum up to 2^32. Array of weight
-- should contain at least 2 elements.
correctWeights :: G.Vector v Word32 => v Word32 -> v Word32
{-# INLINE correctWeights #-}
correctWeights v = G.create $ do
  let
    -- Sum of weights
    s = G.foldl' (flip $ (+) . fromIntegral) 0 v :: Int64
    -- Array size
    n = G.length v
  arr <- G.thaw v
  -- On first pass over array adjust only entries which are larger
  -- than `lim'. On second and subsequent passes `lim' is set to 1.
  --
  -- It's possibly to make this algorithm loop endlessly if all
  -- weights are 1 or 0.
  let loop lim i delta
        | delta == 0 = return ()
        | i >= n     = loop 1 0 delta
        | otherwise  = do
            w <- M.read arr i
            case () of
              _| w < lim   -> loop lim (i+1) delta
               | delta < 0 -> M.write arr i (w + 1) >> loop lim (i+1) (delta + 1)
               | otherwise -> M.write arr i (w - 1) >> loop lim (i+1) (delta - 1)
  loop 255 0 (s - 2^(32::Int))
  return arr

-- | Create a lookup table for the Poisson distibution. Note that
-- table construction may have significant cost. For λ < 100 it
-- takes as much time to build table as generation of 1000-30000
-- variates.
tablePoisson :: Double -> CondensedTableU Int
tablePoisson = tableFromProbabilities . make
  where
    make lam
      | lam < 0    = pkgError "tablePoisson" "negative lambda"
      | lam < 22.8 = U.unfoldr unfoldForward (exp (-lam), 0)
      | otherwise  = U.unfoldr unfoldForward (pMax, nMax)
                  ++ U.tail (U.unfoldr unfoldBackward (pMax, nMax))
      where
        -- Number with highest probability and its probability
        --
        -- FIXME: this is not ideal precision-wise. Check if code
        --        from statistics gives better precision.
        nMax = floor lam :: Int
        pMax = exp $ fromIntegral nMax * log lam - lam - logFactorial nMax
        -- Build probability list
        unfoldForward (p,i)
          | p < minP  = Nothing
          | otherwise = Just ( (i,p)
                             , (p * lam / fromIntegral (i+1), i+1)
                             )
        -- Go down
        unfoldBackward (p,i)
          | p < minP  = Nothing
          | otherwise = Just ( (i,p)
                             , (p / lam * fromIntegral i, i-1)
                             )
    -- Minimal representable probability for condensed tables
    minP = 1.1641532182693481e-10 -- 2**(-33)

-- | Create a lookup table for the binomial distribution.
tableBinomial :: Int            -- ^ Number of tries
              -> Double         -- ^ Probability of success
              -> CondensedTableU Int
tableBinomial n p = tableFromProbabilities makeBinom
  where
  makeBinom
    | n <= 0         = pkgError "tableBinomial" "non-positive number of tries"
    | p == 0         = U.singleton (0,1)
    | p == 1         = U.singleton (n,1)
    | p > 0 && p < 1 = U.unfoldrN (n + 1) unfolder ((1-p)^n, 0)
    | otherwise      = pkgError "tableBinomial" "probability is out of range"
    where
      h = p / (1 - p)
      unfolder (t,i) = Just ( (i,t)
                            , (t * (fromIntegral $ n + 1 - i1) * h / fromIntegral i1, i1) )
        where i1 = i + 1

pkgError :: String -> String -> a
pkgError func err =
    error . concat $ ["System.Random.MWC.CondensedTable.", func, ": ", err]

-- $references
--
-- * Wang, J.; Tsang, W. W.; G. Marsaglia (2004), Fast Generation of
--   Discrete Random Variables, /Journal of Statistical Software,
--   American Statistical Association/, vol. 11(i03).
--   <http://ideas.repec.org/a/jss/jstsof/11i03.html>