-----------------------------------------------------------------------------
--
-- Module      :  Main
-- Copyright   :
-- License     :  AllRightsReserved
--
-- Maintainer  :
-- Stability   :
-- Portability :
--
-- |
--
-----------------------------------------------------------------------------

{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE BangPatterns #-}


module Main (
    main
) where


import Debug.Trace
--import Data.Set
--import qualified Data.Set as S
import qualified Data.Set
import qualified Data.Map
import qualified Data.List
import qualified IO
import qualified Data.Vector as V
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.List.Zipper as Z
import qualified Data.Maybe as DM
import Control.Monad


-- ------------------------
-- Parts of Speech Tagging
-- ------------------------

type Token = String
type Tag = String

data TrainingInstance = TrainingInstance Token Tag
                        deriving Show

-- Splits a list @l@ whe it encounters a seperator @sep@
rsplit :: Eq a => a -> [a] -> ([a], [a])
rsplit sep l = let (ps, xs, _) = rsplit_ sep l in
               (ps, xs)

-- Helper function that scans the string from the right to the left
-- in order to find the last occurrence of the seperator in the
rsplit_ :: Eq a => a -> [a] -> ([a], [a], Bool)
rsplit_ sep = foldr (splitFun sep) ([], [], False)
  where splitFun _ e (px, xs, True) = (e:px, xs, True)
        splitFun sep e (px, xs, False)
                 | e == sep = (px, xs, True)
                 | otherwise = (px, e:xs, False)

{- Tests
rsplit '/' "the/AT"
rsplit '/' "a/b/TEST"
-}

-- Constructor of a @TrainingInstance@. Takes a ngram/token
-- and split it by the last ocurrence of the '/' separator
toTrainingInstance :: String -> TrainingInstance
toTrainingInstance s = let (token, tag) = rsplit '/' s in
                       TrainingInstance token tag

-- Goes through all the ngrams/tokens that have been previously
-- tagged and records the frequency with which each token appears
-- for each of its tags (one token may be assigned multiple tags)
tokenTagFreqs :: [TrainingInstance] -> M.Map Token (M.Map Tag Integer)
tokenTagFreqs = L.foldl' countWord M.empty
  where countWord m (TrainingInstance token tag) =
                    M.insertWith (countTag tag) token (M.singleton tag 1) m
        countTag tag _ old = M.insertWith (\newFreq oldFreq -> newFreq + oldFreq) tag 1 old

{-
h <- IO.openFile "brown-pos-train.txt" IO.ReadMode
c <- IO.hGetContents h
tokenTagFreqs $ take 10 $ map toTrainingInstance $ words c
-}

tokenMostFreqTag :: M.Map k (M.Map Tag Integer) -> M.Map k Tag
tokenMostFreqTag = M.map mostFreqTag
    where
      mostFreqTag = fst . M.foldlWithKey maxTag ("NIL", 0)
      maxTag acc@(_maxTag, maxFreq) tag freq
                 | freq > maxFreq = (tag, freq)
                 | otherwise = acc

-- Scans a list of @TrainingInstance@ and selects
trainFreqTagger :: [TrainingInstance] -> M.Map Token Tag
trainFreqTagger = tokenMostFreqTag . tokenTagFreqs

{-
tokenMostFreqTag $ tokenTagFreqs [TrainingInstance "a" "A", TrainingInstance "a" "B", TrainingInstance "a" "A"]
trainFreqTagger [TrainingInstance "a" "A", TrainingInstance "a" "B", TrainingInstance "a" "A"]
-}

freqTagWord :: M.Map Token Tag -> Token -> Maybe Tag
freqTagWord m t = M.lookup t m

{-
h <- IO.openFile "brown-pos-train.txt" IO.ReadMode
c <- IO.hGetContents h
let model = trainFreqTagger $ map toTrainingInstance $ words c
model
map (freqTagWord model) ["The","cat","is","on","the","mat","."]
map (freqTagWord model) ["That's","right",",","the","mascara", "snake",".","Fast","and","bulbous",".","Also","a","tinned","teardrop","."]
-}

testTagger _ = do
  h <- IO.openFile "brown-pos-train.txt" IO.ReadMode
  c <- IO.hGetContents h
  let model = trainFreqTagger $ map toTrainingInstance $ words c
  putStrLn "Training done"
  let r1 = map (freqTagWord model) ["The","cat","is","on","the","mat","."]
  putStrLn "Result 1"
  print r1
  let r2 = map (freqTagWord model) ["That's","right",",","the","mascara", "snake",".","Fast","and","bulbous",".","Also","a","tinned","teardrop","."]
  putStrLn "Result 2"
  print r2


-- Evaluates the tagger @tagFun@ this tagger must tag a token/unigram
-- All tokens are counted (1st element of tuple). If tag is correct this
-- is counted (second element of the tuple). If it is incorrect then
-- the this is also counted (3rd element of the tuple)
evalTagger :: (Num t) => (Token -> Maybe Tag) -> [TrainingInstance] -> (t, t, t)
evalTagger tagFun = L.foldl' eval (0,0,0)
  where
    eval (n,c,u) (TrainingInstance token correctTag) =
      case tagFun token of
      Just tag -> if tag == correctTag then (n+1,c+1,u)
                                       else (n+1,c,u+1)
      Nothing -> (n+1,c,u+1)

-- Used to tag all tokens with the same @tag@
-- A baseline tagger can be used by tagging all tokens
-- witht the most frequent tag "NN" (common noun)
baselineTagger :: Tag -> Token -> Maybe Tag
baselineTagger tag _ = Just tag

-- This tagger wraps another tagger. If the inner tagger
-- does (cannot) not tag the token, then a constant tag
-- is assigned. Usually we assing the most common tag "NN"
backoffTagger :: (Token -> Maybe Tag) -> Tag -> Token -> Maybe Tag
backoffTagger f bt t = let pick = f t in
                       case pick of
                         Just tag -> Just tag
                         Nothing  -> Just bt

tagginAccuracy (n,c,_) = (c / n) * 100

evalTaggers _ = do
  h <- IO.openFile "brown-pos-train.txt" IO.ReadMode
  c <- IO.hGetContents h
  let model = trainFreqTagger $ map toTrainingInstance $ words c
  putStrLn "Training finished"
  i <- IO.openFile "brown-pos-test.txt" IO.ReadMode
  d <- IO.hGetContents i
  let r1 = evalTagger (baselineTagger "NN") $ map toTrainingInstance $ words d
  putStrLn "Baseline result done."
  let r2 = evalTagger (freqTagWord model) $ map toTrainingInstance $ words d
  putStrLn "Frequent Tag result done"
  let r3 = evalTagger (backoffTagger (freqTagWord model) "NN") $ map toTrainingInstance $ words d
  putStrLn "Back-off tag result done"
  putStrLn "--------------------------------------------"
  putStrLn " Tagger                     |  Accuracy (%) "
  putStrLn "--------------------------------------------"
  putStr   " Baseline                   |"
  print $ tagginAccuracy r1
  putStr   " Frequencey Based           |"
  print $ tagginAccuracy r2
  putStr   " Frequencey Based + Backoff |"
  print $ tagginAccuracy r3



-- ------------------------------------------------------------------
-- Brill parts of speech tagger
-- A simple rule-based part of speech tagger. E. Brill. 1992.
-- Association for Computational Linguistics.
-- ------------------------------------------------------------------

-- First tag is the incorrrect Tag
-- Second tag is the correct Tag
data Replacement = Replacement Tag Tag
                   deriving (Eq, Ord, Show)

data TransformationRule =
      NextTagRule     Replacement Tag
    | PrevTagRule     Replacement Tag
    | SurroundTagRule Replacement Tag Tag
      deriving (Eq, Ord, Show)


rightCursor :: Z.Zipper a -> Maybe a
rightCursor = Z.safeCursor . Z.right

leftCursor :: Z.Zipper a -> Maybe a
leftCursor z = if Z.beginp z then
                   Nothing
               else
                   Z.safeCursor $ Z.left z

---- Generate the next tag rule
instNextTagRule :: Z.Zipper (Tag, Tag) -> Maybe TransformationRule
instNextTagRule z = do
    (_, next) <- rightCursor z
    (correct, incorrect) <- Z.safeCursor z
    return $ NextTagRule (Replacement incorrect correct) next

---- Generate the previous tag rule
instPrevTagRule :: Z.Zipper (Tag, Tag) -> Maybe TransformationRule
instPrevTagRule z = do
    (_, prev)            <- leftCursor z
    (correct, incorrect) <- Z.safeCursor z
    return $ PrevTagRule (Replacement incorrect correct) prev

---- Generate the surround tag rule
instSurroundTagRule :: Z.Zipper (Tag, Tag) -> Maybe TransformationRule
instSurroundTagRule z = do
    (_, next)            <- rightCursor z
    (_, prev)            <- leftCursor z
    (correct, incorrect) <- Z.safeCursor z
    return $ SurroundTagRule (Replacement incorrect correct) prev next


-- Initial learning state uses the miost frequent + back-off taggers
initialLearningState :: [TrainingInstance] -> Z.Zipper (Tag, Tag)
initialLearningState train = Z.fromList $ zip (correct train) (proposed train)
    where proposed    = map tagger . trainTokens
          correct     = map (\(TrainingInstance _ tag) -> tag)
          tagger      = DM.fromJust . backoffTagger (freqTagWord model) "NN"
          trainTokens = map (\(TrainingInstance token _) -> token)
          model       = trainFreqTagger train


-- Apply the learning rules to incorrectly tagged tokens
-- Count the number of times the rule successfully tagged the token
instRules :: [(Z.Zipper (Tag, Tag) -> Maybe TransformationRule)] ->
             Z.Zipper (Tag, Tag) -> M.Map TransformationRule Int
instRules funs = Z.foldlz' applyFuns M.empty
    where applyFuns m z
              | correct == proposed = m
              | otherwise = foldl (applyFun z) m funs
              where (correct, proposed) = Z.cursor z
                    applyFun z m f = case f z of
                                       Nothing -> m
                                       Just r -> M.insertWith' (+) r 1 m

{- Tests
  h <- IO.openFile "brown-pos-train.txt" IO.ReadMode
  c <- IO.hGetContents h
  let proposedRules = instRules [instNextTagRule, instPrevTagRule, instSurroundTagRule] $ initialLearningState $ map toTrainingInstance $ words c
  :t proposedRules
  take 10 $ M.assocs proposedRules
  --Data.Map.size proposedRules
[(NextTagRule (Replacement "(" ",") "JJ",1),
 (NextTagRule (Replacement "(" "NIL") "(",1),
 (NextTagRule (Replacement "(" "NIL") "CD",4),
 (NextTagRule (Replacement "(" "NIL") "NIL",5),
 (NextTagRule (Replacement ")" ",") "TO",1),
 (NextTagRule (Replacement ")" "NIL") ",",1),
 (NextTagRule (Replacement ")" "NIL") ".",4),
 (NextTagRule (Replacement ")" "NIL") "BER",1),
 (NextTagRule (Replacement ")" "NIL") "CC",1),
 (NextTagRule (Replacement ")" "NIL") "IN",1)]
-}

---- Sort the rules (in decreasing order) so that we can select
---- the top n rules (rule most frequently applicable)
sortRules :: M.Map TransformationRule Int -> [(TransformationRule, Int)]
sortRules = L.sortBy (\(_,a) (_,b) -> compare b a) . M.toList


testLearnRules0 train = do
  h <- IO.openFile train IO.ReadMode
  c <- IO.hGetContents h
  let proposedRules = instRules [instNextTagRule, instPrevTagRule, instSurroundTagRule] $ initialLearningState $ map toTrainingInstance $ words c
  --let rs = take 10 $ M.assocs proposedRules
  -- print rs
  -- putStrLn $ show rs
  let rs = take 10 $ sortRules proposedRules
  print rs

{- Tests
 testLearnRules "brown-pos-train.txt"
 [(PrevTagRule (Replacement "TO" "IN") "NN",4918),
  (NextTagRule (Replacement "TO" "IN") "AT",3471),
  (PrevTagRule (Replacement "NN" "VB") "TO",3380),
  (PrevTagRule (Replacement "TO" "IN") "VBN",2176),
  (PrevTagRule (Replacement "TO" "IN") "JJ",1626),
  (PrevTagRule (Replacement "JJ" "NN") "AT",1492),
  (PrevTagRule (Replacement "VBN" "VBD") "NN",1482),
  (PrevTagRule (Replacement "TO" "IN") "NNS",1374),
  (PrevTagRule (Replacement "VBN" "VBD") "PPS",1370),
  (PrevTagRule (Replacement "CS" "WPS") "NN",1340)]
-}


---- Given the training data (in the Zipper incorrectTag, correctTag)
---- See which rule applies and either return the new Just Tag or Nothing
ruleApplication :: TransformationRule -> Z.Zipper (Tag, Tag) -> Maybe Tag
ruleApplication (NextTagRule (Replacement old new) next) z = do
  (_, proposed)     <- Z.safeCursor z
  (_, nextProposed) <- rightCursor z
  if proposed == old && nextProposed == next then Just new else Nothing
ruleApplication (PrevTagRule (Replacement old new) prev) z = do
  (_, proposed)     <- Z.safeCursor z
  (_, prevProposed) <- leftCursor z
  if proposed == old && prevProposed == prev then Just new else Nothing
ruleApplication (SurroundTagRule (Replacement old new) prev next) z = do
  (_, proposed)     <- Z.safeCursor z
  (_, nextProposed) <- rightCursor z
  (_, prevProposed) <- leftCursor z
  if proposed == old && prevProposed == prev &&
      nextProposed == next then Just new else Nothing


---- Calculate how many corrections this rule may generate by
---- taking into account the number of new errors it may induce
scoreRule :: TransformationRule -> Z.Zipper (Tag, Tag) -> Int
scoreRule r z = nCorrect - nIncorrect
    where (nCorrect, nIncorrect) = scoreRule_ r z

scoreRule_ :: TransformationRule -> Z.Zipper (Tag, Tag) -> (Int, Int)
scoreRule_ r = Z.foldlz' (scoreElem r) (0, 0)
    where scoreElem r !s@(!nCorrect, !nIncorrect) z =
              case ruleApplication r z of
                Just tag -> if tag == correct then
                                (nCorrect + 1, nIncorrect)
                            else
                                (nCorrect, nIncorrect + 1)
                Nothing  -> s
              where (correct, _) = Z.cursor z

{- Tests
  h <- IO.openFile "brown-pos-train.txt" IO.ReadMode
  c <- IO.hGetContents h
  let learningState = initialLearningState $ map toTrainingInstance $ words c
  let proposedRules = instRules [instNextTagRule, instPrevTagRule, instSurroundTagRule] learningState
  head $ sortRules proposedRules
  let (firstRule, _) = head $ sortRules proposedRules
  scoreRule firstRule learningState
-}

testLearnRules1 train = do
  h <- IO.openFile train IO.ReadMode
  c <- IO.hGetContents h
  let learningState = initialLearningState $ map toTrainingInstance $ words c
  let proposedRules = instRules [instNextTagRule, instPrevTagRule, instSurroundTagRule] $ learningState
  let sortedRules = sortRules proposedRules
  let (firstRule, firstFreq) = head sortedRules
  print (firstRule, firstFreq)
  let firstScore = scoreRule firstRule learningState
  print (firstRule, firstScore)
--  let (secondRule, secondScore) = (!!) sortedRules 1
--  print (secondRule, secondScore)


{- Result
  (PrevTagRule (Replacement "TO" "IN") "NN",4918)
-}






main = do
   putStrLn "NLP tests"
   evalTaggers ()
   --testLearnRules0 "brown-pos-train.txt"
   testLearnRules1 "brown-pos-train.txt"