module Model where

import BAliPhy.Run
import Bio.Alignment
import Bio.Alphabet
import qualified Data.Set as Set
import IModel
import MCMC
import Options.Applicative
import Probability
import SModel
import SModel.Parsimony
import Tree
import Tree.Newick

gtr_m7_model codons = do
    let nucs = getNucleotides codons

    -- GTR model parameters
    sym <- sample $ symmetricDirichletOn (Set.fromList $ letter_pair_names nucs) 1
    pi <- sample $ symmetricDirichletOn (letterSet nucs) 1

    let posSelModel w = gtr' sym pi nucs +> x3 codons +> dNdS w

    -- Independent uniform priors on the beta mean and normalized variance.
    mu <- sample $ uniform 0 1
    v <- sample $ uniform 0 1
    let m7Model = posSelModel +> m7 mu v 4

    let loggers =
            [ "gtr:sym" %=% sym
            , "gtr:pi" %=% pi
            , "m7:mu" %=% mu
            , "m7:v" %=% v
            ]

    return (m7Model, loggers)

model sequenceData = do
    let taxa = getTaxa sequenceData

    tree <- sample $ uniformLabelledTree taxa (gamma 0.5 (1 / fromIntegral (length taxa)))
    let tlength = treeLength tree

    sigma <- sample $ logLaplace (-3) 1
    indelRates <- fmap (** sigma) <$> sample (iidMap (getUEdgesSet tree) (logNormal 0 1))
    let indelTree = addBranchRates indelRates tree

    scale <- sample $ gamma 0.5 2
    addMove 2 (scaleGroupsSlice [scale] (branchLengths tree))
    addMove 1 (scaleGroupsMH [scale] (branchLengths tree))

    let codons = mkCodons dna (geneticCode "standard")
    (m7_model, log_m7_model) <- gtr_m7_model codons

    rate <- sample $ logLaplace (-4) 0.707
    meanLength <- sample $ shifted_exponential 10 1
    let imodel = IModel.rs07 rate meanLength tree

    let sequenceLengths = getSequenceLengths sequenceData
    (alignment, propertiesA) <- sampleWithProps (phyloAlignment indelTree imodel scale sequenceLengths)
    properties <- observe sequenceData (phyloCTMC tree alignment m7_model scale)

    let alignment_length = alignmentLength alignment
    let num_indels = totalNumIndels alignment
    let total_length_indels = totalLengthIndels alignment
    let prior_A = ln (probability propertiesA)
    let ancStates = prop_anc_cat_states properties
    let ancAlignment = toFasta $ ancestralAlignment tree alignment (getSMap m7_model) codons ancStates
    let substs = parsimony tree (unitCostMatrix codons) (sequenceData, alignment)

    let loggers =
            [ "indelRates:sigma" %=% sigma
            , "S1" %>% log_m7_model
            , "rs07:rate" %=% rate
            , "rs07:mean_length" %=% meanLength
            , "scale" %=% scale
            , "scale*|T|" %=% (scale * tlength)
            , "|A|" %=% alignment_length
            , "#indels" %=% num_indels
            , "|indels|" %=% total_length_indels
            , "#substs" %=% substs
            , "prior_A" %=% prior_A
            ]

    return loggers

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

    runInfo <- initializeModelRun (runMode options)

    sequenceData <-
        mkUnalignedCharacterData (mkCodons dna standard_code)
            <$> loadSequences filename

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

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