content: {-# LANGUAGE OverloadedStrings #-}
  -- | Avegeraged Perceptron Tagger
  --
  -- Adapted from the python implementation found here:
  --
  --  * <https://github.com/sloria/textblob-aptagger/blob/master/textblob_aptagger/taggers.py>
  --
  module NLP.POS.AvgPerceptronTagger
    ( mkTagger
    , trainNew
    , trainOnFiles
    , train
    , trainInt
    , tag
    , tagSentence
    , emptyPerceptron
    , taggerID
    , readTagger
    )
  where
  
  import NLP.Corpora.Parsing (readPOSWith)
  import NLP.ML.AvgPerceptron ( Perceptron, Feature(..)
                              , Class(..), predict, update
                              , emptyPerceptron, averageWeights)
  import NLP.Types
  
  import Control.Monad (foldM)
  import Data.ByteString (ByteString)
  import Data.ByteString.Char8 (pack)
  import Data.List (zipWith4, foldl')
  import Data.Map.Strict (Map)
  import qualified Data.Map.Strict as Map
  import Data.Maybe (fromMaybe)
  import Data.Serialize (encode, decode)
  import Data.Text (Text)
  import qualified Data.Text as T
  import qualified Data.Text.IO as T
  
  import NLP.Tokenize.Chatter (tokenize)
  import NLP.FullStop (segment)
  import System.Random.Shuffle (shuffleM)
  
  taggerID :: ByteString
  taggerID = pack "NLP.POS.AvgPerceptronTagger"
  
  readTagger :: Tag t => ByteString -> Maybe (POSTagger t) -> Either String (POSTagger t)
  readTagger bs backoff = do
    model <- decode bs
    return $ mkTagger model backoff
  
  -- | Create an Averaged Perceptron Tagger using the specified back-off
  -- tagger as a fall-back, if one is specified.
  --
  -- This uses a tokenizer adapted from the 'tokenize' package for a
  -- tokenizer, and Erik Kow's fullstop sentence segmenter
  -- (<http://hackage.haskell.org/package/fullstop>) as a sentence
  -- splitter.
  mkTagger :: Tag t => Perceptron -> Maybe (POSTagger t) -> POSTagger t
  mkTagger per mTgr = POSTagger { posTagger  = tag per
                                , posTrainer = \exs -> do
                                    newPer <- trainInt itterations per exs
                                    return $ mkTagger newPer mTgr
                                , posBackoff = mTgr
                                , posTokenizer = tokenize
                                , posSplitter = (map T.pack) . segment . T.unpack
                                , posSerialize = encode per
                                , posID = taggerID
                                }
  
  itterations :: Int
  itterations = 5
  
  -- | Train a new 'Perceptron'.
  --
  -- The training corpus should be a collection of sentences, one
  -- sentence on each line, and with each token tagged with a part of
  -- speech.
  --
  -- For example, the input:
  --
  -- > "The/DT dog/NN jumped/VB ./.\nThe/DT cat/NN slept/VB ./."
  --
  -- defines two training sentences.
  --
  -- >>> tagger <- trainNew "Dear/jj Sirs/nns :/: Let/vb\nUs/nn begin/vb\n"
  -- >>> tag tagger $ map T.words $ T.lines "Dear sir"
  -- "Dear/jj Sirs/nns :/: Let/vb"
  --
  trainNew :: Tag t => (Text -> t) -> Text -> IO Perceptron
  trainNew parser rawCorpus = train parser emptyPerceptron rawCorpus
  
  -- | Train a new 'Perceptron' on a corpus of files.
  trainOnFiles :: Tag t => (Text -> t) -> [FilePath] -> IO Perceptron
  trainOnFiles parser corpora = foldM step emptyPerceptron corpora
    where
      step :: Perceptron -> FilePath -> IO Perceptron
      step per path = do
        content <- T.readFile path
        train parser per content
  
  -- | Add training examples to a perceptron.
  --
  -- >>> tagger <- train emptyPerceptron "Dear/jj Sirs/nns :/: Let/vb\nUs/nn begin/vb\n"
  -- >>> tag tagger $ map T.words $ T.lines "Dear sir"
  -- "Dear/jj Sirs/nns :/: Let/vb"
  --
  -- If you're using multiple input files, this can be useful to improve
  -- performance (by folding over the files).  For example, see `trainOnFiles`
  --
  train :: Tag t => (Text -> t) -- ^ The POS tag parser.
        -> Perceptron -- ^ The inital model.
        -> Text       -- ^ Training data; formatted with one sentence
                      -- per line, and standard POS tags after each
                      -- space-delimeted token.
        -> IO Perceptron
  train parse per rawCorpus = do
    let corpora = map (readPOSWith parse) $ T.lines rawCorpus
    trainInt itterations per corpora
  
  -- | start markers to ensure all features in context are valid,
  -- even for the first "real" tokens.
  startToks :: [Token]
  startToks = [Token "-START-", Token "-START2-"]
  
  -- | end markers to ensure all features are valid, even for
  -- the last "real" tokens.
  endToks :: [Token]
  endToks = [Token "-END-", Token "-END2-"]
  
  -- | Tag a document (represented as a list of 'Sentence's) with a
  -- trained 'Perceptron'
  --
  -- Ported from Python:
  --
  -- > def tag(self, corpus, tokenize=True):
  -- >     '''Tags a string `corpus`.'''
  -- >     # Assume untokenized corpus has \n between sentences and ' ' between words
  -- >     s_split = nltk.sent_tokenize if tokenize else lambda t: t.split('\n')
  -- >     w_split = nltk.word_tokenize if tokenize else lambda s: s.split()
  -- >     def split_sents(corpus):
  -- >         for s in s_split(corpus):
  -- >             yield w_split(s)
  -- >      prev, prev2 = self.START
  -- >     tokens = []
  -- >     for words in split_sents(corpus):
  -- >         context = self.START + [self._normalize(w) for w in words] + self.END
  -- >         for i, word in enumerate(words):
  -- >             tag = self.tagdict.get(word)
  -- >             if not tag:
  -- >                 features = self._get_features(i, word, context, prev, prev2)
  -- >                 tag = self.model.predict(features)
  -- >             tokens.append((word, tag))
  -- >             prev2 = prev
  -- >             prev = tag
  -- >     return tokens
  --
  tag :: Tag t => Perceptron -> [Sentence] -> [TaggedSentence t]
  tag per corpus = map (tagSentence per) corpus
  
  -- | Tag a single sentence.
  tagSentence :: Tag t => Perceptron -> Sentence -> TaggedSentence t
  tagSentence per sent = let
  
    tags = (map tokenToClass startToks) ++ map (predictPos per) features
  
    features = zipWith4 (getFeatures sent)
               [0..]
               (tokens sent)
               (tail tags)
               tags
  
    in applyTags sent (map (\(Class c) -> parseTag $ T.pack c) $ drop 2 tags)
  
  -- | Train a model from sentences.
  --
  -- Ported from Python:
  --
  -- > def train(self, sentences, save_loc=None, nr_iter=5):
  -- >     self._make_tagdict(sentences)
  -- >     self.model.classes = self.classes
  -- >     prev, prev2 = START
  -- >     for iter_ in range(nr_iter):
  -- >         c = 0
  -- >         n = 0
  -- >         for words, tags in sentences:
  -- >             context = START + [self._normalize(w) for w in words] + END
  -- >             for i, word in enumerate(words):
  -- >                 guess = self.tagdict.get(word)
  -- >                 if not guess:
  -- >                     feats = self._get_features(i, word, context, prev, prev2)
  -- >                     guess = self.model.predict(feats)
  -- >                     self.model.update(tags[i], guess, feats)
  -- >                 prev2 = prev; prev = guess
  -- >                 c += guess == tags[i]
  -- >                 n += 1
  -- >         random.shuffle(sentences)
  -- >         logging.info("Iter {0}: {1}/{2}={3}".format(iter_, c, n, _pc(c, n)))
  -- >     self.model.average_weights()
  -- >     # Pickle as a binary file
  -- >     if save_loc is not None:
  -- >         pickle.dump((self.model.weights, self.tagdict, self.classes),
  -- >                      open(save_loc, 'wb'), -1)
  -- >     return None
  --
  trainInt :: Tag t =>
              Int -- ^ The number of times to iterate over the training
                  -- data, randomly shuffling after each iteration. (@5@
                  -- is a reasonable choice.)
           -> Perceptron -- ^ The 'Perceptron' to train.
           -> [TaggedSentence t] -- ^ The training data. (A list of @[(Text, Tag)]@'s)
           -> IO Perceptron    -- ^ A trained perceptron.  IO is needed
                               -- for randomization.
  trainInt itr per examples = trainCls itr per $ toClassLst $ map unzipTags examples
  
  toClassLst :: Tag t => [(Sentence, [t])] -> [(Sentence, [Class])]
  toClassLst tagged = map (\(x, y)->(x, map (Class . T.unpack . fromTag) y)) tagged
  
  trainCls :: Int -> Perceptron -> [(Sentence, [Class])] -> IO Perceptron
  trainCls itr per examples = do
    trainingSet <- shuffleM $ concat $ take itr $ repeat examples
    return $ averageWeights $ foldl' trainSentence per trainingSet
  
  tokenToClass :: Token -> Class
  tokenToClass = Class . T.unpack . showTok
  
  -- | Train on one sentence.
  --
  -- Adapted from this portion of the Python train method:
  --
  -- >             context = START + [self._normalize(w) for w in words] + END
  -- >             for i, word in enumerate(words):
  -- >                 guess = self.tagdict.get(word)
  -- >                 if not guess:
  -- >                     feats = self._get_features(i, word, context, prev, prev2)
  -- >                     guess = self.model.predict(feats)
  -- >                     self.model.update(tags[i], guess, feats)
  -- >                 prev2 = prev; prev = guess
  -- >                 c += guess == tags[i]
  -- >                 n += 1
  trainSentence :: Perceptron -> (Sentence, [Class]) -> Perceptron
  trainSentence per (sent, ts) = let
  
    tags = (map tokenToClass startToks) ++ ts ++ (map tokenToClass endToks)
  
    features = zipWith4 (getFeatures sent)
                           [0..] -- index
                           (tokens sent)  -- words
                           (tail tags) -- prev1
                           tags  -- prev2
  
    fn :: Perceptron -> (Map Feature Int, Class) -> Perceptron
    fn model (feats, truth) = let
      guess = predictPos model feats
      in update model truth guess $ Map.keys feats
  
    in foldl' fn per (zip features ts)
  
  -- | Predict a Part of Speech, defaulting to the @Unk@ tag, if no
  -- classification is found.
  predictPos :: Perceptron -> Map Feature Int -> Class
  predictPos model feats = fromMaybe (Class "Unk") $ predict model feats
  
  -- | Default feature set.
  --
  -- > def _get_features(self, i, word, context, prev, prev2):
  -- >     '''Map tokens into a feature representation, implemented as a
  -- >     {hashable: float} dict. If the features change, a new model must be
  -- >     trained.
  -- >     '''
  -- >     def add(name, *args):
  -- >         features[' '.join((name,) + tuple(args))] += 1
  -- >      i += len(self.START)
  -- >     features = defaultdict(int)
  -- >     # It's useful to have a constant feature, which acts sort of like a prior
  -- >     add('bias')
  -- >     add('i suffix', word[-3:])
  -- >     add('i pref1', word[0])
  -- >     add('i-1 tag', prev)
  -- >     add('i-2 tag', prev2)
  -- >     add('i tag+i-2 tag', prev, prev2)
  -- >     add('i word', context[i])
  -- >     add('i-1 tag+i word', prev, context[i])
  -- >     add('i-1 word', context[i-1])
  -- >     add('i-1 suffix', context[i-1][-3:])
  -- >     add('i-2 word', context[i-2])
  -- >     add('i+1 word', context[i+1])
  -- >     add('i+1 suffix', context[i+1][-3:])
  -- >     add('i+2 word', context[i+2])
  -- >     return features
  --
  getFeatures :: Sentence -> Int -> Token -> Class -> Class -> Map Feature Int
  getFeatures ctx idx word prev prev2 = let
    context = startToks ++ tokens ctx ++ endToks
  
    i = idx + length startToks
  
    add :: Map Feature Int -> [Text] -> Map Feature Int
    add m args = Map.alter increment (mkFeature $ T.intercalate " " args) m
  
    increment :: Maybe Int -> Maybe Int
    increment Nothing  = Just 1
    increment (Just w) = Just (w + 1)
  
    features :: [[Text]]
    features = [ ["bias", ""]
               , ["i suffix", suffix word ]
               , ["i pref1", T.take 1 $ showTok word ]
               , ["i-1 tag", T.pack $ show prev ]
               , ["i-2 tag", T.pack $ show prev2 ]
               , ["i tag+i-2 tag", T.pack $ show prev, T.pack $ show prev2 ]
               , ["i word",     showTok (context!!i) ]
               , ["i-1 tag+i word", T.pack $ show prev, showTok (context!!i) ]
               , ["i-1 word",   showTok (context!!(i-1)) ]
               , ["i-1 suffix",  suffix (context!!(i-1)) ]
               , ["i-2 word",   showTok (context!!(i-2)) ]
               , ["i+1 word",   showTok (context!!(i+1)) ]
               , ["i+1 suffix",  suffix (context!!(i+1)) ]
               , ["i+2 word",   showTok (context!!(i+2)) ]
               ]
    -- in trace ("getFeatures: "++show (ctx, idx, word, prev, prev2)) $
    in foldl' add Map.empty features
  
  mkFeature :: Text -> Feature
  mkFeature txt = Feat $ T.copy txt
  