module Model where

import           BAliPhy.Run
import           MCMC (runMCMC)
import           Options.Applicative
import           Probability
import           Bio.Alignment
import           Bio.Alphabet
import           Bio.Sequence
import           Tree
import           Tree.Newick
import           SModel
import           IModel

branch_length_dist topology branch = gamma (1/2) (2/fromIntegral n) where n = numBranches topology

model seq_data = do
    let taxa            = getTaxa seq_data
        tip_seq_lengths = getSequenceLengths seq_data

    -- Tree
    scale <- prior $ gamma (1/2) 2
    tree  <- prior $ uniformLabelledTree'' taxa branch_length_dist

    -- Indel model
    indel_rate   <- prior $ logLaplace (-4) 0.707
    mean_length <- (1 +) <$> sample (exponential 10)
    let imodel = rs07 indel_rate mean_length tree

    -- Substitution model
    freqs  <- prior $ symmetricDirichletOn (letterSet dna) 1
    kappa1 <- prior $ logNormal 0 1
    kappa2 <- prior $ logNormal 0 1
    let tn93_model = tn93' dna kappa1 kappa2 freqs

    -- Alignment
    alignment <- prior $ phyloAlignment tree imodel scale tip_seq_lengths

    -- Observation
    observe seq_data $ phyloCTMC tree alignment tn93_model scale

    return
        [ "tree" %=% writeNewick tree
        , "log(indel_rate)" %=% log indel_rate
        , "mean_length" %=% mean_length
        , "kappa1" %=% kappa1
        , "kappa2" %=% kappa2
        , "frequencies" %=% freqs
        , "scale" %=% scale
        , "|T|" %=% treeLength tree
        , "scale*|T|" %=% treeLength tree * scale
        , "|A|" %=% alignmentLength alignment
        ]

main = do
    (options, filename) <- execParser $
      modelRunParserWith "Model" 200000 $
        strArgument (metavar "SEQUENCES" <> help "Unaligned DNA sequences")

    runInfo <- initializeModelRun (runMode options)

    seq_data <- mkUnalignedCharacterData dna <$> loadSequences filename

    mcmcState <- makeLoggedMCMCState runInfo (logFormats options) $ model seq_data

    case runInfo of
      TestRun -> printInitialModel (logFormats options) mcmcState
      MCMCRun directory -> do
        reportModelRun (iterations options) (logFormats options) directory
        runMCMC (iterations options) mcmcState