----------------------------------------------------
-- |
-- Module     :  AI.Network
-- License    :  GPL
--
-- Maintainer :  Kiet Lam <ktklam9@gmail.com>
--
--
-- This module provides a generic module for
-- initiialization and training of neural networks
--
-- User must provide the needed functions
--
--
----------------------------------------------------


module AI.Model.GenericModel (
  GenericModel(..),
  initializeModel,
  getOutput,
  trainModel
  ) where

import System.Random
import Data.Packed.Vector
import Data.Packed.Matrix

import AI.Signatures
import AI.Calculation
import AI.Network
import AI.Training


-- | Generic neural network model for expansion
data GenericModel = GenericModel
                    {
                      cost :: Cost,   -- ^ The cost model of the model
                      net  :: Network -- ^ The neural network to be used for modeling
                    }


-- | Initialize neural network model with the weights
-- randomized within [-1.0,1.0]
initializeModel :: Activation      -- ^ The activation model of each neuron
                   -> Cost         -- ^ The cost model of the output neurons
                                   --   compared to the expected output
                   -> [Int]        -- ^ The architecture of the network
                                   --   e.g., a 2-3-1 architecture would be [2,3,1]
                   -> Double       -- ^ The regularization constant
                                   --   should be 0 if you do not want regularization
                   -> StdGen       -- ^ The random generator
                   -> GenericModel -- ^ Returns the initialized model
initializeModel ac co arch la gen =
  let n = foldl (+) 0 [((x + 1) * y) | (x,y) <- zip arch (tail arch)]
      ws = (fromList . take n) (randomRs (-1.0, 1.0) gen :: [Double])
  in
   GenericModel { cost = co,
                  net = Network { activation = getActivation ac,
                                  derivative = getDerivative ac,
                                  lambda = la,
                                  weights = ws,
                                  architecture = arch
                                }
                }


-- | Get the output of the model
getOutput :: GenericModel     -- ^ The model of interest
             -> Vector Double -- ^ The input vector to the input layer
             -> Vector Double -- ^ The output of the network model
getOutput (GenericModel {net = nn}) input = networkOutput nn input


-- | Train the model given the parameters and the training algorithm
trainModel :: GenericModel          -- ^ The model to be trained
              -> TrainingAlgorithm  -- ^ The training algorithm to be used
              -> Double             -- ^ The precision to train with regards to
                                    --   the cost function
              -> Int                -- ^ The maximum amount of epochs to train
              -> Matrix Double      -- ^ The input matrix
              -> Matrix Double      -- ^ The expected output matrix
              -> GenericModel       -- ^ Returns the trained model
trainModel (GenericModel {cost = co, net = nn}) algo prec epochs inMat exMat =
  let trainedNet = trainNetwork algo co backpropagation nn prec epochs inMat exMat in
  GenericModel { net = trainedNet,
                 cost = co
               }