content: {-# LANGUAGE ScopedTypeVariables #-}
  
  module AI.Learning.NeuralNetwork
      (
        -- * Representation
        NeuralNetwork
        -- * Prediction and Training
      , nnPredict
      , nnTrain
      , nnTrainIO
      ) where
  
  import Control.Monad.Random hiding (fromList)
  import Numeric.LinearAlgebra
  import Numeric.LinearAlgebra.Util
  import System.IO.Unsafe
  
  import AI.Learning.Core
  import AI.Util.Matrix
  
  -----------------------
  -- NN Representation --
  ----------------------
  
  -- |Representation for a single-hidden-layer neural network. If the network
  --  has K input nodes, H hidden nodes and L output layers then the dimensions
  --  of the matrices theta0 and theta1 are
  --
  --  * size Layer1 = (K+1) x H
  --  * size Layer2 = (H+1) x L
  --
  --  the (+1)s account for the addition of bias nodes in the input layer and
  --  the hidden layer. Therefore the total number of parameters is
  --  H * (K + L) + H + L
  data NeuralNetwork = NN (Matrix Double) (Matrix Double)
  
  -- |Three-tuple describing the shape of a network: (input, hidden, output).
  type NNShape = (Int,Int,Int)
  
  instance Show NeuralNetwork where
      show (NN theta0 theta1) = "Neural Net:\n\n" ++ dispf 3 theta0 ++
                                             "\n" ++ dispf 3 theta1
  
  fromVector :: NNShape -> Vector Double -> NeuralNetwork
  fromVector (k,h,l) vec = NN theta0 theta1
      where theta0 = reshape h $ takeVector ((k + 1) * h) vec
            theta1 = reshape l $ dropVector ((k + 1) * h) vec
  
  toVector :: Matrix Double -> Matrix Double -> Vector Double
  toVector theta0 theta1 = join [flatten theta0, flatten theta1]
  
  ----------------------
  -- NN Train/Predict --
  ----------------------
  
  -- |Make predictions using a neural network.
  nnPredict :: NeuralNetwork -> Matrix Double -> Matrix Double
  nnPredict nn x = h where (_,_,h) = nnForwardProp nn x
  
  -- |Train a neural network from a training set. Note that you must supply
  --  an initial vector of weights. Supplying initial weights all equal to
  --  zero will generally give poor performance.
  nnTrain :: Int              -- number of hidden neurons
          -> Matrix Double    -- y
          -> Matrix Double    -- x
          -> Double           -- lambda
          -> Vector Double    -- initial weights
          -> NeuralNetwork
  nnTrain h y x lambda initialVec = fromVector shape vec
      where shape = (cols x, h, cols y)
            vec   = minimizeS cost grad initialVec
            cost  = fst . nnCostGradient shape y x lambda
            grad  = snd . nnCostGradient shape y x lambda
  
  -- |Train a neural network from a training set.
  --
  --  Note that this is implemented as an IO action, because it randomly sets the
  --  weights in the network before performing numerical optimization (this is to
  --  break the symmetry that exists in the network when all weights are zero).
  --  As a result, it is possible to get different results when training the
  --  network multiple times with the same data.
  nnTrainIO :: Int              -- number of hidden neurons
            -> Matrix Double    -- y
            -> Matrix Double    -- x
            -> Double           -- lambda
            -> IO NeuralNetwork
  nnTrainIO h y x lambda = nnTrain h y x lambda `fmap` initialVec (cols x, h, cols y)
  
  -- |Choose initial random weights for a neural network.
  initialVec :: NNShape -> IO (Vector Double)
  initialVec (k,h,l) = do
      let len = h * (k + l) + h + l
      xs <- getRandomRs (0.0, 0.01)
      return . fromList $ take len xs
  
  ------------------------------
  -- Forward/Back Propagation --
  ------------------------------
  
  -- |Perform forward propagation through a neural network, returning the matrices
  --  created in the process.
  nnForwardProp :: NeuralNetwork                                 -- neural net
                -> Matrix Double                                 -- design matrix (x)
                -> (Matrix Double, Matrix Double, Matrix Double) -- results of fwd prop
  nnForwardProp (NN theta0 theta1) x = (a0,a1,a2)
      where a0 = addOnes $ x
            a1 = addOnes $ sigmoid (a0 <> theta0)
            a2 = sigmoid (a1 <> theta1)
  
  -- |Perform backward propagiation through a neural network. You must supply the
  --  target values and the results of forward propagation for each layer, and
  --  the function returns the gradient matrices for the neural network.
  nnBackProp :: NeuralNetwork                                 -- neural net
             -> Matrix Double                                 -- target (y)
             -> (Matrix Double, Matrix Double, Matrix Double) -- results of fwd prop
             -> (Matrix Double, Matrix Double)                -- gradient (delta0, delta1)
  nnBackProp (NN _ theta1) y (a0,a1,a2) = (dropColumns 1 delta0, delta1)
      where
          d2     = a2 - y
          d1     = (d2 <> trans theta1) * a1 * (1 - a1)
          delta0 = trans a0 <> d1
          delta1 = trans a1 <> d2
  
  -- |Perform back and forward propagation through a neural network, returning the
  --  final predictions (variable /a2/) and the gradient matrices (variables
  --  /delta0/ and /delta1/) produced.
  nnFwdBackProp :: NeuralNetwork -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double, Matrix Double)
  nnFwdBackProp nn@(NN theta0 theta1) y x = (a2, delta0, delta1)
      where
          (a0,a1,a2)      = nnForwardProp nn x
          (delta0,delta1) = nnBackProp nn y (a0,a1,a2)
  
  -- |Compute the penalty function and gradient vector for a neural network
  --  given a training set.
  nnCostGradient :: NNShape                   -- (K,H,L)
                 -> Matrix Double             -- targets (y)
                 -> Matrix Double             -- design matrix (x)
                 -> Double                    -- regularization parameter (lambda)
                 -> Vector Double             -- neural network
                 -> (Double, Vector Double)   -- (cost, gradient)
  nnCostGradient shape y x lambda vec = (cost, grad)
      where
          m = fromIntegral (rows x)
          nn@(NN theta0 theta1) = fromVector shape vec
          (h, delta0, delta1)   = nnFwdBackProp nn y x
  
          cost  = (cost1 + cost2) / m
          cost1 = negate $ sumMatrix $ y * log h + (1-y) * log (1-h)
          cost2 = lambda/2 * (normMatrix theta0 + normMatrix theta1)
  
          grad  = (1/m) `scale` (grad1 + grad2)
          grad1 = toVector delta0 delta1
          grad2 = lambda `scale` toVector (insertNils theta0) (insertNils theta1)
  
          normMatrix m = sumMatrix $ (dropRows 1 m) ^ 2
          insertNils m = vertcat [0, dropRows 1 m]
  
  nnCost shape y x lambda = fst . nnCostGradient shape y x lambda
  nnGrad shape y x lambda = snd . nnCostGradient shape y x lambda
  
  -- |Use central differencing to compute an approximation to the gradient
  --  vector for a neural network. This is mainly used for checking the
  --  implementation of backprop.
  nnGradApprox :: NNShape -> Matrix Double -> Matrix Double -> Double -> Vector Double -> Vector Double
  nnGradApprox shape y x lambda vec = fromList $ g `map` [0..n-1]
      where
          h = 1e-6
          n = dim vec
          f v = nnCost shape y x lambda v
          g i = (f (vec + e i) - f (vec - e i)) / (2*h)
          e i = fromList $ replicate i 0 ++ [h] ++ replicate (n-i-1) 0
  
  -------------
  -- Testing --
  -------------
  
  testNN :: NeuralNetwork
  testNN = NN t0 t1
      where t0 = fromLists [[11.9934, -5.1396], [-7.7162, 10.1512], [-7.668, 10.1835]]
            t1 = fromLists [[-16.8806], [10.0445], [8.7476]]
  
  testFwdProp :: IO ()
  testFwdProp = do
      putStrLn "***\nCompare forward propagation to the MATLAB implementation.\n"
      let theta0 = fromLists [[0], [-10]]
          theta1 = fromLists [[5],[-10]]
          nn = NN theta0 theta1
          x = fromLists [[-1],[0],[1]]
      -- sigmoid [10, 0, -10] = [1, 0.5, 0]
      -- sigmoid [-5, 0, 5]  = [0.0, 0.5, 1]
      let y = nnPredict nn x
      putStrLn "Predictions (should be roughly 0.0, 0.5, 1.0)"
      disp 2 y
      
  testBackProp :: IO ()
  testBackProp = do
      putStrLn "***\nCompare back propagation to the MATLAB implementation.\n"
      let nn = testNN
          x  = fromLists [[0, 0], [0, 1], [1, 0], [1, 1]]
          y  = fromLists [[0], [1], [1], [0]]
          (h, delta0, delta1) = nnFwdBackProp nn y x
      putStrLn "Predictions (should be 0.0011, 0.8487, 0.8476, 0.0004)"
      disp 4 h
      putStrLn "Delta 0 (should be -0.0401, -0.0171, -0.0205, -0.0088, -0.0195, -0.0084)"
      disp 4 delta0
      putStrLn "Delta 1 (should be -0.3022, -0.2985, -0.3013)"
      disp 4 delta1
  
  testCostGradient :: IO ()
  testCostGradient = do
      putStrLn "***\nTest cost/gradient against the MATLAB implementation.\n"
      let nn@(NN theta0 theta1) = testNN
          x = fromLists [[0, 0], [0, 1], [1, 0], [1, 1]]
          y = fromLists [[0], [1], [1], [0]]
          lambda = 1e-4
          (cost, grad) = nnCostGradient (2,2,1) y x lambda (toVector theta0 theta1)
      putStrLn "Cost (should be around 0.0890)"
      print cost
      putStrLn "Gradient (should be -0.0100, -0.0043, -0.0053, -0.0019, -0.0051, -0.0019, -0.0755, -0.0744, -0.0751)"
      disp 4 (asColumn grad)
  
  test :: Int -> Double -> IO ()
  test n lambda = do
      putStrLn "***\nLearning XOR function.\n"
      x <- rand n 2
      e <- fmap (0.01*) (rand n 1)
      let y = mapMatrix (\x -> if x > 0.5 then 1.0 else 0.0) (xor x)
      nn <- nnTrainIO 4 y x lambda
      let ypred = nnPredict nn x
      -- Show predictions
      putStrLn "Predictions:"
      disp 2 $ takeRows 10 $ horzcat [x, y, ypred]
      -- Show neural net
      print nn
      -- Final test; should approximately compute xor function
      let xx = fromLists [[0,0],[0,1],[1,0],[1,1]]
          yy = nnPredict nn xx
      putStrLn "Exclusive or:"
      disp 2 $ horzcat [xx,yy]
      
  xor :: Matrix Double -> Matrix Double
  xor x = let [u,v] = toColumns x in asColumn (u + v - 2 * u * v)
      
  