{-# LANGUAGE FlexibleContexts,
             MultiParamTypeClasses #-}
module Network.Trainer.BackpropTrainer
( BackpropTrainer(..)

, backprop
, inputs
, outputs
, deltas
, hiddenDeltas
, calculateNablas
, fit
, evaluate
) where

import Network.Network
import Network.Network.FeedForwardNetwork
import Network.Neuron
import Network.Layer
import Network.Trainer

import Numeric.LinearAlgebra

-- | A BackpropTrainer performs simple backpropagation on a neural network.
--   It can be used as the basis for more complex trainers.
data BackpropTrainer = BackpropTrainer { eta :: Double
                                       , cost :: CostFunction
                                       , cost' :: CostFunction'
                                       }

-- | Declare the BackpropTrainer to be an instance of Trainer.
instance Trainer BackpropTrainer FeedForwardNetwork where
  fit s t n examples = foldl (backprop t) n $ s examples
  -- | Use the cost function to determine the error of a network
  evaluate t n example = (cost t) (snd example) (predict (fst example) n)

-- | Perform backpropagation on a single training data instance.
backprop :: BackpropTrainer -> FeedForwardNetwork -> [TrainingData] -> FeedForwardNetwork
backprop t n es =
  updateNetwork (length es) t (foldl (calculateNablas t n) emptyFeedForwardNetwork es) n

-- | Given the size of the minibatch, the trainer, the nablas for each layer, given
--   as a network, and the network itself, return a network with updated wieghts.
updateNetwork :: Int -> BackpropTrainer -> FeedForwardNetwork -> FeedForwardNetwork -> FeedForwardNetwork
updateNetwork mag t nablas n = addFeedForwardNetworks n
  (FeedForwardNetwork $ map (scaleLayer $ -1 * (eta t) / (fromIntegral mag)) (layers nablas))

-- | Calculate the nablas for a minibatch and return them as a network (so each
--   weight and bias gets its own nabla).
calculateNablas :: BackpropTrainer -> FeedForwardNetwork -> FeedForwardNetwork -> TrainingData -> FeedForwardNetwork
calculateNablas t n nablas e = FeedForwardNetwork $ map (updateLayer t) (zip3 (layers n) ds os)
  where ds = deltas t n e
        os = outputs (fst e) n

-- | The mapped function to update the weight and biases in a single layer
updateLayer :: BackpropTrainer -> (Layer, Vector Double, Vector Double) -> Layer
updateLayer t (l, delta, output) = Layer newWeight newBias n
  where n = neuron l
        newWeight = ((reshape 1 delta) <> (reshape (dim output) output))
        newBias = delta

-- | The outputs function scans over each layer of the network and stores the
--   activated results
outputs :: Vector Double -> FeedForwardNetwork -> [Vector Double]
outputs input network = scanl apply input (layers network)

-- | The inputs function performs a similar task to outputs, but returns a list
--   of vectors of unactivated inputs
inputs :: Vector Double -> FeedForwardNetwork -> [Vector Double]
inputs input network = if null (layers network) then []
  else unactivated : inputs activated (FeedForwardNetwork (tail $ layers network))
  where unactivated = weightMatrix layer <> input + biasVector layer
        layer = head $ layers network
        activated = mapVector (activation (neuron layer)) unactivated

-- | The deltas function returns a list of layer deltas.
deltas :: BackpropTrainer -> FeedForwardNetwork -> TrainingData -> [Vector Double]
deltas t n example = hiddenDeltas
  (FeedForwardNetwork (reverse (layers n))) outputDelta (tail $ reverse is)
    ++ [outputDelta]
  where outputDelta = costd (snd example) output *
          mapVector activationd lastInput
        costd = cost' t
        activationd = activation' (neuron (last (layers n)))
        output = last os
        lastInput = last is
        is = inputs (fst example) n
        os = outputs (fst example) n

-- | Compute the hidden layer deltas
hiddenDeltas :: FeedForwardNetwork -> Vector Double -> [Vector Double] -> [Vector Double]
hiddenDeltas n prevDelta is = if length (layers n) <= 1 then []
  else delta : hiddenDeltas rest delta (tail is)
  where rest = FeedForwardNetwork (tail $ layers n)
        delta = (trans w) <> prevDelta * spv
        w = weightMatrix (head $ layers n)
        spv = mapVector (activation' (neuron (head $ layers n))) (head is)
