-- |
-- Module      :  ELynx.MarkovProcess.Nucleotide
-- Description :  Substitution models using nucleotides
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Jan 24 08:33:26 2019.
--
-- XXX: Maybe rename to something like /DNA substitution models/. Nucleotide ~
-- Alphabet; DNA ~ Character.
--
-- The order of nucleotides is A, C, G, T; see 'ELynx.Character.Nucleotide'.
--
-- For the different DNA substitution models, please see
-- https://en.wikipedia.org/wiki/Models_of_DNA_evolution
module ELynx.MarkovProcess.Nucleotide
  ( jc,
    f81,
    hky,
    gtr4,
  )
where

import qualified Data.Vector.Storable as V
import ELynx.Alphabet.Alphabet
import ELynx.MarkovProcess.RateMatrix
import ELynx.MarkovProcess.SubstitutionModel
import Numeric.LinearAlgebra hiding (normalize)

-- XXX: Another idea of structuring the code. This would probably be cleaner in
-- the long run.

-- data PhyloModel = MixtureModel | SubstitutionModel

--
-- I think it could simply be:
-- data PhyloModel = [(Weight, SubstitutionModel)]
--

-- data MixtureModel = [(Weight, SubstitutionModel)]

-- data SubstitutionModel = SMDNA DNASubstitutionModel | SMAA AASubstitutionModel

-- data DNASubstitutionModel = JC | HKY Double StationaryDistribution

-- data AASubstitutionModel = LG | ...

n :: Int
-- n = length (alphabet :: [Nucleotide])
-- Hard code this here. Reduces model dependencies, and number of nucleotides
-- will not change.
n :: Int
n = Int
4

-- | JC model matrix.
jcExch :: ExchangeabilityMatrix
jcExch :: ExchangeabilityMatrix
jcExch =
  (Int
n forall a. Storable a => Int -> Int -> [a] -> Matrix a
>< Int
n)
    [ R
0.0,
      R
1.0,
      R
1.0,
      R
1.0,
      R
1.0,
      R
0.0,
      R
1.0,
      R
1.0,
      R
1.0,
      R
1.0,
      R
0.0,
      R
1.0,
      R
1.0,
      R
1.0,
      R
1.0,
      R
0.0
    ]

uniformVec :: Vector Double
uniformVec :: Vector R
uniformVec = forall a. Storable a => Int -> a -> Vector a
V.replicate Int
n (R
1 forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)

-- | JC substitution model.
jc :: Normalize -> SubstitutionModel
jc :: Normalize -> SubstitutionModel
jc Normalize
nz = Alphabet
-> Name
-> [R]
-> Normalize
-> Vector R
-> ExchangeabilityMatrix
-> SubstitutionModel
substitutionModel Alphabet
DNA Name
"JC" [] Normalize
nz Vector R
uniformVec ExchangeabilityMatrix
jcExch

-- | F81 substitution model.
f81 :: Normalize -> StationaryDistribution -> SubstitutionModel
f81 :: Normalize -> Vector R -> SubstitutionModel
f81 Normalize
nz Vector R
d = Alphabet
-> Name
-> [R]
-> Normalize
-> Vector R
-> ExchangeabilityMatrix
-> SubstitutionModel
substitutionModel Alphabet
DNA Name
"F81" [] Normalize
nz Vector R
d ExchangeabilityMatrix
jcExch

hkyExch :: Double -> ExchangeabilityMatrix
hkyExch :: R -> ExchangeabilityMatrix
hkyExch R
k =
  (Int
n forall a. Storable a => Int -> Int -> [a] -> Matrix a
>< Int
n)
    [R
0.0, R
1.0, R
k, R
1.0, R
1.0, R
0.0, R
1.0, R
k, R
k, R
1.0, R
0.0, R
1.0, R
1.0, R
k, R
1.0, R
0.0]

-- | HKY substitution model.
hky :: Normalize -> Double -> StationaryDistribution -> SubstitutionModel
hky :: Normalize -> R -> Vector R -> SubstitutionModel
hky Normalize
nz R
k Vector R
d = Alphabet
-> Name
-> [R]
-> Normalize
-> Vector R
-> ExchangeabilityMatrix
-> SubstitutionModel
substitutionModel Alphabet
DNA Name
"HKY" [R
k] Normalize
nz Vector R
d ExchangeabilityMatrix
e where e :: ExchangeabilityMatrix
e = R -> ExchangeabilityMatrix
hkyExch R
k

-- | HKY substitution model.
gtr4 :: Normalize -> [Double] -> StationaryDistribution -> SubstitutionModel
gtr4 :: Normalize -> [R] -> Vector R -> SubstitutionModel
gtr4 Normalize
nz [R]
es Vector R
d = Alphabet
-> Name
-> [R]
-> Normalize
-> Vector R
-> ExchangeabilityMatrix
-> SubstitutionModel
substitutionModel Alphabet
DNA Name
"GTR" [R]
es Normalize
nz Vector R
d ExchangeabilityMatrix
e
  where
    e :: ExchangeabilityMatrix
e = forall a.
(RealFrac a, Container Vector a) =>
Int -> [a] -> Matrix a
exchFromListUpper Int
n [R]
es