module HiddenMarkovModel.Hardwired where

import qualified HiddenMarkovModel as HMMF
import HiddenMarkovModel
         (NamedGaussian, Gaussian, ShapeInt, ShapeState,
          State(State), state, 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.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Hermitian as Hermitian
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector

import qualified Data.Array.Comfort.Boxed as Array
import Data.Array.Comfort.Boxed (Array)

import qualified Data.NonEmpty as NonEmpty
import Data.Map (Map)
import Data.Semigroup ((<>))



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

numberOfStates :: Int
numberOfStates = 6

statesShape :: ShapeState
statesShape = HMMF.statesShape numberOfStates

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

labelFromStateMap :: Array ShapeState String
labelFromStateMap =
   Array.fromList statesShape $
      Label.pause :
      Label.clickBegin :
      Label.clickEnd :
      Label.chirpingMain :
      Label.chirpingPause :
      Label.growling :
      []

stateFromLabelMap :: Map String State
stateFromLabelMap =
   inverseMap labelFromStateMap



type Pattern = Pat.T ShapeState Double

infixr 7 *<>

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


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

pattern :: Pattern
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 :: Gaussian
hmm = hmmTrained

hmmTrained :: Gaussian
hmmTrained =
   HMM.Cons {
      HMM.initial =
         Vector.fromList statesShape [0.0,0.0,0.0,1.0,0.0,0.0],
      HMM.transition =
         Square.fromGeneral $ Matrix.fromRowArray statesShape $
         Array.fromList statesShape $ fmap (Vector.fromList statesShape) $
            [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 $ Array.fromList statesShape $
            (Vector.autoFromList [0.9513191890047871], covariance [[0.17689006357223516]]) :
            (Vector.autoFromList [1.5879408507110250], covariance [[0.600575479836784]]) :
            (Vector.autoFromList [0.7454942099113683], covariance [[0.4088353694711163]]) :
            (Vector.autoFromList [1.0231037870319346], covariance [[0.19801719658707737]]) :
            (Vector.autoFromList [0.6214106323233616], covariance [[0.3085570412459857]]) :
            (Vector.autoFromList [1.5574159338071116], covariance [[0.6221472768351596]]) :
            []}

hmmPattern :: Gaussian
hmmPattern =
   HMM.finishTraining $ flip Pat.finish pattern $
      Distr.gaussianTrained $ Array.fromList statesShape $
         map (\(center,cov) -> (1,center,cov)) $
            (Vector.autoFromList [1.00], covariance [[0.17]]) :
            (Vector.autoFromList [1.60], covariance [[0.60]]) :
            (Vector.autoFromList [0.75], covariance [[0.40]]) :
            (Vector.autoFromList [1.00], covariance [[0.20]]) :
            (Vector.autoFromList [0.60], covariance [[0.30]]) :
            (Vector.autoFromList [1.60], covariance [[0.60]]) :
            []

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


type HermitianMatrix = Hermitian.Hermitian ShapeInt

covariance :: [[Double]] -> HermitianMatrix Double
covariance =
   maybe
      (Hermitian.autoFromList MatrixShape.RowMajor [])
      (Hermitian.gramian . Matrix.fromRowsNonEmpty) .
   NonEmpty.fetch . map Vector.autoFromList


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