{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- | Average Perceptron implementation of Part of speech tagging,
-- adapted for Haskell from this python implementation, which is described on the blog post:
--
--  * <http://honnibal.wordpress.com/2013/09/11/a-good-part-of-speechpos-tagger-in-about-200-lines-of-python/>
--
-- The Perceptron code can be found on github:
--
--  * <https://github.com/sloria/TextBlob/blob/dev/text/_perceptron.py>
--
module NLP.ML.AvgPerceptron
  ( Perceptron(..)
  , Class(..)
  , Weight
  , Feature(..)
  , emptyPerceptron
  , predict
  , train
  , update
  , averageWeights
  )
where

import Control.DeepSeq (NFData)
import Data.List (foldl')
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Data.Maybe (fromMaybe)
import Data.Serialize (Serialize)
import Data.Text (Text)
import System.Random.Shuffle (shuffleM)
import GHC.Generics

import NLP.Types ()

newtype Feature = Feat Text
    deriving (Read, Show, Eq, Ord, Generic, NFData)

instance Serialize Feature

-- | The classes that the perceptron assigns are represnted with a
-- newtype-wrapped String.
--
-- Eventually, I think this should become a typeclass, so the classes
-- can be defined by the users of the Perceptron (such as custom POS
-- tag ADTs, or more complex classes).
newtype Class = Class String
    deriving (Read, Show, Eq, Ord, Generic, NFData)

instance Serialize Class

-- | Typedef for doubles to make the code easier to read, and to make
-- this simple to change if necessary.
type Weight = Double

infinity :: Weight
infinity = recip 0

-- | An empty perceptron, used to start training.
emptyPerceptron :: Perceptron
emptyPerceptron = Perceptron { weights = Map.empty
                             , totals = Map.empty
                             , tstamps = Map.empty
                             , instances = 0 }

-- | The perceptron model.
data Perceptron = Perceptron {
    -- | Each feature gets its own weight vector, so weights is a
    -- dict-of-dicts
    weights :: Map Feature (Map Class Weight)

    -- | The accumulated values, for the averaging. These will be
    -- keyed by feature/clas tuples
    , totals :: Map (Feature, Class) Weight

    -- | The last time the feature was changed, for the averaging. Also
    -- keyed by feature/clas tuples
    -- (tstamps is short for timestamps)
    , tstamps :: Map (Feature, Class) Int

    -- | Number of instances seen
    , instances :: Int
    } deriving (Read, Show, Eq, Generic)

instance Serialize Perceptron
instance NFData Perceptron

incrementInstances :: Perceptron -> Perceptron
incrementInstances p = p { instances = 1 + (instances p) }

getTimestamp :: Perceptron -> (Feature, Class) -> Int
getTimestamp p param = Map.findWithDefault 0 param (tstamps p)

getTotal :: Perceptron -> (Feature, Class) -> Weight
getTotal p param = Map.findWithDefault 0 param (totals p)

getFeatureWeight :: Perceptron -> Feature -> Map Class Weight
getFeatureWeight p f = Map.findWithDefault Map.empty f (weights p)

-- | Predict a class given a feature vector.
--
-- Ported from python:
--
-- > def predict(self, features):
-- >     '''Dot-product the features and current weights and return the best label.'''
-- >     scores = defaultdict(float)
-- >     for feat, value in features.items():
-- >         if feat not in self.weights or value == 0:
-- >             continue
-- >         weights = self.weights[feat]
-- >         for label, weight in weights.items():
-- >             scores[label] += value * weight
-- >     # Do a secondary alphabetic sort, for stability
-- >     return max(self.classes, key=lambda label: (scores[label], label))
--
predict :: Perceptron -> Map Feature Int -> Maybe Class
predict per features = -- find highest ranked score in finalScores:
    -- trace ("features: "++ show features ++ "\n") $
    -- trace ("finalScores: "++ show sortedScores ++ "\n") $
    sortedScores
    where
      sortedScores :: Maybe Class
      sortedScores = fst $ Map.foldlWithKey ranker (Nothing, negate infinity) finalScores

      ranker r@(_, ow) nc nw | nw > ow   = (Just nc, nw)
                             | otherwise = r

      finalScores :: Map Class Weight
      finalScores = Map.foldlWithKey fn Map.empty features

      fn :: Map Class Weight -> Feature -> Int -> Map Class Weight
      fn scores f  v
         | v > 0 = case Map.lookup f (weights per) of
                     Just vec -> Map.foldlWithKey (doProd v) scores vec
                     Nothing  -> scores
         | otherwise = scores

      doProd :: Int -> Map Class Weight -> Class -> Weight -> Map Class Weight
      doProd value scores label weight =
        Map.alter (updater (weight * (fromIntegral value))) label scores

      updater :: Weight -> Maybe Weight -> Maybe Weight
      updater newVal Nothing  = Just newVal
      updater newVal (Just v) = Just (v + newVal)

-- | Update the perceptron with a new example.
--
-- > update(self, truth, guess, features)
-- >    ...
-- >         self.i += 1
-- >         if truth == guess:
-- >             return None
-- >         for f in features:
-- >             weights = self.weights.setdefault(f, {}) -- setdefault is Map.findWithDefault, and destructive.
-- >             upd_feat(truth, f, weights.get(truth, 0.0), 1.0)
-- >             upd_feat(guess, f, weights.get(guess, 0.0), -1.0)
-- >         return None
--
update :: Perceptron -> Class -> Class -> [Feature] -> Perceptron
update per truth guess features
  | truth == guess = incrementInstances per
  | otherwise      = foldr loopBody per features
    where
      loopBody :: Feature -> Perceptron -> Perceptron
      loopBody f p = let
        fweights  = getFeatureWeight p f
        cweight c = Map.findWithDefault 0 c fweights
        in upd_feat guess f (cweight guess) (-1)
             (upd_feat truth f (cweight truth) 1 p)

-- | ported from python:
--
-- > def update(self, truth, guess, features):
-- >        '''Update the feature weights.'''
-- >        def upd_feat(c, f, w, v):
-- >           param = (f, c)
-- >           self._totals[param] += (self.i - self._tstamps[param]) * w
-- >           self._tstamps[param] = self.i
-- >           self.weights[f][c] = w + v
upd_feat :: Class -> Feature -> Weight -> Weight -> Perceptron -> Perceptron
upd_feat c f w v p = let
    newInstances = 1 + (instances p) -- increment the instance counter.

    -- self._totals[param] += (self.i - self._tstamps[param]) * w
    paramTstamp = newInstances - getTimestamp p (f, c)
    tmpTotal  = (getTotal p (f, c)) + ((fromIntegral paramTstamp) * w)
    newTotals = Map.insert (f, c) tmpTotal (totals p)

    -- self._tstamps[param] = self.i
    newTstamps = Map.insert (f, c) newInstances (tstamps p)

    -- self.weights[f][c] = w + v
    newWeights = Map.insert f (Map.insert c (w + v) (getFeatureWeight p f)) (weights p)

    in p { totals = newTotals
         , tstamps = newTstamps
         , weights = newWeights }


-- | Average the weights
--
-- Ported from Python:
--
-- > def average_weights(self):
-- >     for feat, weights in self.weights.items():
-- >         new_feat_weights = {}
-- >         for clas, weight in weights.items():
-- >             param = (feat, clas)
-- >             total = self._totals[param]
-- >             total += (self.i - self._tstamps[param]) * weight
-- >             averaged = round(total / float(self.i), 3)
-- >             if averaged:
-- >                 new_feat_weights[clas] = averaged
-- >         self.weights[feat] = new_feat_weights
-- >     return None
--
averageWeights :: Perceptron -> Perceptron
averageWeights per | (instances per) == 0 = per
                   | otherwise = per { weights = Map.mapWithKey avgWeights $ weights per
                                     , tstamps = Map.mapWithKey (\_ _->0) (tstamps per)
                                     , instances = 0 }
  where
    avgWeights :: Feature -> Map Class Weight -> Map Class Weight
    avgWeights feat ws = Map.foldlWithKey (doAvg feat) Map.empty ws

    doAvg :: Feature -> Map Class Weight -> Class -> Weight -> Map Class Weight
    doAvg f acc c w = let
      param = (f, c)
      paramTotal = instances per - getTimestamp per param

      total :: Weight
      total = (getTotal per param) + ((fromIntegral paramTotal) * w)
      averaged = roundTo 3 (total / (fromIntegral $ instances per))
      in if 0 == averaged
           then acc
           else Map.insert c averaged acc

-- | round a fractional number to a specified decimal place.
--
-- >>> roundTo 2 3.1459
-- 3.15
--
roundTo :: RealFrac a => Int -> a -> a
roundTo n f = (fromInteger $ round $ f * (10^n)) / (10.0^^n)


-- Train a perceptron
--
-- Ported from Python:
-- > def train(nr_iter, examples):
-- >     model = Perceptron()
-- >     for i in range(nr_iter):
-- >         random.shuffle(examples)
-- >         for features, class_ in examples:
-- >             scores = model.predict(features)
-- >             guess, score = max(scores.items(), key=lambda i: i[1])
-- >             if guess != class_:
-- >                 model.update(class_, guess, features)
-- >     model.average_weights()
-- >     return model
train :: Int -> Perceptron -> [(Map Feature Int, Class)] -> IO Perceptron
train itr model exs = do
  trainingSet <- shuffleM $ concat $ take itr $ repeat exs
  return $ averageWeights $ foldl' trainEx model trainingSet

trainEx :: Perceptron -> (Map Feature Int, Class) -> Perceptron
trainEx model (feats, truth) = let
  guess = fromMaybe (Class "Unk") $ predict model feats
  in update model truth guess $ Map.keys feats