{- |
Example of an HMM with continuous emissions with two-dimensional observations.
We train a model to accept a parametric curve of a circle with a certain speed.
This is like "Math.HiddenMarkovModel.Example.SineWave" but in two dimensions.

The four hidden states correspond to the four quadrants.
-}
module Math.HiddenMarkovModel.Example.Circle
{-# 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 Data.Packed.Vector (Vector)

import qualified System.Random as Rnd

import qualified Control.Monad.Trans.State as MS
import Control.Monad (liftM2, replicateM)

import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import Data.Function.HT (nest)
import Data.NonEmpty ((!:))



hmm :: HMM.Gaussian Double
hmm =
   HMM.Cons {
      HMM.initial = Vector.fromList [1/4, 1/4, 1/4, 1/4],
      HMM.transition =
         Matrix.fromLists $
            [0.9, 0.0, 0.0, 0.1] :
            [0.1, 0.9, 0.0, 0.0] :
            [0.0, 0.1, 0.9, 0.0] :
            [0.0, 0.0, 0.1, 0.9] :
            [],
      HMM.distribution =
         let cov0 = Matrix.fromLists [[0.10, -0.09],[-0.09, 0.10]]
             cov1 = Matrix.fromLists [[0.10,  0.09],[ 0.09, 0.10]]
         in  Distr.gaussian $
                (Vector.fromList [ 0.5,  0.5], cov0) :
                (Vector.fromList [-0.5,  0.5], cov1) :
                (Vector.fromList [-0.5, -0.5], cov0) :
                (Vector.fromList [ 0.5, -0.5], cov1) :
                []
   }

circleLabeled :: NonEmpty.T [] (HMM.State, Vector Double)
circleLabeled =
   NonEmpty.mapTail (take 200) $
   fmap
      (\x ->
         (HMM.state $ mod (floor (x*2/pi)) 4,
          Vector.fromList [cos x, sin x])) $
   NonEmptyC.iterate (0.1+) 0

circle :: NonEmpty.T [] (Vector Double)
circle = fmap snd circleLabeled

revealed :: NonEmpty.T [] HMM.State
revealed = HMM.reveal hmm circle

{- |
Sample multivariate normal distribution and reconstruct it from the samples.
You should obtain the same parameters.
-}
reconstructDistribution :: HMM.Gaussian Double
reconstructDistribution =
   let s0 = HMM.state 0
       gen = Distr.generate (HMM.distribution hmm) s0
   in  HMM.finishTraining $ HMM.trainSupervised 1 $ fmap ((,) s0) $
       flip MS.evalState (Rnd.mkStdGen 23) $
       liftM2 (!:) gen $ replicateM 1000 gen


hmmTrainedSupervised :: HMM.Gaussian Double
hmmTrainedSupervised =
   HMM.finishTraining $ HMM.trainSupervised 4 circleLabeled

hmmTrainedUnsupervised :: HMM.Gaussian Double
hmmTrainedUnsupervised =
   HMM.finishTraining $ HMM.trainUnsupervised hmm circle

hmmIterativelyTrained :: HMM.Gaussian Double
hmmIterativelyTrained =
   nest 100
      (HMM.finishTraining . flip HMM.trainUnsupervised circle)
      hmm