{- |
This is an example of an HMM with discrete emissions.
We model a traffic light consisting of the colors red, yellow, green,
where only one lamp can be switched on at every point in time.
This way, when it is yellow you cannot tell immediately
whether it will switch to green or red.
We can only infer this from the light seen before.

There are four hidden states:
0 emits red, 1 emits yellow between red and green,
2 emits green, 3 emits yellow between green and red.

We quantise time in time steps.
The transition matrix of the model 'hmm' encodes
the expected duration of every state counted in time steps
and what states follow after each other.
E.g. transition probability of 0.8 of a state to itself means
that the expected duration of the state is 5 time steps (1/(1-0.8)).
However, it is a geometric distribution,
that is, shorter durations are always more probable.

The distribution of 'hmm' encodes which lights a state activates.
In our case everything is deterministic:
Every state can switch exactly one light on.

Given a sequence of observed lights
the function 'HMM.reveal' tells us the most likely sequence of states.
We test this with the light sequences in 'stateSequences'
where we already know the hidden states
as they are stored in 'labeledSequences'.
'verifyRevelation' compares the computed state sequence with the given one.

We also try some trainings in 'hmmTrainedSupervised' et.al.
-}
module Math.HiddenMarkovModel.Example.TrafficLight
{-# WARNING "do not import that module, it is only intended for demonstration" #-}
   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)