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


module AI.HsANN.ComplexNetwork
    (
    -- * Types
      ComplexNetwork(..)
    , ActivationFunction
    , ActivationFunctionDerivative
    , Sample
    , Samples
    , (-->)

    -- * Creating a neural network
    , createComplexNetwork
    , fromWeightMatrices

    -- * Computing a neural network's output
    , output
    , sigmoid
    , sigmoid'
    , complexSigmoid
    , complexSigmoid'
    , randComplex

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

    -- * Loading and saving a neural network
    , loadComplexNetwork
    , saveComplexNetwork
    ) 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 qualified Data.Complex as C
import Debug.Trace

import System.Random.MWC
import Numeric.LinearAlgebra
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 ComplexNetwork a = ComplexNetwork
                 { matrices   :: V.Vector (Matrix a) -- ^ the weight matrices
                 } deriving Show

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


--instance (Complex a, Variate a) => Variate (Complex a) where
--	uniform g = uniform g :+ uniform g
--	uniformR (x,y) g = uniformR (x,y) g :+ uniformR (x,y) g




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


--

randComplex :: (Variate a, Num a) => IO(Complex a)
randComplex = withSystemRandom . asGenST $ \gen -> do
  r <- uniformR (-1,1) gen
  i <- uniformR (-1,1) gen
  return (r :+ i)

randComplexList :: (Variate a, Num a) => Int -> IO([Complex a])
randComplexList 0 = return []
randComplexList n = do
  c <- randComplex
  lst <- randComplexList (n-1)
  return (c : lst)

randComplexMatrix :: (Variate a, Storable a, Num a) => (Int,Int) -> IO(Matrix (Complex a))
randComplexMatrix (rows,cols) = do
  lst <- randComplexList (rows*cols)
  return (reshape cols $ Numeric.LinearAlgebra.fromList lst)

-- | 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.uniformR'.
-- > createComplexNetwork n l k
createComplexNetwork :: (Variate a, Storable a, Num a) => Int -> [Int] -> Int -> IO (ComplexNetwork (Complex a))
createComplexNetwork nInputs hiddens nOutputs =
  fmap ComplexNetwork $ go dimensions V.empty
  where
        go [] !ms         = return ms
        go ((!n,!m):ds) ms = do
          !mat <- randComplexMatrix (n,m)
          go ds (ms `V.snoc` mat)
        dimensions      = zip (hiddens ++ [nOutputs]) $
                              (nInputs : hiddens)
{-# INLINE createComplexNetwork #-}



-- | 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 (Complex a)) -> ComplexNetwork (Complex a)
fromWeightMatrices ws = ComplexNetwork 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."


--without the trick (of constant input)

output :: (Floating (Vector a), Product a, Storable a, Num (Vector a)) => ComplexNetwork a -> ActivationFunction a -> Vector a -> Vector a
output (ComplexNetwork{..}) act input = V.foldl' f input matrices
  where f !inp m = mapVector 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), Product a, Storable a, Num (Vector a)) => ComplexNetwork a -> ActivationFunction a -> Vector a -> V.Vector (Vector a)
outputs (ComplexNetwork{..}) act input = V.scanl f input matrices
  where f !inp m = mapVector act $ m <> inp
{-# INLINE outputs #-}


conju (x:+y) = x:+(-y)

deltas :: (Floating b, Floating (Vector a), Floating a, Product a, Storable a, Num (Vector a), Container Vector a, a ~ Complex b) => ComplexNetwork a -> ActivationFunctionDerivative a -> V.Vector (Vector a) -> Vector a -> V.Vector (Matrix a)
deltas (ComplexNetwork{..}) act' os expected = V.zipWith outer (V.tail ds) (V.init (V.map (mapVector conju) os))
  where !dl = (V.last os - expected) * (deriv $ mapVector conju (V.last os)) -- = (out_last - target)*(f'(netin_last)) = dev_last * deriv(out_last) = partial_last
        !ds = V.scanr f dl (V.zip os matrices) -- generates dev_i
        f (!o, m) !del = deriv o * ((ctrans m) <> del) -- dev_k = deriv(out_k)*dev_k = deriv(out_k)*(WEIGHTS <> partial_k+1)
        deriv = mapVector act'
{-# INLINE deltas #-}

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

-- | 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 b, Floating a, Product a, Num (Vector a), Container Vector a, a ~ Complex b) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork a
backpropOnce rate act act' n samples = foldl' (updateComplexNetwork 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 b, Floating a, Product a, Num (Vector a), Container Vector a, a ~ Complex b) => (Int -> ComplexNetwork a -> Samples a -> Bool) -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork 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 b, Floating a, Product a, Num (Vector a), Container Vector a, a ~ Complex b) => Int -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork 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 b, Floating a, Num (Vector a), Num (RealOf a), Product a, a ~ Complex b) => ActivationFunction a -> ComplexNetwork a -> Samples a -> RealOf a
quadError act net samples = foldl' (\err (inp, out) -> err + (norm2 $ 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 unappropriate activation
-- functions.
trainUntilErrorBelow :: (Floating (Vector a), Floating b, Floating a, Product a, Num (Vector a), Ord a, Container Vector a, Num (RealOf a), a ~ RealOf a, Show a, a ~ Complex b) => a -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork 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' #-}


complexSigmoid :: Floating a => Complex a -> Complex a
complexSigmoid (x :+ y) = (sigmoid x) :+ (sigmoid y)

complexSigmoid' :: Floating a => Complex a -> Complex a
complexSigmoid' (x :+ y) = (sigmoid' x) :+ (sigmoid' y)

-- | 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.
loadComplexNetwork :: (Storable a, Element a, Binary a) => FilePath -> IO (ComplexNetwork a)
loadComplexNetwork fp = decode . decompress <$> B.readFile fp
{-# INLINE loadComplexNetwork #-}

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