module Math.HiddenMarkovModel.Example.TrafficLightPrivate where

import qualified Math.HiddenMarkovModel as HMM
import qualified Math.HiddenMarkovModel.Distribution as Distr

import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector

import Text.Read.HT (maybeRead)

import Control.Monad (liftM2)

import qualified Data.Map as Map
import qualified Data.NonEmpty as NonEmpty
import qualified Data.List.HT as ListHT
import Data.NonEmpty ((!:))



data Color = Red | Yellow | Green
   deriving (Eq, Ord, Enum, Show, Read)

{- |
Using 'show' and 'read' is not always a good choice
since they must format and parse Haskell expressions
which is not of much use to the outside world.
-}
instance Distr.CSVSymbol Color where
   cellFromSymbol = show
   symbolFromCell = maybeRead


hmm :: HMM.Discrete Double Color
hmm =
   HMM.Cons {
      HMM.initial = Vector.fromList [1/3, 1/6, 1/3, 1/6],
      HMM.transition =
         Matrix.fromLists $
            [0.8, 0.0, 0.0, 0.2] :
            [0.2, 0.8, 0.0, 0.0] :
            [0.0, 0.2, 0.8, 0.0] :
            [0.0, 0.0, 0.2, 0.8] :
            [],
      HMM.distribution =
         Distr.Discrete $ Map.fromList $
            (Red,    Vector.fromList [1,0,0,0]) :
            (Yellow, Vector.fromList [0,1,0,1]) :
            (Green,  Vector.fromList [0,0,1,0]) :
            []
   }

hmmDisturbed :: HMM.Discrete Double Color
hmmDisturbed =
   HMM.Cons {
      HMM.initial = Vector.fromList [1/4, 1/4, 1/4, 1/4],
      HMM.transition =
         Matrix.fromLists $
            [0.3, 0.2, 0.2, 0.3] :
            [0.3, 0.3, 0.2, 0.2] :
            [0.2, 0.3, 0.3, 0.2] :
            [0.2, 0.2, 0.3, 0.3] :
            [],
      HMM.distribution =
         Distr.Discrete $ Map.fromList $
            (Red,    Vector.fromList [0.6, 0.2, 0.2, 0.2]) :
            (Yellow, Vector.fromList [0.2, 0.6, 0.2, 0.6]) :
            (Green,  Vector.fromList [0.2, 0.2, 0.6, 0.2]) :
            []
   }


red, yellowRG, green, yellowGR :: (HMM.State, Color)
red      = (HMM.state 0, Red)
yellowRG = (HMM.state 1, Yellow)
green    = (HMM.state 2, Green)
yellowGR = (HMM.state 3, Yellow)

labeledSequences :: NonEmpty.T [] (NonEmpty.T [] (HMM.State, Color))
labeledSequences =
   (red !: red : red : red :
    yellowRG : yellowRG :
    green : green : green : green : green :
    yellowGR :
    red : red : red :
    []) !:
   (green !: green : green :
    yellowGR :
    red : red : red : red :
    yellowRG :
    green : green : green : green : green :
    yellowGR : yellowGR :
    []) :
   []

{- |
Construct a Hidden Markov model by watching a set
of manually created sequences of emissions and according states.
-}
hmmTrainedSupervised :: HMM.Discrete Double Color
hmmTrainedSupervised =
   HMM.trainMany (HMM.trainSupervised 4) labeledSequences


stateSequences :: NonEmpty.T [] (NonEmpty.T [] Color)
stateSequences = fmap (fmap snd) labeledSequences

{- |
Construct a Hidden Markov model starting from a known model
and a set of sequences that contain only the emissions, but no states.
-}
hmmTrainedUnsupervised :: HMM.Discrete Double Color
hmmTrainedUnsupervised =
   HMM.trainMany (HMM.trainUnsupervised hmm) stateSequences

{- |
Repeat unsupervised training until convergence.
-}
hmmIterativelyTrained :: HMM.Discrete Double Color
hmmIterativelyTrained =
   snd $ head $ dropWhile fst $
   ListHT.mapAdjacent (\hmm0 hmm1 -> (HMM.deviation hmm0 hmm1 > 1e-5, hmm1)) $
   iterate
      (flip HMM.trainMany stateSequences . HMM.trainUnsupervised)
      hmmDisturbed


verifyRevelation ::
   HMM.Discrete Double Color -> NonEmpty.T [] (HMM.State, Color) -> Bool
verifyRevelation model xs =
   fmap fst xs == HMM.reveal model (fmap snd xs)

verifyRevelations :: [Bool]
verifyRevelations =
   liftM2 verifyRevelation
      [hmm, hmmDisturbed, hmmTrainedSupervised, hmmTrainedUnsupervised]
      (NonEmpty.flatten labeledSequences)