{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.FFT (
  fft1d
, fft2d
, fft3d
, Direction(..)
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C
import Data.Vector.Image.Color.RGB
import System.IO.Unsafe(unsafePerformIO)

foreign import ccall unsafe "fft1d" c_fft1d :: Ptr (Complex Double) -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "fft2d" c_fft2d :: Ptr (Complex Double) -> CInt -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "fft3d" c_fft3d :: Ptr (Complex Double) -> CInt -> CInt -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "&fftw_free" c_p_freeComplexDouble :: FunPtr(Ptr (Complex Double) -> IO ())

data Direction = Forward | Backword

fft1d :: Vector (Complex Double) -> Direction -> Vector (Complex Double)
fft1d vec dir = unsafePerformIO $ do
  let n = V.length vec
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft1d v (fromIntegral n) 1
      Backword -> c_fft1d v (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 n

fft2d :: Vector (Complex Double) -> Int -> Direction -> Vector (Complex Double)
fft2d vec m dir = unsafePerformIO $ do
  let nm = V.length vec
  let n = nm `div` m
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft2d v (fromIntegral m) (fromIntegral n) 1
      Backword -> c_fft2d v (fromIntegral m) (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 nm

fft3d :: Vector (Complex Double) -> Int -> Int -> Direction -> Vector (Complex Double)
fft3d vec l m dir = unsafePerformIO $ do
  let lmn = V.length vec
  let n = lmn `div` (l*m)
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft3d v (fromIntegral l) (fromIntegral m) (fromIntegral n) 1
      Backword -> c_fft3d v (fromIntegral l) (fromIntegral m) (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 lmn


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}
-- |
-- Module    : Statistics.Distribution
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Type classes for probability distributions

module Statistics.Distribution
    (
      -- * Type classes
      Distribution(..)
    , DiscreteDistr(..)
    , ContDistr(..)
      -- ** Distribution statistics
    , MaybeMean(..)
    , Mean(..)
    , MaybeVariance(..)
    , Variance(..)
    , MaybeEntropy(..)
    , Entropy(..)
    , FromSample(..)
      -- ** Random number generation
    , ContGen(..)
    , DiscreteGen(..)
    , genContinuous
    , genContinous
      -- * Helper functions
    , findRoot
    , sumProbabilities
    ) where

import Control.Applicative ((<$>), Applicative(..))
import Control.Monad.Primitive (PrimMonad,PrimState)
import Prelude hiding (sum)
import Statistics.Function (square)
import Statistics.Sample.Internal (sum)
import System.Random.MWC (Gen, uniform)
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G

-- | Type class common to all distributions. Only c.d.f. could be
-- defined for both discrete and continuous distributions.
class Distribution d where
    -- | Cumulative distribution function.  The probability that a
    -- random variable /X/ is less or equal than /x/,
    -- i.e. P(/X/≤/x/). Cumulative should be defined for
    -- infinities as well:
    --
    -- > cumulative d +∞ = 1
    -- > cumulative d -∞ = 0
    cumulative :: d -> Double -> Double

    -- | One's complement of cumulative distribution:
    --
    -- > complCumulative d x = 1 - cumulative d x
    --
    -- It's useful when one is interested in P(/X/>/x/) and
    -- expression on the right side begin to lose precision. This
    -- function have default implementation but implementors are
    -- encouraged to provide more precise implementation.
    complCumulative :: d -> Double -> Double
    complCumulative d x = 1 - cumulative d x

-- | Discrete probability distribution.
class Distribution  d => DiscreteDistr d where
    -- | Probability of n-th outcome.
    probability :: d -> Int -> Double
    probability d = exp . logProbability d

    -- | Logarithm of probability of n-th outcome
    logProbability :: d -> Int -> Double
    logProbability d = log . probability d

-- | Continuous probability distribution.
--
--   Minimal complete definition is 'quantile' and either 'density' or
--   'logDensity'.
class Distribution d => ContDistr d where
    -- | Probability density function. Probability that random
    -- variable /X/ lies in the infinitesimal interval
    -- [/x/,/x+/δ/x/) equal to /density(x)/⋅δ/x/
    density :: d -> Double -> Double
    density d = exp . logDensity d

    -- | Inverse of the cumulative distribution function. The value
    -- /x/ for which P(/X/≤/x/) = /p/. If probability is outside
    -- of [0,1] range function should call 'error'
    quantile :: d -> Double -> Double

    -- | 1-complement of @quantile@:
    --
    -- > complQuantile x ≡ quantile (1 - x)
    complQuantile :: d -> Double -> Double
    complQuantile d x = quantile d (1 - x)

    -- | Natural logarithm of density.
    logDensity :: d -> Double -> Double
    logDensity d = log . density d

-- | Type class for distributions with mean. 'maybeMean' should return
--   'Nothing' if it's undefined for current value of data
class Distribution d => MaybeMean d where
    maybeMean :: d -> Maybe Double

-- | Type class for distributions with mean. If a distribution has
--   finite mean for all valid values of parameters it should be
--   instance of this type class.
class MaybeMean d => Mean d where
    mean :: d -> Double

-- | Type class for distributions with variance. If variance is
--   undefined for some parameter values both 'maybeVariance' and
--   'maybeStdDev' should return Nothing.
--
--   Minimal complete definition is 'maybeVariance' or 'maybeStdDev'
class MaybeMean d => MaybeVariance d where
    maybeVariance :: d -> Maybe Double
    maybeVariance d = (*) <$> x <*> x where x = maybeStdDev d
    maybeStdDev   :: d -> Maybe Double
    maybeStdDev = fmap sqrt . maybeVariance

-- | Type class for distributions with variance. If distribution have
--   finite variance for all valid parameter values it should be
--   instance of this type class.
--
--   Minimal complete definition is 'variance' or 'stdDev'
class (Mean d, MaybeVariance d) => Variance d where
    variance :: d -> Double
    variance d = square (stdDev d)
    stdDev   :: d -> Double
    stdDev = sqrt . variance

-- | Type class for distributions with entropy, meaning Shannon entropy
--   in the case of a discrete distribution, or differential entropy in the
--   case of a continuous one.  'maybeEntropy' should return 'Nothing' if
--   entropy is undefined for the chosen parameter values.
class (Distribution d) => MaybeEntropy d where
  -- | Returns the entropy of a distribution, in nats, if such is defined.
  maybeEntropy :: d -> Maybe Double

-- | Type class for distributions with entropy, meaning Shannon
--   entropy in the case of a discrete distribution, or differential
--   entropy in the case of a continuous one.  If the distribution has
--   well-defined entropy for all valid parameter values then it
--   should be an instance of this type class.
class (MaybeEntropy d) => Entropy d where
  -- | Returns the entropy of a distribution, in nats.
  entropy :: d -> Double

-- | Generate discrete random variates which have given
--   distribution.
class Distribution d => ContGen d where
  genContVar :: PrimMonad m => d -> Gen (PrimState m) -> m Double

-- | Generate discrete random variates which have given
--   distribution. 'ContGen' is superclass because it's always possible
--   to generate real-valued variates from integer values
class (DiscreteDistr d, ContGen d) => DiscreteGen d where
  genDiscreteVar :: PrimMonad m => d -> Gen (PrimState m) -> m Int

-- | Estimate distribution from sample. First parameter in sample is
--   distribution type and second is element type.
class FromSample d a where
  -- | Estimate distribution from sample. Returns nothing is there's
  --   not enough data to estimate or sample clearly doesn't come from
  --   distribution in question. For example if there's negative
  --   samples in exponential distribution.
  fromSample :: G.Vector v a => v a -> Maybe d

-- | Generate variates from continuous distribution using inverse
--   transform rule.
genContinuous :: (ContDistr d, PrimMonad m) => d -> Gen (PrimState m) -> m Double
genContinuous d gen = do
  x <- uniform gen
  return $! quantile d x

-- | Backwards compatibility with genContinuous.
genContinous :: (ContDistr d, PrimMonad m) => d -> Gen (PrimState m) -> m Double
genContinous = genContinuous
{-# DEPRECATED genContinous "Use genContinuous" #-}

data P = P {-# UNPACK #-} !Double {-# UNPACK #-} !Double

-- | Approximate the value of /X/ for which P(/x/>/X/)=/p/.
--
-- This method uses a combination of Newton-Raphson iteration and
-- bisection with the given guess as a starting point.  The upper and
-- lower bounds specify the interval in which the probability
-- distribution reaches the value /p/.
findRoot :: ContDistr d =>
            d                   -- ^ Distribution
         -> Double              -- ^ Probability /p/
         -> Double              -- ^ Initial guess
         -> Double              -- ^ Lower bound on interval
         -> Double              -- ^ Upper bound on interval
         -> Double
findRoot d prob = loop 0 1
  where
    loop !(i::Int) !dx !x !lo !hi
      | abs dx <= accuracy || i >= maxIters = x
      | otherwise                           = loop (i+1) dx'' x'' lo' hi'
      where
        err                   = cumulative d x - prob
        P lo' hi' | err < 0   = P x hi
                  | otherwise = P lo x
        pdf                   = density d x
        P dx' x' | pdf /= 0   = P (err / pdf) (x - dx)
                 | otherwise  = P dx x
        P dx'' x''
            | x' < lo' || x' > hi' || pdf == 0 = let y = (lo' + hi') / 2
                                                 in  P (y-x) y
            | otherwise                        = P dx' x'
    accuracy = 1e-15
    maxIters = 150

-- | Sum probabilities in inclusive interval.
sumProbabilities :: DiscreteDistr d => d -> Int -> Int -> Double
sumProbabilities d low hi =
  -- Return value is forced to be less than 1 to guard against roundoff errors.
  -- ATTENTION! this check should be removed for testing or it could mask bugs.
  min 1 . sum . U.map (probability d) $ U.enumFromTo low hi


{-# LANGUAGE BangPatterns,
             ScopedTypeVariables,
             RecordWildCards,
             FlexibleContexts,
             TypeFamilies #-}

-- |
-- Module       : AI.HNN.FF.Network
-- Copyright    : (c) 2012 Alp Mestanogullari
-- License      : BSD3
-- Maintainer   : alpmestan@gmail.com
-- Stability    : experimental
-- Portability  : GHC
--
-- An implementation of feed-forward neural networks in pure Haskell.
--
-- It uses weight matrices between each layer to represent the connections between neurons from
-- a layer to the next and exports only the useful bits for a user of the library.
--
-- Here is an example of using this module to create a feed-forward neural network with 2 inputs,
-- 2 neurons in a hidden layer and one neuron in the output layer, with random weights, and compute
-- its output for [1,2] using the sigmoid function for activation for all the neurons.
--
-- > import AI.HNN.FF.Network
-- > import Numeric.LinearAlgebra
-- >
-- > main = do
-- >   n <- createNetwork 2 [2] 1 :: IO (Network Double)
-- >   print $ output n sigmoid (fromList [1, 1])
--
-- /Note/: Here, I create a @Network Double@, but you can replace 'Double' with any number type
-- that implements the appropriate typeclasses you can see in the signatures of this module.
-- Having your number type implement the @Floating@ typeclass too is a good idea, since that's what most of the
-- common activation functions require.
--
-- /Note 2/: You can also give some precise weights to initialize the neural network with, with
-- 'fromWeightMatrices'. You can also restore a neural network you had saved using 'loadNetwork'.
--
-- Here is an example of how to train a neural network to learn the XOR function.
-- ( for reference: XOR(0, 0) = 0, XOR(0, 1) = 1, XOR(1, 0) = 1, XOR(1, 1) = 0 )
--
-- First, let's import hnn's feedforward neural net module, and hmatrix's vector types.
--
-- > import AI.HNN.FF.Network
-- > import Numeric.LinearAlgebra
--
-- Now, we will specify our training set (what the net should try to learn).
--
-- > samples :: Samples Double
-- > samples = [ fromList [0, 0] --> fromList [0]
-- >           , fromList [0, 1] --> fromList [1]
-- >           , fromList [1, 0] --> fromList [1]
-- >           , fromList [1, 1] --> fromList [0]
-- >           ]
--
-- You can see that this is basically a list of pairs of vectors, the first vector being
-- the input given to the network, the second one being the expected output. Of course,
-- this imply working on a neural network with 2 inputs, and a single neuron on the output layer. Then,
-- let's create one!
--
-- > main = do
-- >   net <- createNetwork 2 [2] 1
--
-- You may have noticed we haven't specified a signature this time, unlike in the earlier snippet.
-- Since we gave a signature to samples, specifying we're working with 'Double' numbers, and since
-- we are going to tie 'net' and 'samples' by a call to a learning function, GHC will gladly figure out
-- that 'net' is working with 'Double'.
--
-- Now, it's time to train our champion. But first, let's see how bad he is now. The weights are most likely
-- not close to those that will give a good result for simulating XOR. Let's compute the output of the net on
-- the input vectors of our samples, using 'tanh' as the activation function.
--
-- >   mapM_ (print . output net tanh . fst) samples
--
-- Ok, you've tested this, and it gives terrible results. Let's fix this by letting 'trainNTimes' teach our neural net
-- how to behave. Since we're using 'tanh' as our activation function, we will tell it to the training function,
-- and also specify its derivative.
--
-- >   let smartNet = trainNTimes 1000 0.8 tanh tanh' net samples
--
-- So, this tiny piece of code will run the backpropagation algorithm on the samples 1000 times, with a learning rate
-- of 0.8. The learning rate is basically how strongly we should modify the weights when we try to correct the error the net makes
-- on our samples. The bigger it is, the more the weights are going to change significantly. Depending on the case, it can be good,
-- but sometimes it can make the backprop algorithm oscillate around good weight values without actually getting to them.
-- You usually want to test several values and see which ones get you the nicest neural net, which generalizes well to samples
-- that are not in the training set while giving decent results on the training set.
--
-- Now, let's see how that worked out for us:
--
-- >   mapM_ (print . output smartNet tanh . fst) samples
--
-- You could even save that neural network's weights to a file, so that you don't need to train it again in the future, using 'saveNetwork':
--
-- >   saveNetwork "smartNet.nn" smartNet
--
-- Please note that 'saveNetwork' is just a wrapper around zlib compression + serialization using the binary package.
-- AI.HNN.FF.Network also provides a 'Data.Binary.Binary' instance for 'Network', which means you can also simply use
-- 'Data.Binary.encode' and 'Data.Binary.decode' to have your own saving/restoring routines, or to simply get a bytestring
-- we can send over the network, for example.
--
-- Here's a run of the program we described on my machine (with the timing): first set of
-- fromList's is the output of the initial neural network, the second one is the output of
-- 'smartNet' :-)
--
-- > fromList [0.574915179613429]
-- > fromList [0.767589097192215]
-- > fromList [0.7277396607146663]
-- > fromList [0.8227114080561128]
-- > ------------------
-- > fromList [6.763498312099933e-2]
-- > fromList [0.9775186355284375]
-- > fromList [0.9350823296850516]
-- > fromList [-4.400205702560454e-2]
-- >
-- > real    0m0.365s
-- > user    0m0.072s
-- > sys     0m0.016s
--
-- Rejoyce! Feel free to play around with the library and report any bug, feature request and whatnot to us on
-- our github repository <https://github.com/alpmestan/hnn/issues> using the appropriate tags. Also, you can
-- see the simple program we studied here with pretty colors at <https://github.com/alpmestan/hnn/blob/master/examples/ff/xor.hs>
-- and other ones at <https://github.com/alpmestan/hnn/tree/master/examples/ff>.

module AI.HNN.FF.Network
    (
    -- * Types
      Network(..)
    , ActivationFunction
    , ActivationFunctionDerivative
    , Sample
    , Samples
    , (-->)

    -- * Creating a neural network
    , createNetwork
    , fromWeightMatrices

    -- * Computing a neural network's output
    , output
    , tanh
    , tanh'
    , sigmoid
    , sigmoid'

    -- * Training a neural network
    , trainUntil
    , trainNTimes
    , trainUntilErrorBelow
    , quadError

    -- * Loading and saving a neural network
    , loadNetwork
    , saveNetwork
    ) where

import Codec.Compression.Zlib     (compress, decompress)
import Data.Binary                (Binary(..), encode, decode)
import Data.List                  (foldl')
import Foreign.Storable           (Storable)
import qualified Data.ByteString.Lazy  as B
import qualified Data.Vector           as V

import System.Random.MWC
import Numeric.LinearAlgebra.HMatrix hiding (corr)
import Data.Functor ((<$>))

-- | Our feed-forward neural network type. Note the 'Binary' instance, which means you can use
--   'encode' and 'decode' in case you need to serialize your neural nets somewhere else than
--   in a file (e.g over the network)
newtype Network a = Network
                 { matrices   :: V.Vector (Matrix a) -- ^ the weight matrices
                 } deriving Show

instance (Element a, Binary a) => Binary (Network a) where
  put (Network ms) = put . V.toList $ ms
  get = (Network . V.fromList) `fmap` get

-- | The type of an activation function, mostly used for clarity in signatures
type ActivationFunction a = a -> a

-- | The type of an activation function's derivative, mostly used for clarity in signatures
type ActivationFunctionDerivative a = a -> a

-- | The following creates a neural network with 'n' inputs and if 'l' is [n1, n2, ...]
--   the net will have n1 neurons on the first layer, n2 neurons on the second, and so on
--   ending with k neurons on the output layer, with random weight matrices as a courtesy of
-- 'System.Random.MWC.uniformVector'.
--
-- > createNetwork n l k
createNetwork :: (Variate a, Storable a) => Int -> [Int] -> Int -> IO (Network a)
createNetwork nInputs hiddens nOutputs =
  fmap Network $ withSystemRandom . asGenST $ \gen -> go gen dimensions V.empty
  where
        go _ [] !ms         = return ms
        go gen ((!n,!m):ds) ms = do
          !mat <- randomMat n m gen
          go gen ds (ms `V.snoc` mat)
        randomMat n m g = reshape m `fmap` uniformVector g (n*m)
        dimensions      = zip (hiddens ++ [nOutputs]) $
                              (nInputs+1 : hiddens)
{-# INLINE createNetwork #-}

-- | Creates a neural network with exactly the weight matrices given as input here.
--   We don't check that the numbers of rows/columns are compatible, etc.
fromWeightMatrices :: Storable a => V.Vector (Matrix a) -> Network a
fromWeightMatrices ws = Network ws
{-# INLINE fromWeightMatrices #-}

-- The `join [input, 1]' trick  below is a courtesy of Alberto Ruiz
-- <http://dis.um.es/~alberto/>. Per his words:
--
-- "The idea is that the constant input in the first layer can be automatically transferred to the following layers
-- by the learning algorithm (by setting the weights of a neuron to 1,0,0,0,...). This allows for a simpler
-- implementation and in my experiments those networks are able to easily solve non linearly separable problems."

-- | Computes the output of the network on the given input vector with the given activation function
output :: (Floating (Vector a), Numeric a, Storable a, Num (Vector a)) => Network a -> ActivationFunction a -> Vector a -> Vector a
output (Network{..}) act input = V.foldl' f (vjoin [input, 1]) matrices
  where f !inp m = cmap act $ m #> inp
{-# INLINE output #-}

-- | Computes and keeps the output of all the layers of the neural network with the given activation function
outputs :: (Floating (Vector a), Numeric a, Storable a, Num (Vector a)) => Network a -> ActivationFunction a -> Vector a -> V.Vector (Vector a)
outputs (Network{..}) act input = V.scanl f (vjoin [input, 1]) matrices
  where f !inp m = cmap act $ m #> inp
{-# INLINE outputs #-}

deltas :: (Floating (Vector a), Floating a, Numeric a, Container Vector a, Num (Vector a)) => Network a -> ActivationFunctionDerivative a -> V.Vector (Vector a) -> Vector a -> V.Vector (Matrix a)
deltas (Network{..}) act' os expected = V.zipWith outer (V.tail ds) (V.init os)
  where !dl = (V.last os - expected) * (deriv $ V.last os)
        !ds = V.scanr f dl (V.zip os matrices)
        f (!o, m) !del = deriv o * (tr m #> del)
        deriv = cmap act'
{-# INLINE deltas #-}

updateNetwork :: (Floating (Vector a), Floating a, Numeric a, Storable a, Num (Vector a), Container Vector a) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Sample a -> Network a
updateNetwork alpha act act' n@(Network{..}) (input, expectedOutput) = Network $ V.zipWith (+) matrices corr
    where !xs = outputs n act input
          !ds = deltas n act' xs expectedOutput
          !corr = V.map (scale (-alpha)) ds
{-# INLINE updateNetwork #-}

-- | Input vector and expected output vector
type Sample a = (Vector a, Vector a)

-- | List of 'Sample's
type Samples a = [Sample a]

-- | Handy operator to describe your learning set, avoiding unnecessary parentheses. It's just a synonym for '(,)'.
--   Generally you'll load your learning set from a file, a database or something like that, but it can be nice for
--   quickly playing with hnn or for simple problems where you manually specify your learning set.
--   That is, instead of writing:
--
-- > samples :: Samples Double
-- > samples = [ (fromList [0, 0], fromList [0])
-- >           , (fromList [0, 1], fromList [1])
-- >           , (fromList [1, 0], fromList [1])
-- >           , (fromList [1, 1], fromList [0])
-- >           ]
--
--   You can write:
--
-- > samples :: Samples Double
-- > samples = [ fromList [0, 0] --> fromList [0]
-- >           , fromList [0, 1] --> fromList [1]
-- >           , fromList [1, 0] --> fromList [1]
-- >           , fromList [1, 1] --> fromList [0]
-- >           ]
(-->) :: Vector a -> Vector a -> Sample a
(-->) = (,)

backpropOnce :: (Floating (Vector a), Floating a, Numeric a, Num (Vector a), Container Vector a) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
backpropOnce rate act act' n samples = foldl' (updateNetwork rate act act') n samples
{-# INLINE backpropOnce #-}

-- | Generic training function.
--
-- The first argument is a predicate that will tell the backpropagation algorithm when to stop.
-- The first argument to the predicate is the epoch, i.e the number of times the backprop has been
-- executed on the samples. The second argument is /the current network/, and the third is the list of samples.
-- You can thus combine these arguments to create your own criterion.
--
-- For example, if you want to stop learning either when the network's quadratic error on the samples,
-- using the tanh function, is below 0.01, or after 1000 epochs, whichever comes first, you could
-- use the following predicate:
--
-- > pred epochs net samples = if epochs == 1000 then True else quadError tanh net samples < 0.01
--
-- You could even use 'Debug.Trace.trace' to print the error, to see how the error evolves while it's learning,
-- or redirect this to a file from your shell in order to generate a pretty graphics and what not.
--
-- The second argument (after the predicate) is the learning rate. Then come the activation function you want,
-- its derivative, the initial neural network, and your training set.
-- Note that we provide 'trainNTimes' and 'trainUntilErrorBelow' for common use cases.
trainUntil :: (Floating (Vector a), Floating a, Numeric a, Num (Vector a), Container Vector a) => (Int -> Network a -> Samples a -> Bool) -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
trainUntil pr learningRate act act' net samples = go net 0
  where go n !k | pr k n samples = n
                | otherwise      = case backpropOnce learningRate act act' n samples of
                                    n' -> go n' (k+1)
{-# INLINE trainUntil #-}

-- | Trains the neural network with backpropagation the number of times specified by the 'Int' argument,
-- using the given learning rate (second argument).
trainNTimes :: (Floating (Vector a), Floating a, Numeric a, Num (Vector a), Container Vector a) => Int -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
trainNTimes n = trainUntil (\k _ _ -> k > n)
{-# INLINE trainNTimes #-}

-- | Quadratic error on the given training set using the given activation function. Useful to create
-- your own predicates for 'trainUntil'.
quadError :: (Floating (Vector a), Floating a, Fractional (RealOf a), Normed (Vector a), Numeric a) => ActivationFunction a -> Network a -> Samples a -> RealOf a
quadError act net samples = realToFrac $ foldl' (\err (inp, out) -> err + (norm_2 $ output net act inp - out)) 0 samples
{-# INLINE quadError #-}

-- | Trains the neural network until the quadratic error ('quadError') comes below the given value (first argument),
-- using the given learning rate (second argument).
--
-- /Note/: this can loop pretty much forever when you're using a bad architecture for the problem, or inappropriate activation
-- functions.
trainUntilErrorBelow :: (Floating (Vector a), Floating a, Numeric a, Normed (Vector a), Ord a, Container Vector a, Num (RealOf a), a ~ RealOf a, Show a) => a -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> Network a -> Samples a -> Network a
trainUntilErrorBelow x rate act = trainUntil (\_ n s -> quadError act n s < x) rate act
{-# INLINE trainUntilErrorBelow #-}

-- | The sigmoid function:  1 / (1 + exp (-x))
sigmoid :: Floating a => a -> a
sigmoid !x = 1 / (1 + exp (-x))
{-# INLINE sigmoid #-}

-- | Derivative of the sigmoid function: sigmoid x * (1 - sigmoid x)
sigmoid' :: Floating a => a -> a
sigmoid' !x = case sigmoid x of
  s -> s * (1 - s)
{-# INLINE sigmoid' #-}

-- | Derivative of the 'tanh' function from the Prelude.
tanh' :: Floating a => a -> a
tanh' !x = case tanh x of
  s -> 1 - s**2
{-# INLINE tanh' #-}

-- | Loading a neural network from a file (uses zlib compression on top of serialization using the binary package).
--   Will throw an exception if the file isn't there.
loadNetwork :: (Storable a, Element a, Binary a) => FilePath -> IO (Network a)
loadNetwork fp = decode . decompress <$> B.readFile fp
{-# INLINE loadNetwork #-}

-- | Saving a neural network to a file (uses zlib compression on top of serialization using the binary package).
saveNetwork :: (Storable a, Element a, Binary a) => FilePath -> Network a -> IO ()
saveNetwork fp = B.writeFile fp . compress . encode
{-# INLINE saveNetwork #-}


{-|
Module: MachineLearning.NeuralNetwork.Topology
Description: Neural Network's Topology
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Neural Network's Topology
-}

module MachineLearning.NeuralNetwork.Topology
(
  Topology
  , LossFunc(..)
  , makeTopology
  , loss
  , propagateForward
  , propagateBackward
  , numberOutputs
  , initializeTheta
  , initializeThetaIO
  , initializeThetaM
  , flatten
  , unflatten
)

where

import Control.Monad (zipWithM)
import Data.List (foldl')
import qualified Control.Monad.Random as RndM
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Utils (listOfTuplesToList)
import MachineLearning.NeuralNetwork.Layer (Layer(..), Cache(..))
import MachineLearning.NeuralNetwork.Regularization (Regularization, forwardReg, backwardReg)

-- | Loss function's type.
-- Takes x, weights and y.
type LossFunc = Matrix -> Matrix -> R

-- | Neural network topology has at least 2 elements: numver of input and number of outputs.
-- And sizes of hidden layers between 2 elements.
data Topology = Topology [(Int, Int)] [Layer] LossFunc

-- | Makes Neural Network's Topology.
-- Takes number of inputs, list of hidden layers, output layer and loss function.
makeTopology :: Int -> [Layer] -> Layer -> LossFunc -> Topology
makeTopology nInputs hiddenLayers outputLayer lossFunc =
  let layers = hiddenLayers ++ [outputLayer]
      layerSizes = nInputs : (map lUnits layers)
      sizes = getThetaSizes layerSizes
  in Topology sizes layers lossFunc

-- | Calculates loss for the given topology.
-- Takes topology, regularization, x, weights, y.
loss :: Topology -> Regularization -> Matrix -> [(Matrix, Matrix)] -> Matrix -> R
loss (Topology _ _ lf) reg x weights y =
  let lossValue = lf x y
      regValue = forwardReg reg weights
  in lossValue + regValue

-- | Implementation of forward propagation algorithm.
propagateForward :: Topology -> Matrix -> [(Matrix, Matrix)] -> (Matrix, [Cache])
propagateForward (Topology _ layers _) x thetaList = foldl' f (x, []) $ zip thetaList layers
  where f (a, cs) (theta, hl) =
          let (a', cache) = forwardPass hl a theta
          in (a', cache:cs)

-- | Makes one forward step for the given layer.
forwardPass :: Layer -> Matrix -> (Matrix, Matrix) -> (Matrix, Cache)
forwardPass layer a (b, w) = (a', Cache z a w)
  where z = lForward layer a b w
        a' = lActivation layer z

-- | Implementation of backward propagation algorithm.
propagateBackward :: Topology -> Regularization -> Matrix -> [Cache] -> Matrix -> [(Matrix, Matrix)]
propagateBackward (Topology _ layers _) reg scores (cache:cacheList) y = gradientList
  where cache' = Cache scores (cacheX cache) (cacheW cache)
        cacheList' = cache':cacheList
        gradientList = snd $ foldl' f (y, []) $ zip cacheList' $ reverse layers
        f (da, grads) (cache, hl) =
          let (da', db, dw) = backwardPass hl reg da cache
          in (da', (db, dw):grads)

-- | Makes one backward step for the given layer.
backwardPass :: Layer -> Regularization -> Matrix -> Cache -> (Matrix, Matrix, Matrix)
backwardPass layer reg da cache = (da', db, dw')
  where delta = lActivationGradient layer (cacheZ cache) da
        (da', db, dw) = lBackward layer delta cache
        dw' = dw + (backwardReg reg (cacheW cache))

-- | Returns number of outputs of the given topology.
numberOutputs :: Topology -> Int
numberOutputs (Topology nnt _ _) = fst $ last nnt

-- | Returns dimensions of weight matrices for given neural network topology
getThetaSizes :: [Int] -> [(Int, Int)]
getThetaSizes nn = zipWith (\r c -> (r, c)) (tail nn) nn

-- | Create and initialize weights vector with random values
-- for given neural network topology.
-- Takes a seed to initialize generator of random numbers as a first parameter.
initializeTheta :: Int -> Topology -> Vector
initializeTheta seed topology = RndM.evalRand (initializeThetaM topology) gen
  where gen = RndM.mkStdGen seed

-- | Create and initialize weights vector with random values
-- for given neural network topology inside IO Monad.
initializeThetaIO :: Topology -> IO Vector
initializeThetaIO = RndM.evalRandIO . initializeThetaM

-- | Create and initialize weights vector with random values
-- for given neural network topology inside RandomMonad.
initializeThetaM :: RndM.RandomGen g => Topology -> RndM.Rand g Vector
initializeThetaM topology = flatten <$> initializeThetaListM topology

-- | Create and initialize list of weights matrices with random values
-- for given neural network topology.
initializeThetaListM :: RndM.RandomGen g => Topology -> RndM.Rand g [(Matrix, Matrix)]
initializeThetaListM (Topology sizes layers _) = zipWithM lInitializeThetaM layers sizes

-- | Flatten list of matrices into vector.
flatten :: [(Matrix, Matrix)] -> Vector
flatten ms = V.concat $ map LA.flatten $ listOfTuplesToList ms

-- | Unflatten vector into list of matrices for given neural network topology.
unflatten :: Topology -> Vector -> [(Matrix, Matrix)]
unflatten (Topology sizes _ _) v =
  let offsets = reverse $ foldl' (\os (r, c) -> (r+r*c + head os):os) [0] (init sizes)
      ms = zipWith (\o (r, c) -> (LA.reshape r (slice o r), LA.reshape c (slice (o+r) (r*c)))) offsets sizes
      slice o n = V.slice o n v
  in ms


{-# LANGUAGE BangPatterns #-}
module STCPinwheel 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: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/STCPinwheel"
      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
      periodEnvelope = periodR2 ^ 2 * 2
      maxScale = sqrt periodEnvelope
  removePathForcibly folderPath
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath

  hist <-
    if flag
      then do printCurrentTime "read from files..."
              decodeFile histFilePath
      else do printCurrentTime "Start computing coefficients..."
              initialise []
              devs <- M.mapM device deviceIDs
              ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
              ptxs <- M.mapM createTargetFromContext ctxs
              runMonteCarloFourierCoefficientsGPU
                deviceIDs
                numThread
                1000000
                500000
                thetaSigma
                scaleSigma
                0
                s
                tau
                deltaT
                phiFreqs
                rhoFreqs
                thetaFreqs
                scaleFreqs
                maxScale
                histFilePath
      -- else sampleCartesian
      --        histFilePath
      --        folderPath
      --        ptxs
      --        numPoints
      --        periodEnvelope
      --        delta
      --        numOrientation
      --        initScale
      --        thetaSigma
      --        tau
      --        threshold
      --        s
      --        phiFreq
      --        rhoFreq
      --        thetaFreq
      --        scaleFreq
  printCurrentTime "Done"
  printCurrentTime "Start Convloution.."
  plan <-
    makePlanDiscrete
      folderPath
      emptyPlan
      numPointsRecon
      numPointsRecon
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
      (2 * phiFreq + 1)
      (2 * rhoFreq + 1)
  let coefficients = getNormalizedHistogramArr $ hist
  -- coefficients <- hollowCoefficients plan radius periodEnvelope coefficients'
      -- coefficients = getNormalizedHistogramArr $ hist
        -- normalizeFreqArr' std phiFreqs rhoFreqs .
  harmonicsArray <-
    createHarmonics
      numR2Freq
      phiFreq
      rhoFreq
      thetaFreq
      scaleFreq
      (-s)
      periodR2
      periodEnvelope
      coefficients
      -- harmonicsArray1 =
      --   createHarmonics
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-1.25)
      --     periodR2
      --     periodEnvelope
      --     coefficients
  let !initialSourceDistribution =
        computeInitialDistributionFourierPinwheel
          numR2Freq
          periodR2
          periodEnvelope
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          initSource
      !initialSinkDistribution =
        computeInitialDistributionFourierPinwheel
          numR2Freq
          periodR2
          periodEnvelope
          phiFreq
          rhoFreq
          thetaFreq
          scaleFreq
          initSink
      -- harmonicsArray' =
      --   centerHollow numR2Freq $
      --   pinwheelFourierCoefficientsAnatical
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-s)
      --     periodR2
      --     periodEnvelope
      -- harmonicsArray1' =
      --   -- L.map (centerHollowVector numR2Freq) $
      --   pinwheelFourierCoefficientsAnaticalList1
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-s)
      --     periodR2
      --     periodEnvelope
      -- harmonicsArray2' =
      --   -- L.map (centerHollowVector numR2Freq) $
      --   pinwheelFourierCoefficientsAnaticalList2
      --     numR2Freq
      --     phiFreq
      --     rhoFreq
      --     thetaFreq
      --     scaleFreq
      --     (-s)
      --     periodR2
      --     periodEnvelope
      -- initialSourceDistribution =
      --   computeInitialDistributionFull'
      --     numR2Freq
      --     periodR2
      --     phiFreq
      --     rhoFreq
      --     initSource
      -- initialSourceDistribution =
      --   computeInitialDistributionPowerMethodPinwheelBasis'
      --     numR2Freq
      --     (-s)
      --     periodR2
      --     phiFreq
      --     rhoFreq
      --     initSource
      -- initialSinkDistribution =
      --   computeInitialDistributionFull'
      --     numR2Freq
      --     periodR2
      --     phiFreq
      --     rhoFreq
      --     initSink
  -- source <-
  --   convoluvePinhweelBasisTest'
  --     coefficients
  --     harmonicsArray1'
  --     harmonicsArray2'
  --     initialSourceDistribution
  -- sink <-
  --   timeReversal' <$>
  --   convoluvePinhweelBasisTest'
  --     coefficients
  --     harmonicsArray1'
  --     harmonicsArray2'
  --     initialSinkDistribution
  -- let source =
  --       convoluvePinhweelBasis'
  --         coefficients
  --         harmonicsArray'
  --         initialSourceDistribution
  --     sink =
  --       timeReversal' $
  --       convoluvePinhweelBasis'
  --         coefficients
  --         harmonicsArray'
  --         initialSinkDistribution
  --     sourceMat =
  --       A.transpose . A.use . fromDFTArray . getDFTArrayVector $ source
  --     sinkMat = A.transpose . A.use . fromDFTArray . getDFTArrayVector $ sink-
  source <- convolve harmonicsArray initialSourceDistribution
  -- source1 <- convolve harmonicsArray1 initialSourceDistribution
  let numLogPolarFreq = ((2 * scaleFreq + 1) * (2 * thetaFreq + 1))
      -- sourceMat = A.transpose . toMatrixAcc $ source
      -- sourceMat1 = A.transpose . toMatrixAcc $ source1
      -- sourceMat2 =
      --   A.transpose . toMatrixAcc $
      --   parZipWithFPArray (VS.zipWith (-)) source source1
  -- sourceR2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sourceMat
  -- plotImageRepa (folderPath </> "SourcePower.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   sourceR2
  sourceR2 <- plotFPArray plan (folderPath </> "SourcePower.png") source
  -- source1R2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sourceMat1
  -- plotImageRepa (folderPath </> "SourcePower1.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   source1R2
  -- source2R2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sourceMat2
  -- plotImageRepa (folderPath </> "SourcePower2.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   source2R2
  sink <- convolve harmonicsArray initialSinkDistribution
  let sinkMat = A.transpose . toMatrixAcc $ sink
  -- completion <- completionField' plan source (timeReversal' sink)
  -- let completionMat = createCuMat . getDFTArrayVector $ completion
  -- print completionMat
  sinkR2 <- plotFPArray plan (folderPath </> "SinkPower.png") sink
  -- sinkR2 <-
  --   computeFourierSeriesR2StreamAcc
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     numLogPolarFreq
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     sinkMat
  -- plotImageRepa (folderPath </> "SinkPower.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   VU.map sqrt . toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
  --   sinkR2
  -- completionR2 <-
  --   computeFourierSeriesR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freq
  --     numPointsRecon
  --     periodR2
  --     deltaRecon
  --     numBatchR2
  --     [transposeCuMat completionMat]
  printCurrentTime "Start ploting.."
  -- plotDFTArrayPower
  --   (folderPath </> "CoefficientsPower.png")
  --   numR2Freq
  --   numR2Freq
  --   source
  -- plotImageRepaComplex (folderPath </> "Coefficients.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numR2Freq :. numR2Freq) .
  --   VU.convert . L.foldl1' (VS.zipWith (+)) . getDFTArrayVector $
  --   source
  -- plotImageRepaComplex (folderPath </> "Source.png") .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
  --   toUnboxed . sumS . rotate3D $
  --   sourceR2
  completionR2 <-
    completionFieldRepa
      plan
      (R.reshape
         (Z :. (L.length scaleFreqs) :. (L.length thetaFreqs) :. numPointsRecon :.
          numPointsRecon) $
       sourceR2)
      (timeReversalRepa thetaFreqs .
       R.reshape
         (Z :. (L.length scaleFreqs) :. (L.length thetaFreqs) :. numPointsRecon :.
          numPointsRecon) $
       sinkR2)
  plotImageRepa (folderPath </> "CompletionPower.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPointsRecon :. numPointsRecon) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  printCurrentTime "Done"
  -- sparse
  let harmonicsArraySparse =
        computeHarmonicsArraySparse
          numPointsRecon
          deltaRecon
          numPointsRecon
          deltaRecon
          phiFreqs
          rhoFreqs
          thetaFreqs
          scaleFreqs
          -- (2 * pi)
          (log (periodR2 / 4))
          64
          s
      init =
        computeInitialDistribution'
          numPointsRecon
          numPointsRecon
          phiFreqs
          rhoFreqs
          0
          initSource
  harmonicsArrayDFT <-
    fmap (listArray (bounds harmonicsArraySparse)) .
    dftExecuteBatchP
      plan
      (DFTPlanID DFT1DG [numPointsRecon, numPointsRecon] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    harmonicsArraySparse
  sourceSparse <- convolve' Source plan coefficients harmonicsArrayDFT init
  plotDFTArrayPower
    (folderPath </> "SourcePower_Sparse.png")
    numPointsRecon
    numPointsRecon
    sourceSparse


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : First attempt at approximate counting
-- | Creator: Xiao Ling
-- | Created: 12/08/2015
-- | TODO   : test standard deviation of alpha, beta, and final version
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module Morris (

    morris
  , morris'

  , morris1
  , morris1'

  ) where

import Control.Monad.Random
import Control.Monad.Random.Class

import Data.Conduit
import Data.List.Split
import qualified Data.Conduit.List as Cl

import Core
import Statistics

{-----------------------------------------------------------------------------
    I. Morris Algorithm list of counter
------------------------------------------------------------------------------}

---- * Count the number of items in `as` to within `eps` of actual
---- * with confidence `delta`
morris' :: EpsDelta -> Batch a IO Counter
morris' ed = morris ed `using` evalRandIO

morris :: MonadRandom m => EpsDelta -> Streaming a m Counter
morris (ED e d) = medianOfMeans $ go cs
  where cs            = replicate (t*m) 0
        t             = round $ 1/(e^2*d)  :: Int
        m             = round . log $ 1/d
        go cs         = (\c -> 2^(round c) - 1) `ffmap` Cl.foldM (\cs _ -> traverse incr cs) cs
        ffmap         = fmap . fmap
        medianOfMeans = fmap median' . (fmap . fmap) mean' . fmap (chunksOf t)

{-----------------------------------------------------------------------------
    II. Morris Algorithm list of list of counter
------------------------------------------------------------------------------}

---- * Count the number of items in `as` to within `eps` of actual
---- * with confidence `delta`
morris1' :: EpsDelta -> Batch a IO Counter
morris1' ed = morris1 ed `using` evalRandIO

-- * Run on stream inputs `as` for t independent trials for `t = 1/eps^2 * d`,
-- * and `m` times in parallel, for `m = log(1/d)` and take the median
morris1 :: MonadRandom m => EpsDelta -> Streaming a m Counter
morris1 (ED e d) = medianOfMeans $ go ccs
      where
        medianOfMeans = fmap median' . (fmap . fmap) mean'
        ccs           = replicate m $ replicate t 0
        t             = round $ 1/(e^2*d)
        m             = round . log $ 1/d
        go ccs        = (\x -> 2^(round x) - 1) `fffmap` Cl.foldM (\xs _ -> incrs' xs) ccs
        fffmap        = fmap . fmap . fmap

-- * given a list of list of counters toss a coin for each counter and incr
-- * this can be flattened
incrs' :: MonadRandom m => [[Counter]] -> m [[Counter]]
incrs' = sequence . fmap (sequence . fmap incr)

{-----------------------------------------------------------------------------
    III. Utils
------------------------------------------------------------------------------}

-- * Increment a counter `x` with probability 1/2^x
incr :: MonadRandom m => Counter -> m Counter
incr x = do
  h <- toss . coin $ 0.5^(round x)
  return $ if isHead h then (seq () succ x) else seq () x

mean', median' :: (Floating a, Ord a, RealFrac a) => [a] -> Float
mean'   = fromIntegral . round . mean
median' = fromIntegral . round . median


{-# LANGUAGE DataKinds, NamedFieldPuns, GADTs, KindSignatures, TypeFamilies, TypeOperators, TypeApplications, ScopedTypeVariables, MultiParamTypeClasses, OverloadedLabels, InstanceSigs, FlexibleContexts, AllowAmbiguousTypes, FlexibleInstances, RankNTypes, UndecidableInstances, ConstraintKinds, OverlappingInstances, TypeFamilyDependencies #-}

module TraceTypes where

import           Data.Row.Records        hiding ( Disjoint
                                                , sequence
                                                , map
                                                , zip
                                                )
import           Data.Row.Internal              ( Subset
                                                , type (.\\)
                                                , type (.\/)
                                                )
import           Data.Default
import           Numeric.Log                   as Log
import qualified Data.Vector.Sized             as V
import           Data.Vector.Sized              ( Vector )
import           GHC.OverloadedLabels
import           Data.Finite
import           GHC.TypeLits
import           Control.Monad
import           Control.Monad.Bayes.Class     as Bayes
import Statistics.Distribution (logDensity, complCumulative)
import Statistics.Distribution.Gamma
import Statistics.Distribution.Beta
import Statistics.Distribution.Normal

-- The "Constraint" type is intended only
-- for internal use, and is not user-facing.
-- A Constraint of a certain record type is
-- a record where each field type `a` has been
-- replaced with the type `Maybe a`; a value of
-- `Nothing` indicates an unconstrained value,
-- whereas `Just x` indicates an observed (i.e.
-- constrained) value.
type Constraint r = Rec (Map Maybe r)

-- (ConcatTo l r s) means that l, r, and s are all
-- valid record types, and that l .+ r = r .+ l = s.
type ConcatTo l r s
  = ( WellBehaved l
    , WellBehaved r
    , WellBehaved s
    , Subset l s
    , Subset r s
    , (l .+ r) ≈ s
    , (r .+ l) ≈ s
    , s .\\ l ≈ r
    , s .\\ r ≈ l
    )

-- Are the label sets of l and r disjoint?
type Disjoint l r
  = ( ConcatTo l r (l .+ r)
    , ConcatTo (Map Maybe l) (Map Maybe r) (Map Maybe (l .+ r))
    , Forall (Map Maybe l) Default
    , Forall (Map Maybe r) Default
    )

-- Three main types: probability distributions (D), unnormalized density-carrying measures (U),
-- and probabilistic programs (P).
data U m a = U
  { usampler :: m a
  , udensity :: a -> Log Double
  }
data D m a = D
  { dsampler :: m a
  , ddensity :: a -> Log Double
  }
data P m t a = P
  { traced    :: D m (Rec t)
  , retval    :: Rec t -> a
  , constrain :: Constraint t -> U m (Rec t)
  }

class DensityCarrying f m where
  sampler :: f m a -> m a
  density :: f m a -> a -> Log Double

instance DensityCarrying U m where
  sampler = usampler
  density = udensity

instance DensityCarrying D m where
  sampler = dsampler
  density = ddensity

dist sampler density = D { dsampler = sampler, ddensity = density }
udist sampler density = U { usampler = sampler, udensity = density }

-- DEFAULT INFERENCE
observe
  :: forall m a t s
   . (Disjoint t s, Monad m)
  => P m (t .+ s) a
  -> Rec t
  -> U m (Rec s)
observe p obs =
  let
    -- Turn the observation record `obs` of type `t` into a Constraint of
    -- type `t .+ s`, by filling in the `s` piece of the record with `Nothing`.
      constraints :: Constraint (t .+ s)
      constraints = map' Just obs .+ (default' @Default def :: Constraint s)

      -- Obtain a constrained distribution over traces.
      u :: U m (Rec (t .+ s))
      u = constrain p constraints
  in
    -- Modify the distribution u so that it is over only the unobserved values
    -- (Rec s), not full-but-constrained traces (Rec (t .+ s)).
      udist (fmap restrict $ sampler u) (density u . ((.+) obs))

traceSampler = sampler . traced
traceDensity = density . traced

-- UNIT INTERVAL TYPE

newtype Unit = Sigmoid {logit :: Double}

fromUnit :: Unit -> Double
fromUnit (Sigmoid x) = 1.0 / (1.0 + exp (-x))

mkUnit :: Double -> Unit
mkUnit u = if 0 < u && u < 1
  then Sigmoid (log (u / (1.0 - u)))
  else error "Number must be between 0 and 1."

instance Show Unit where
  show x = show (fromUnit x)

------------------------------
-- STPL LANGUAGE CONSTRUCTS --
------------------------------

-- RETURN

pReturn :: Monad m => a -> P m Empty a
pReturn x =
  let sampler = return Data.Row.Records.empty
      density _ = 1
      traced = dist sampler density
      constrain _ = udist sampler density
      retval _ = x
  in  P { traced, retval, constrain }

-- BIND

pBind :: (Disjoint t s, Monad m) => P m t a -> (a -> P m s b) -> P m (t .+ s) b
pBind f g =
  let
    bindSampler = do
      fTrace <- sampler $ traced f
      let x = retval f fTrace
      gTrace <- sampler $ traced (g x)
      return (fTrace .+ gTrace)

    bindDensity tr =
      let fDensity = density (traced f) (restrict tr)
          x        = retval f (restrict tr)
          gDensity = density (traced $ g x) (restrict tr)
      in  fDensity * gDensity

    bindTraced = dist bindSampler bindDensity

    bindRet tr = let x = retval f (restrict tr) in retval (g x) (restrict tr)

    constrainedSampler c = do
      fTrace <- sampler $ constrain f (restrict c)
      let x = retval f fTrace
      gTrace <- sampler $ constrain (g x) (restrict c)
      return (fTrace .+ gTrace)

    constrainedDensity c tr =
      let fDensity = density (constrain f (restrict c)) (restrict tr)
          x        = retval f (restrict tr)
          gDensity = density (constrain (g x) (restrict c)) (restrict tr)
      in  fDensity * gDensity

    bindConstrain c = udist (constrainedSampler c) (constrainedDensity c)
  in
    P { traced = bindTraced, retval = bindRet, constrain = bindConstrain }

-- SAMPLE

singleLabelConstrain
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> D m (Rec (l .== a))
  -> Constraint (l .== a)
  -> U m (Rec (l .== a))
singleLabelConstrain lbl traced c = case c .! lbl of
  Just x ->
    let t                  = lbl .== x
        constrainedSampler = do
          score $ density traced t
          return t
        constrainedDensity _ = density traced t
    in  udist constrainedSampler constrainedDensity

  Nothing -> udist (sampler traced) (density traced)

sample :: (KnownSymbol l, MonadInfer m) => Label l -> D m a -> P m (l .== a) a
sample lbl d =
  let sampleTraced = dist
        (do
          x <- sampler d
          return (lbl .== x)
        )
        (\t -> density d (t .! lbl))

      sampleRet t = t .! lbl

      sampleConstrain = singleLabelConstrain lbl sampleTraced
  in  P { traced    = sampleTraced
        , retval    = sampleRet
        , constrain = sampleConstrain
        }

-- WITH PROBABILITY

withProbability
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> Unit
  -> P m t a
  -> P m s a
  -> P m (l .== Either (Rec t) (Rec s)) a
withProbability lbl logit f g =
  let p         = fromUnit logit

      wpSampler = do
        r <- random
        if r < p
          then do
            t <- sampler $ traced f
            return (lbl .== Left t)
          else do
            t <- sampler $ traced g
            return (lbl .== Right t)

      wpDensity tr = case tr .! lbl of
        Left  t -> Exp (log p) * density (traced f) t
        Right t -> Exp (log (1.0 - p)) * density (traced g) t

      wpTraced = dist wpSampler wpDensity

      wpRet t = case t .! lbl of
        Left  t -> retval f t
        Right t -> retval g t

      wpConstrain = singleLabelConstrain lbl wpTraced
  in  P { traced = wpTraced, retval = wpRet, constrain = wpConstrain }

-- FOR EACH

forEach
  :: (KnownSymbol l, KnownNat n, MonadInfer m)
  => Label l
  -> Vector n a
  -> (a -> P m t b)
  -> P m (l .== Vector n (Rec t)) (Vector n b)
forEach lbl xs body =
  let
    forEachSampler = do
      traces <- V.mapM (sampler . traced . body) xs
      return (lbl .== traces)

    forEachDensity t = V.product
      (V.map (\(x, tr) -> density (traced $ body x) tr) (V.zip xs (t .! lbl)))

    forEachTraced = dist forEachSampler forEachDensity

    forEachRet t = V.map (\(x, tr) -> retval (body x) tr) (V.zip xs (t .! lbl))

    forEachConstrain = singleLabelConstrain lbl forEachTraced
  in
    P { traced    = forEachTraced
      , retval    = forEachRet
      , constrain = forEachConstrain
      }

-- FOR ... IN RANDOM RANGE ...

forRandomRange
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> D m Int
  -> (Int -> P m t a)
  -> P m (l .==  [Rec t]) [a]
forRandomRange lbl d body =
  let
    -- Sample a trace of the for loop.
    forSampler = do
      n      <- sampler d
      traces <- sequence $ map (sampler . traced . body) [0 .. n - 1]
      return (lbl .== traces)

    -- Compute the density of a trace of the for loop.
    forDensity t =
      density d (length $ t .! lbl)
        * product
            (map (\(tr, i) -> density (traced $ body i) tr)
                 (zip (t .! lbl) [0 ..])
            )

    -- The for loop's density-carrying distribution over traces
    forTraced = dist forSampler forDensity

    -- Compute the return value of the for loop, given a trace of an execution.
    -- The return value is just a list of the return values from each iteration.
    forRet t = map (\(tr, i) -> retval (body i) tr) (zip (t .! lbl) [0 ..])

    -- Constrain the random choices made inside the for loop.
    forConstrain = singleLabelConstrain lbl forTraced
  in
    P { traced = forTraced, retval = forRet, constrain = forConstrain }

{- WHILE:

    A stochastic while loop.
    Takes as input:
    - lbl    : The label at which to trace the while loop's random choices.
    - init   : The initial state of the while loop.
    - pi     : A function computing a probability of continuing, based on the current state.
    - pi_max : An upper bound beyond which the probability of continuing is truncated.
    - body   : The body of the while loop, which simulates a new state based on the current state.

    The resulting probabilistic program has type P m (l .== [Rec t]) a:
    - The trace contains, at the specified label, a list of subtraces, one for each iteration.
    - The return value is the final state at the finish of the while loop.

 -}

while
  :: (KnownSymbol l, MonadInfer m)
  => Label l
  -> a
  -> (a -> Unit)
  -> Unit
  -> (a -> P m t a)
  -> P m (l .== [Rec t]) a
while lbl init pi pi_max body =
  let
    -- Sample a trace of the while loop.
    whileSampler state = do
      shouldContinue <- bernoulli $ min (fromUnit (pi state)) (fromUnit pi_max)
      if shouldContinue
        then do
          nextTrace <- sampler (traced $ body state)
          let nextState = retval (body state) nextTrace
          restOfTraces <- whileSampler nextState
          return (lbl .== (nextTrace : (restOfTraces .! lbl)))
        else return (lbl .== [])

    -- Given traces from each iteration, and an initial state,
    -- compute the final state.
    retvalFromTraces init ts = case ts of
      []     -> init
      t : ts -> retvalFromTraces (retval (body init) t) ts

    -- Given a trace of the entire while loop, compute the final
    -- state.
    whileRet t = retvalFromTraces init (t .! lbl)

    -- Given an initial state and traces from each iteration,
    -- compute the density.
    densityFromTraces init ts = case ts of
        [] -> Exp (log (1.0 - (fromUnit (pi init))))
        t : ts ->
          Exp (log (fromUnit (pi init)))
            * (density (traced $ body init) t)
            * densityFromTraces (retval (body init) t) ts

    -- Given a trace of the entire while loop, compute the density.
    whileDensity t = densityFromTraces init (t .! lbl)

    -- The while loop's density-carrying distribution over traces.
    whileTraced    = dist (whileSampler init) whileDensity

    -- Constrain the choices made inside the loop.
    whileConstrain = singleLabelConstrain lbl whileTraced

  in  P { traced = whileTraced, retval = whileRet, constrain = whileConstrain }

-------------------
-- DISTRIBUTIONS --
-------------------

nrm :: MonadInfer m => Double -> Log Double -> D m Double
nrm mu (Exp sigln) = dist (normal mu (exp sigln)) (normalPdf mu (exp sigln))

brn :: MonadSample m => Unit -> D m Bool
brn logit =
  let p = fromUnit logit
  in  dist (bernoulli p) (\b -> Exp . log $ if b then p else (1.0 - p))

gmma :: MonadSample m => Log Double -> Log Double -> D m (Log Double)
gmma shape scale =
  let toDouble (Exp l) = exp l
      toLogDouble d = Exp (log d)
  in
    dist (fmap toLogDouble (gamma (toDouble shape) (toDouble scale)))
         (Exp . (logDensity $ Statistics.Distribution.Gamma.gammaDistr
                                (toDouble shape) (toDouble scale))
              . toDouble)

geom :: MonadSample m => Unit -> D m Int
geom p' =
  let p = fromUnit p'
  in dist (geometric p)
          (\n -> Exp (fromIntegral n * (log p) + (log (1.0 - p))))

pois :: MonadSample m => Int -> D m Int
pois = undefined

bta :: MonadSample m => Log Double -> Log Double -> D m Unit
bta (Exp lna) (Exp lnb) =
  let
    a = exp lna
    b = exp lnb
  in
    dist (fmap mkUnit (beta a b))
         (Exp . (logDensity $ Statistics.Distribution.Beta.betaDistr a b)
              . fromUnit)

unif :: MonadSample m => D m Unit
unif = dist (fmap mkUnit Bayes.random) (\x -> 1.0 :: Log Double)

cat :: (KnownNat n, MonadSample m) => Vector n Unit -> D m (Finite n)
cat probs = undefined

lognormal :: forall m. MonadInfer m => Log Double -> Log Double -> D m (Log Double)
lognormal (Exp logmu) sigma =
  dist (fmap Exp $ sampler $ nrm logmu sigma)
       (\(Exp logx) -> density (nrm logmu sigma :: D m Double) logx / (Exp logx))

truncnrm :: MonadInfer m => Log Double -> Log Double -> D m (Log Double)
truncnrm (Exp logmu) (Exp logsigma) =
  let mu = exp logmu
      sigma = exp logsigma
      unbounded = nrm mu (Exp logsigma)

      truncNormSampler = do
        x <- sampler unbounded
        if x < 0 then truncNormSampler else return (Exp $ log x)

      truncNormPdf (Exp logy) =
        let y = exp logy
            z = complCumulative (normalDistr mu sigma) 0
        in (density unbounded y / (Exp (log z)))
  in dist truncNormSampler truncNormPdf

----------------------------
-- PROGRAMMABLE INFERENCE --
----------------------------

-- IMPORTANCE SAMPLING

importance
  :: (Disjoint t s, MonadInfer m)
  => P m (t .+ s) a
  -> Rec t
  -> P m s b
  -> m (Rec s)
importance p t q =
  let target = observe p t
  in  do
        tr <- sampler $ traced q
        score ((density target tr) / (density (traced q) tr))
        return tr

-- SMC

unroll
  :: (MonadInfer m, Disjoint t s)
  => P m i a
  -> (a -> P m (t .+ s) a)
  -> [Rec t]
  -> m (Rec i, [Rec s])
unroll init next steps = do
  tInit <- sampler $ traced init
  let processStep soFar t = do {
    (s, results) <- soFar;
    tNext <- sampler $ observe (next s) t;
    return (retval (next s) (tNext .+ t), tNext : results)
  }
  (s, l) <- foldl processStep (return (retval init tInit, [])) steps
  return (tInit, reverse l)

particleFilter
  :: (MonadInfer m, Disjoint t s,
      Disjoint Empty i, (Empty .+ i) ≈ i)
  => P m i a
  -> P m i b
  -> (a -> P m (t .+ s) a)
  -> [(Rec t, a -> P m s c)]
  -> m (Rec i, [Rec s])
particleFilter init qInit next steps = do
  tInit <- importance init Data.Row.Records.empty qInit
  let processStep soFar (t, q) = do {
    (s, results) <- soFar;
    tNext <- importance (next s) t (q s);
    return (retval (next s) (tNext .+ t), tNext : results)
  }
  (s, l) <- foldl processStep (return (retval init tInit, [])) steps
  return (tInit, reverse l)

-- MCMC

type K m t (s :: Row *) = U m (Rec t) -> Rec t -> m (Rec t)

mh
  :: forall s t m a
   . (Disjoint t s, MonadInfer m)
  => (Rec (t .+ s) -> P m t a)
  -> K m (t .+ s) t
mh q p old =
  let proposal = do
        new <- sampler (traced $ q old)
        return $ new .+ (restrict old :: Rec s)
      rho (x, y) =
          ((density p y) * (density (traced $ q y) (restrict x)))
            / ((density p x) * (density (traced $ q x) (restrict y)))
  in  do
        proposed <- proposal
        r        <- Bayes.random
        if Exp (log r) < (min 1 (rho (old, proposed)))
          then return proposed
          else return old

seqK :: Monad m => K m t s -> K m t r -> K m t (s .\/ r)
seqK k1 k2 p old = do
  t' <- k1 p old
  k2 p t'

mixK :: MonadSample m => Double -> K m t s -> K m t r -> K m t (s .\/ r)
mixK p k1 k2 model old = do
  r <- Bayes.random
  if r < p then k1 model old else k2 model old

repeatK :: Monad m => Int -> K m t s -> K m t s
repeatK n k p old = if n == 0
  then return old
  else do
    t <- k p old
    repeatK (n - 1) k p t

ifK
  :: (Disjoint t s, Monad m)
  => (Rec t -> Bool)
  -> K m (t .+ s) s
  -> K m (t .+ s) s
ifK b k p old = if (b (restrict old)) then k p old else return old

-- VARIATIONAL (STUBBED)

svi
  :: (Disjoint t s, Monad m)
  => P m (t .+ s) a
  -> Rec t
  -> (theta -> P m s b)
  -> theta
  -> m theta
svi p t q theta = undefined
-- do tr <- traceSampler $ q theta
--    let sco  = density (observe p t) tr
--    let gradient = grad (\th -> traceDensity (q th) tr) theta
--    return theta + stepSize * sco * gradient

trainAmortized
  :: (Disjoint t s, Monad m)
  => P m (t .+ s) a
  -> (theta -> Rec t -> P m s b)
  -> theta
  -> m theta
trainAmortized p q theta = undefined
-- do tr <- traceSampler p
--    let gradient = grad (\th -> traceDensity (q th (restrict tr)) (restrict tr)) theta
--    return theta + step * gradient

ifThenElse c t f = if c then t else f


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLabels      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Add
Description : Layer to add some constant, with shape broadcasting
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}

module Grenade.Layers.Add
  (
  -- * Layer instances
    Add (..)

  -- * Helper functions
  , initAdd
  ) where

import           Control.DeepSeq              (NFData (..))

import           Data.Kind                    (Type)
import           Data.Maybe                   (fromJust)
import           Data.Proxy
import           Data.Serialize
import           GHC.TypeLits

import qualified Numeric.LinearAlgebra        as LA
import           Numeric.LinearAlgebra.Static (R)
import qualified Numeric.LinearAlgebra.Static as H

import           Lens.Micro                   ((^.))

import           Grenade.Core
import           Grenade.Layers.Internal.Add
import           Grenade.Onnx

data Add :: Nat -- The number of channels of the bias
         -> Nat -- The number of rows of the bias
         -> Nat -- The number of columns of the bias
         -> Type where
  Add  :: ( KnownNat channels
          , KnownNat rows
          , KnownNat columns )
          => R (channels * rows * columns)
          -> Add channels rows columns

instance UpdateLayer (Add c h w) where
  type Gradient (Add c h w) = ()
  runUpdate _ x _  = x
  reduceGradient _ = ()

instance (KnownNat c, KnownNat h, KnownNat w ) => RandomLayer (Add c h w) where
  createRandomWith _ _ = pure initAdd

-- | Initialize an Add layer with 0 bias
initAdd :: forall c h w. ( KnownNat c, KnownNat h, KnownNat w )
        => Add c h w
initAdd =
  let c'      = fromIntegral $ natVal (Proxy :: Proxy c)
      h'      = fromIntegral $ natVal (Proxy :: Proxy h)
      w'      = fromIntegral $ natVal (Proxy :: Proxy w)
      zeroes  = replicate (c' * h' * w') 0
      bias    = H.fromList zeroes   :: R (c * h * w)
  in Add bias

instance ( KnownNat c, KnownNat h, KnownNat w ) => Serialize (Add c h w) where
  put (Add bias) = putListOf put . LA.toList . H.extract $ bias
  get            = do
    bias <- maybe (fail "Vector of incorrect size") return . H.create . LA.fromList =<< getListOf get
    return $ Add bias

-- | A three dimensions image can have a vector added to it with size equal to the number of channels
--   of the image, this corresponds to adding the nth element of the bias to every pixel of the nth
--   channel of the input image
instance ( KnownNat i, KnownNat j, KnownNat k ) => Layer (Add k 1 1) ('D3 i j k) ('D3 i j k) where
  type Tape (Add k 1 1) ('D3 i j k) ('D3 i j k) = ()

  runForwards (Add b) (S3D m)
    = let c  = fromIntegral $ natVal (Proxy :: Proxy k)
          h  = fromIntegral $ natVal (Proxy :: Proxy i)
          w  = fromIntegral $ natVal (Proxy :: Proxy j)
          m' = H.extract m
          b' = H.extract b
          r  = addPerChannel c h w m' b'
      in  ((), S3D . fromJust . H.create $ r)

  runBackwards = undefined

instance OnnxOperator (Add c h w) where
  onnxOpTypeNames _ = ["Add"]

instance (KnownNat c, KnownNat h, KnownNat w) => OnnxLoadable (Add c h w) where
  loadOnnxNode inits node = case node ^. #input of
    [bias, _] -> do
      loadedBias <- readInitializerTensorIntoVector inits bias
      return $ Add loadedBias
    _ -> onnxIncorrectNumberOfInputs

instance NFData (Add c h w) where
  rnf (Add bias) = rnf bias


{-# LANGUAGE GADTs, FlexibleContexts, RankNTypes #-}

module Numeric.Trainee.Neural (
	Net, net, netSeed,
	(⭃),
	ndup, npar, ndupWith, nsum, nfold,
	pad, pad2,
	maxPool, maxPool2,
	fc, conv, conv2,
	dconv, dconv2,

	sigma, relu, softplus,
	summator, biaser, activator, convolver, convolver2,
	normVar,

	module Numeric.Trainee.Types,
	module Numeric.Trainee.Learnee
	) where

import Prelude hiding ((.), id)
import Prelude.Unicode

import Control.Monad (replicateM, liftM2)
import Data.Random
import qualified Data.Vector as V
import Numeric.LinearAlgebra hiding (conv, conv2)
import System.Random (newStdGen, StdGen)

import Numeric.Trainee.Types
import Numeric.Trainee.Gradee
import Numeric.Trainee.Learnee

type Net a = Learnee (Vector a) (Vector a)

net ∷ RVar a → IO a
net n = netSeed <$> newStdGen <*> pure n

netSeed ∷ StdGen → RVar a → a
netSeed g n = fst $ sampleState n g

(⭃) ∷ RVar (Learnee a b) → RVar (Learnee b c) → RVar (Learnee a c)
n ⭃ l = liftM2 (⇉) n l

-- | Split layers, just wraps @vdup@
ndup ∷ Num a ⇒ Int → RVar (Learnee a (V.Vector a))
ndup = return ∘ computee ∘ vdup

-- | Process vector of layers
npar ∷ Int → RVar (Learnee a b) → RVar (Learnee (V.Vector a) (V.Vector b))
npar n = fmap (parallel ∘ V.fromList) ∘ replicateM n

-- | Split layers, justs wraps @vdupWith@
ndupWith ∷ (a → a → a) → Int → RVar (Learnee a (V.Vector a))
ndupWith fn = return ∘ computee ∘ vdupWith fn

-- | Join layers with sum
nsum ∷ Num a ⇒ RVar (Learnee (V.Vector a) a)
nsum = return (computee vsum)

-- | Join layers with custom sum
nfold ∷ (a → a → a) → RVar (Learnee (V.Vector a) a)
nfold = return ∘ computee ∘ vfold

-- | Padding
pad ∷ Numeric a ⇒ Int → RVar (Learnee (Vector a) (Vector a))
pad = return ∘ computee ∘ padVec

-- | 2d padding
pad2 ∷ Numeric a ⇒ Int → Int → RVar (Learnee (Matrix a) (Matrix a))
pad2 w h = return ∘ computee $ padMat w h

-- | Max pool layer
maxPool ∷ (RealFrac a, Numeric a) ⇒ Int → RVar (Learnee (Vector a) (Vector a))
maxPool = return ∘ computee ∘ maxPoolVec

-- | 2d max pool layer
maxPool2 ∷ (RealFrac a, Numeric a) ⇒ Int → Int → RVar (Learnee (Matrix a) (Matrix a))
maxPool2 w h = return ∘ computee $ maxPoolMat w h

-- | Fully connected layer
fc ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Vector a) → Int → Int → RVar (Net a)
fc f inputs outputs = do
	s ← summator inputs outputs
	b ← biaser outputs
	return $ s ⇉ b ⇉ activator f

-- | Convolution layer 1-d
conv ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Vector a) → Int → RVar (Net a)
conv f w = do
	c ← convolver w
	b ← do
		bs ← normVar
		return $ learnee biasVec bs
	return $ c ⇉ b ⇉ activator f

-- | Convolution layer 2-d
conv2 ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Matrix a) → (Int, Int) → RVar (Learnee (Matrix a) (Matrix a))
conv2 f (w, h) = do
	c ← convolver2 w h
	b ← do
		bs ← normVar
		return $ learnee biasMat bs
	return $ c ⇉ b ⇉ activator f

-- | Depth (with several input and output channels) convolution 1-d
dconv ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Vector a) → Int → Int → Int → RVar (Learnee (V.Vector (Vector a)) (V.Vector (Vector a)))
dconv f inputs outputs width =
	ndupWith (V.zipWith (+)) outputs ⭃
	npar outputs (npar inputs (conv f width) ⭃ nsum)

-- | Depth (with several input and output channels) convolution 2-d
dconv2 ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Unary (Matrix a) → Int → Int → (Int, Int) → RVar (Learnee (V.Vector (Matrix a)) (V.Vector (Matrix a)))
dconv2 f inputs outputs (w, h) =
	ndupWith (V.zipWith (+)) outputs ⭃
	npar outputs (npar inputs (conv2 f (w, h)) ⭃ nsum)

sigma ∷ Floating a ⇒ a → a
sigma t = 1 / (1 + exp (negate t))

relu ∷ Fractional a ⇒ a → a
relu t = 0.5 * (1 + signum t) * t

softplus ∷ Floating a ⇒ a → a
softplus t = log (1 + exp t)

summator ∷ (Distribution Normal a, Fractional a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → Int → RVar (Net a)
summator inputs outputs = do
	ws ← replicateM (inputs * outputs) normVar
	return $ learnee matVec ((outputs >< inputs) ws)

biaser ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → RVar (Net a)
biaser sz = do
	bs ← replicateM sz normVar
	return $ learnee odot (fromList bs)

activator ∷ Num a ⇒ Unary a → Learnee a a
activator f = computee (unary f)

convolver ∷ (Distribution Normal a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → RVar (Net a)
convolver w = do
	ws ← replicateM w normVar
	return $ learnee corrVec (fromList ws)

convolver2 ∷ (Distribution Normal a, Fractional a, Numeric a, Parametric (Vector a), Parametric a) ⇒ Int → Int → RVar (Learnee (Matrix a) (Matrix a))
convolver2 w h = do
	ws ← replicateM (w * h) normVar
	return $ learnee corrMat ((w >< h) ws)

normVar ∷ (Distribution Normal a, Fractional a) ⇒ RVar a
normVar = normal 0.0 0.25