module HiddenMarkovModel.Hardwired where

import HiddenMarkovModel (inverseMap)
import qualified Label

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel.Pattern as Pat
import qualified Math.HiddenMarkovModel.Named as HMMNamed
import qualified Math.HiddenMarkovModel as HMM

import qualified Numeric.Container as NC
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Data.Packed.Matrix (Matrix)

import qualified Data.Map as Map; import Data.Map (Map)
import Data.Semigroup ((<>))
import Data.Tuple.HT (mapFst)



pause, clickBegin, clickEnd, chirping, chirpingPause, growling :: HMM.State
pause         = HMM.state 0
clickBegin    = HMM.state 1
clickEnd      = HMM.state 2
chirping      = HMM.state 3
chirpingPause = HMM.state 4
growling      = HMM.state 5

numberOfStates :: Int
numberOfStates = 6

formatState :: Distr.State -> String
formatState (Distr.State s) =
   case s of
      1 -> "click begin"
      2 -> "click end"
      3 -> "chirping loop"
      4 -> "chirping pause"
      5 -> "growling"
      _ -> "pause"

labelFromStateMap :: Map HMM.State String
labelFromStateMap =
   Map.fromList $ map (mapFst HMM.state) $
      (0, Label.pause) :
      (1, Label.clickBegin) :
      (2, Label.clickEnd) :
      (3, Label.chirpingMain) :
      (4, Label.chirpingPause) :
      (5, Label.growling) :
      []

stateFromLabelMap :: Map String HMM.State
stateFromLabelMap =
   inverseMap labelFromStateMap



infixr 7 *<>

(*<>) :: Int -> Pat.T Double -> Pat.T Double
(*<>) = Pat.replicate


rasping :: Pat.T Double
rasping =
   15 *<>
      (600 *<> Pat.atom clickBegin
       <>
       600 *<> Pat.atom clickEnd)

pattern :: Pat.T Double
pattern =
   10000 *<> Pat.atom pause
   <>
   15 *<>
      (rasping
       <>
       6000 *<> Pat.atom chirping
       <>
       1500 *<> Pat.atom chirpingPause)
   <>
   rasping
   <>
   60000 *<> Pat.atom pause
   <>
   7 *<>
      (150 *<> Pat.atom growling
       <>
       1000 *<> Pat.atom pause)



hmm :: HMM.Gaussian Double
hmm = hmmTrained

hmmTrained :: HMM.Gaussian Double
hmmTrained =
   HMM.Cons {
      HMM.initial =
         Vector.fromList [0.0,0.0,0.0,1.0,0.0,0.0],
      HMM.transition =
         Matrix.fromLists $
            [0.9994586913864266,0.0,2.100090303883067e-5,0.0,0.0,1.0218978102189781e-2] :
            [0.0,0.9855812349085892,4.09517609257198e-3,0.0,2.4915465385299874e-3,0.0] :
            [0.0,1.4418765091410832e-2,0.9956108112648844,0.0,0.0,0.0] :
            [0.0,0.0,2.730117395047987e-4,0.9994628194305887,0.0,0.0] :
            [0.0,0.0,0.0,5.371805694114036e-4,0.99750845346147,0.0] :
            [5.413086135733135e-4,0.0,0.0,0.0,0.0,0.9897810218978101] :
            [],
      HMM.distribution =
         Distr.gaussian $
            (Vector.fromList [0.9513191890047871], covariance [[0.17689006357223516]]) :
            (Vector.fromList [1.5879408507110250], covariance [[0.600575479836784]]) :
            (Vector.fromList [0.7454942099113683], covariance [[0.4088353694711163]]) :
            (Vector.fromList [1.0231037870319346], covariance [[0.19801719658707737]]) :
            (Vector.fromList [0.6214106323233616], covariance [[0.3085570412459857]]) :
            (Vector.fromList [1.5574159338071116], covariance [[0.6221472768351596]]) :
            []}

hmmPattern :: HMM.Gaussian Double
hmmPattern =
   (HMM.finishTraining $
    Pat.finish numberOfStates (Distr.GaussianTrained Map.empty) pattern)
       {HMM.distribution =
          Distr.gaussian $
            (Vector.fromList [1.00], covariance [[0.17]]) :
            (Vector.fromList [1.60], covariance [[0.60]]) :
            (Vector.fromList [0.75], covariance [[0.40]]) :
            (Vector.fromList [1.00], covariance [[0.20]]) :
            (Vector.fromList [0.60], covariance [[0.30]]) :
            (Vector.fromList [1.60], covariance [[0.60]]) :
            []}

hmmNamed :: HMMNamed.Gaussian Double
hmmNamed =
   HMMNamed.Cons {
      HMMNamed.model = hmm,
      HMMNamed.nameFromStateMap = labelFromStateMap,
      HMMNamed.stateFromNameMap = stateFromLabelMap
   }


covariance :: [[Double]] -> Matrix Double
covariance xs =
   let m = Matrix.fromLists xs
   in  Matrix.trans m NC.<> m


scaleStdDev :: Double -> HMM.Gaussian Double -> HMM.Gaussian Double
scaleStdDev k model =
   model {
      HMM.distribution =
         let Distr.Gaussian arr = HMM.distribution model
         in  Distr.Gaussian $
             fmap (\(center,dev,c) -> (center, NC.scale k dev, c/k)) arr
   }