-- |
-- Module      :  ELynx.Data.MarkovProcess.MixtureModel
-- Description :  Mixture models are a set of substitution models with weights
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Tue Jan 29 19:17:40 2019.
--
-- To be imported qualified.
module ELynx.Data.MarkovProcess.MixtureModel
  ( -- * Types
    Weight,
    Component (weight, substModel),
    MixtureModel (name, alphabet, components),

    -- * Getters
    getWeights,
    getSubstitutionModels,

    -- * Building mixture models
    fromSubstitutionModels,

    -- * Transformations
    concatenate,
    scale,
    normalize,
    appendNameComponents,
  )
where

import qualified Data.Vector as V
import ELynx.Data.Alphabet.Alphabet hiding (all)
import qualified ELynx.Data.MarkovProcess.SubstitutionModel as S
import Prelude

-- | Mixture model component weight.
type Weight = Double

-- | A mixture model component has a weight and a substitution model.
data Component = Component
  { weight :: Weight,
    substModel :: S.SubstitutionModel
  }
  deriving (Show, Read)

-- | A mixture model with its components.
data MixtureModel = MixtureModel
  { -- | Name
    name :: S.Name,
    alphabet :: Alphabet,
    components :: V.Vector Component
  }
  deriving (Show, Read)

-- | Get weights.
getWeights :: MixtureModel -> V.Vector Weight
getWeights = V.map weight . components

-- | Get substitution models.
getSubstitutionModels :: MixtureModel -> V.Vector S.SubstitutionModel
getSubstitutionModels = V.map substModel . components

-- | Create a mixture model from a list of substitution models.
fromSubstitutionModels :: S.Name -> V.Vector Weight -> V.Vector S.SubstitutionModel -> MixtureModel
fromSubstitutionModels n ws sms
  | null ws = error "fromSubstitutionModels: No weights given."
  | length ws /= length sms = error "fromSubstitutionModels: Number of weights and substitution models does not match."
  | not $ allEqual alphs = error "fromSubstitutionModels: alphabets of substitution models are not equal."
  | otherwise = MixtureModel n (V.head alphs) comps
  where
    comps = V.zipWith Component ws sms
    alphs = V.map S.alphabet sms
    allEqual xs | V.null xs = True
                | otherwise = V.all (== V.head xs) xs

-- | Concatenate mixture models.
concatenate :: S.Name -> V.Vector MixtureModel -> MixtureModel
concatenate n mms = fromSubstitutionModels n ws sms
  where
    comps = V.concatMap components mms
    ws = V.map weight comps
    sms = V.map substModel comps

scaleComponent :: Double -> Component -> Component
scaleComponent s c = c {substModel = s'} where s' = S.scale s $ substModel c

-- | Scale all substitution models of the mixture model.
scale :: Double -> MixtureModel -> MixtureModel
scale s m = m {components = cs'}
  where
    cs = components m
    cs' = V.map (scaleComponent s) cs

-- | Globally normalize a mixture model so that on average one event happens per
-- unit time.
normalize :: MixtureModel -> MixtureModel
normalize mm = scale (1 / c) mm
  where
    c = sum $ V.zipWith (*) weights scales
    weights = getWeights mm
    scales = V.map S.totalRate $ getSubstitutionModels mm

appendNameComponent :: S.Name -> Component -> Component
appendNameComponent n c = c {substModel = s'}
  where
    s' = S.appendName n $ substModel c

-- | Append byte string to all substitution models of mixture model.
appendNameComponents :: S.Name -> MixtureModel -> MixtureModel
appendNameComponents n m = m {components = cs'}
  where
    cs = components m
    cs' = V.map (appendNameComponent n) cs