module Math.HiddenMarkovModel.Utility where

import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Hermitian as Hermitian
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Triangular (Diagonal)
import Numeric.LAPACK.Matrix.Array (ArrayMatrix)
import Numeric.LAPACK.Vector (Vector, (.*|))

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable as StorableArray
import qualified Data.Array.Comfort.Boxed as Array
import qualified Data.Array.Comfort.Shape as Shape

import Foreign.Storable (Storable)

import qualified System.Random as Rnd

import qualified Control.Monad.Trans.State as MS


normalizeProb :: (Shape.C sh, Class.Real a) => Vector sh a -> Vector sh a
normalizeProb = snd . normalizeFactor

normalizeFactor :: (Shape.C sh, Class.Real a) => Vector sh a -> (a, Vector sh a)
normalizeFactor xs =
   let c = Vector.sum xs
   in  (c, recip c .*| xs)

-- see htam:Stochastic
randomItemProp ::
   (Rnd.RandomGen g, Rnd.Random b, Num b, Ord b) =>
   [(a,b)] -> MS.State g a
randomItemProp props =
   let (keys,ps) = unzip props
   in  do p <- MS.state (Rnd.randomR (0, sum ps))
          return $
             fst $ head $ dropWhile ((0<=) . snd) $
             zip keys $ tail $ scanl (-) p ps

attachOnes :: (Num b) => [a] -> [(a,b)]
attachOnes = map (flip (,) 1)


vectorDim :: Shape.C sh => Vector sh a -> Int
vectorDim = Shape.size . StorableArray.shape


hermitianFromList ::
   (Shape.C sh, Storable a) => sh -> [a] -> Hermitian.Hermitian sh a
hermitianFromList = Hermitian.fromList MatrixShape.RowMajor


squareConstant ::
   (Shape.C sh, Class.Real a) => sh -> a -> Matrix.Square sh a
squareConstant =
   (ArrMatrix.fromVector .) .
   Vector.constant . MatrixShape.square MatrixShape.RowMajor

squareFromLists ::
   (Shape.C sh, Eq sh, Storable a) => sh -> [Vector sh a] -> Matrix.Square sh a
squareFromLists sh =
   Square.fromGeneral . Matrix.fromRowArray sh . Array.fromList sh

diagonal :: (Shape.C sh, Class.Real a) => Vector sh a -> Diagonal sh a
diagonal = Triangular.diagonal MatrixShape.RowMajor


newtype Distance f a = Distance {getDistance :: f a -> f a -> a}

distance ::
   (Shape.C sh, Eq sh, Class.Real a) =>
   Vector sh a -> Vector sh a -> a
distance =
   getDistance $
   Class.switchReal
      (Distance $ (Vector.normInf .) . Vector.sub)
      (Distance $ (Vector.normInf .) . Vector.sub)

matrixDistance ::
   (Shape.C sh, Eq sh, Class.Real a) =>
   ArrayMatrix sh a -> ArrayMatrix sh a -> a
matrixDistance a b = distance (ArrMatrix.toVector a) (ArrMatrix.toVector b)