// $Id: dlmodel.h,v 1.72 2018/01/03 21:32:57 mkkuhner Exp $

/*
Copyright 2002-2005 Mary K. Kuhner, Peter Beerli, and Joseph Felsenstein

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

*/

/* Authors: Peter Beerli, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


/*********************************************************************
 DLModel is the key control class for polymorphism based on the
 type of data likelihood model in use.  It is subclassed based on
 model.

 A DataModel is an Abstract Factory (Design Patterns, pg 87)
 responsible for creating the data model output report.

 A DataModel also provides the implementation of a particular data
 likelihood model.  Exactly what needs to be provided is laid out in
 the appropriate DLCalculator class.  For example, each realization of
 a nucleotide based model must provide 4 functions:
 RescaleLengths(), ComputeSiteDLs(), ComputeSubtreeDL(),
 and ComputeCatDL();

 The DataToLikes() function allows input data to be translated to
 likelihoods in a model-appropriate way.

 Each instantiable DataModel should have a distinct value of GetModelType()
 taken from the model_type enum.

 NB:  Lamarc includes two similar, but distinct, ways of looking at
 data analysis.  Classes such as DataType and its dependencies
 (DLCalculators, DLCells, etc.) classify the input data based on
 how it was generated (is it full DNA, SNPs, electrophoretic alleles,
 etc?)  DataModel, in contrast, classifies the data based on what
 model will be used to analyze it.  These hierarchies are distinct,
 because the same kind of data can be analyzed using different
 models (e.g. F84 versus GTR for DNA data) and the same model can
 occasionally be used for different kinds of data (e.g. F84 works
 for both DNA and SNPs).  However, they are far from independent,
 since any given model will only cover a few data types.

 Written by Jim Sloan, rewritten by Jon Yamato.
 winter 2001--removed Dialog factory capability--Mary Kuhner
 2002/01/02 moved DataToLikes into this class--Mary Kuhner
 2002/07/08 added K-Allele model--Mary Kuhner
 2004/09/09 Added data simulation capability

**************************************************************/

#ifndef DLMODEL_H
#define DLMODEL_H

#include <cmath>
#include <deque>
#include <string>
#include <vector>

#include "vectorx.h"
#include "constants.h"
#include "datatype.h" //for the data_type enum
#include "types.h"   // for Cell_ptr type

class Locus;
class Cell;
class DataType;
class DLCalculator;
class Tree;

typedef std::map<std::string, long> alleleMap;
typedef std::vector<alleleMap> alleleVec;
typedef alleleMap::iterator alleleIt;
typedef alleleMap::const_iterator alleleConstIt;

//------------------------------------------------------------------------------------
// The Data model base class

class DataModel
{
  private:
    DataModel();        // undefined, as we want to create only
    // models with good member data values

  protected:
    long            m_ncategories;
    DoubleVec1d     m_catrates;
    DoubleVec1d     m_catprobs;
    double          m_acratio;
    double          m_ucratio;
    bool            m_normalize;
    long            m_nbins;
    long            m_nmarkers;
    DoubleVec1d     m_catcells;  // working storage for category computations
    double          m_relmurate; // relative mutation rate of genetic marker

    // to adjust cat probabilities that don't sum to 1.0
    void ScaleCatProbabilities();
    void ScaleCatRates();

    // to try to catch & fix impossible subtrees, used by ComputeSubtreeDL()
    // this function will throw either a datalikenorm_error or a zero_dl_error!
    void TryToNormalizeAndThrow(long posn, model_type mtype);

    // Validation
    virtual bool IsValidDataModel()      const;

  public:
    DataModel(long nmarkers,
              long numCategories,
              DoubleVec1d categoryRates,
              DoubleVec1d categoryProbabilities,
              double userAutoCorrelationValue,
              bool doNormalize,
              long nbins,
              double relmurate);
    virtual ~DataModel()                         {};
    virtual DataModel* Clone() const = 0;

    virtual model_type GetModelType() const = 0;
    // Data input
    string GetDataModelName()        const;
    string GetDataModelShortName()   const;
    virtual vector<double> DataToLikes(const string& datum, long marker = 0) const = 0;

    // Access
    long        GetNcategories()          const   {return m_ncategories;};
    DoubleVec1d GetCatRates()             const   {return m_catrates;};
    DoubleVec1d GetCatProbabilities()     const   {return m_catprobs;};
    double      GetUserAcratio()          const   {return 1.0/m_acratio;};
    virtual bool ShouldNormalize()        const   {return m_normalize;};
    long        GetNbins()                const   {return m_nbins;};
    virtual double GetRelMuRate()         const   {return m_relmurate;}

    virtual void SetNormalize(bool norm)                 {m_normalize = norm;};

    // SNP support function, used in setting up the SNP invarmodel
    void SetNmarkers(long nmark)                         {m_nmarkers = nmark;};

    // Likelihood Calculation
    void ResetCatCells();   // this needs to be called before each tree evaluation
    virtual double ComputeCatDL(long startmarker, long endmarker) = 0;
    //opt datamodel, default do nothing
    virtual bool OptimizeDataModel(Tree*, const Locus&) {return false;};
    virtual void WriteAlpha(std::ofstream&, long, long, long) {};
    virtual void ResetAlpha() {};
    virtual void SetAlpha(double, long, long) {};

    // Factories

    virtual StringVec1d      CreateDataModelReport() const;
    virtual StringVec1d      ToXML(std::size_t nspaces) const;

    // Simulation
    // Select a vector of rates at random, weighted by their probabilities
    //  and correlated according to m_acratio.
    DoubleVec1d ChooseRandomRates(long nsites) const;

    // Select an ancestral state at random
    virtual DoubleVec1d ChooseAncestralState(long marker) = 0;

    // Simulate data on a branch
    virtual DoubleVec1d SimulateMarker(double branchlength, long whichmarker,
                                       const DoubleVec1d & state) const = 0;

    //Convert bin number to a data string
    virtual string CellToData(Cell_ptr cell, long marker) const = 0;

};

//------------------------------------------------------------------------------------
// Nucleotide model

class NucModel : public DataModel
{
  private:
    NucModel();         // undefined, as we want to create only
    // models with good member data values
  protected:
    DoubleVec1d m_basefreqs;
    bool m_freqsfromdata;
    double  m_perBaseErrorRate;

    // Validation
    void NormalizeBaseFrequencies();
    virtual bool IsValidDataModel()               const;

  public:
    NucModel(long nmarkers,
             long numCategories,
             DoubleVec1d categoryRates,
             DoubleVec1d categoryProbabilities,
             double userAutoCorrelationValue,
             bool doNormalize,
             double relMuRate,
             double freqA,
             double freqC,
             double freqG,
             double freqT,
             bool freqsFromData,
             double perBaseErrorRate);
    virtual ~NucModel() {};

    // Data input
    static vector<double> StaticDataToLikes(const string& datum, double perBaseErrorRate);
    virtual vector<double> DataToLikes(const string& datum, long marker = 0) const;

    // Access
    DoubleVec1d GetBaseFrequencies()  const {return m_basefreqs;};
    bool        FreqsFromData()       const {return m_freqsfromdata;};
    double      GetPerBaseErrorRate() const {return m_perBaseErrorRate;};

    // Data Likelihood Calculation
    virtual void    RescaleLengths(double length1, double length2)      = 0;
    virtual double** ComputeSiteDLs(double** siteDL1, double** siteDL2)      = 0;
    virtual double  ComputeSubtreeDL(Cell& rootdls, double** startmarker,
                                     double** endmarker, long posn) = 0;

    // Factories
    virtual StringVec1d      CreateDataModelReport() const;
    virtual StringVec1d      ToXML(std::size_t nspaces) const;
    virtual DoubleVec1d      ChooseAncestralState(long marker);
    virtual string CellToData(Cell_ptr cell, long marker) const;

};

//------------------------------------------------------------------------------------
// Felsenstein '84 model

class F84Model : public NucModel
{
  private:
    F84Model();     // undefined
    // User-defined parameters
    double m_ttratio;

    // Buffers for data likelihood calculation
    bool        computed1,computed2;
    double      freqar, freqcy, freqgr, freqty;
    DoubleVec1d xcatrates, ycatrates;
    DoubleVec1d expA1, expA2, expB1, expB2, expC1, expC2;
    // we own the following:
    double **basefreqarray;          // used in memcpy(), ComputeSiteDL()
    // this dimensionality is needed by memcpy!
    double **daughter1, **daughter2; // used in a memcpy(), ComputeCatDL()
    double **target;                 // used as return storage, ComputeSiteDL()
    //    which is then used in a memcpy(),
    //    NucCell::SetSiteDLs()
    DoubleVec2d catlikes;           // dim: m_nmarkers * m_ncategories

    void EmptyBuffers();
    void AllocateBuffers();
    void CopyMembers(const F84Model& src);
    void CopyBuffers(const F84Model& src);

    // helper functions for RescaleLengths(double, double);
    void     RescaleLength1(double length1);
    void     RescaleLength2(double length2);

  protected:
    // Validation
    virtual bool IsValidDataModel()               const;

  public:
    F84Model(long nmarkers,
             long numCategories,
             DoubleVec1d categoryRates,
             DoubleVec1d categoryProbabilities,
             double userAutoCorrelationValue,
             bool doNormalize,
             double relMuRate,
             double freqA,
             double freqC,
             double freqG,
             double freqT,
             double ttRatio,
             bool calculateFrequenciesFromData,
             double perBaseErrorRate);
    virtual ~F84Model();
    virtual DataModel* Clone() const;

    F84Model& operator=(const F84Model& src);
    F84Model(const F84Model& src);

    virtual model_type GetModelType() const         { return F84; };

    // Access
    void SetTTratio(double tr);

    double      GetTTratio()          const {return m_ttratio;};

    // Data Likelihood Calculation
    virtual void Finalize();
    virtual void     RescaleLengths(double length1, double length2);
    virtual double** ComputeSiteDLs(double** siteDL1, double** siteDL2);
    virtual double   ComputeSubtreeDL(Cell& rootdls, double** startmarker,
                                      double** endmarker, long posn);
    virtual double   ComputeCatDL(long startmarker, long endmarker);

    // Factories
    virtual StringVec1d  CreateDataModelReport() const;
    virtual StringVec1d  ToXML(std::size_t nspaces) const;

    virtual DoubleVec1d  SimulateMarker(double branchlength,
                                        long whichmarker, const DoubleVec1d & state) const;
};

//------------------------------------------------------------------------------------
// GTR model

class GTRModel : public NucModel
{
  private:
    GTRModel();     // undefined
    // user supplied info
    double AC, AG, AT, CG, CT, TG;

    // stuff derived from user data and set up in Finalize
    DoubleVec1d eigvals;             // eigen values
    DoubleVec2d eigvecs1, eigvecs2;  // eigvecs1 is the simple matrix of
    // eigen vectors.
    // eigvecs2 is the inverted and transposed
    // eigvecs1.

    // used to communicate between RescaleLengths() and ComputeSiteDLs()
    std::deque<bool> computed;
    DoubleVec4d pchange;             // transition probs scaled by length
    // dim: 2 X m_ncategories X rate matrix

    // used to communicate between ComputeSubtreeDL() and ComputeCatDL()
    DoubleVec2d catlikes;            // dim: m_nmarkers X m_ncategories

    // we own the following, used by ComputeSiteDLs():
    double **basefreqarray;         // used in memcpy(), ComputeSiteDL()
    // this dimensionality is needed by memcpy!
    double ***daughter;             // used in a memcpy(), ComputeSiteDL()
    double **target;                // used as return storage, ComputeSiteDL()
    //    which is then used in a memcpy(),
    //    NucCell::SetSiteDLs()

    // helpers for ctor's and operator=
    void AllocateBuffers();
    void EmptyBuffers();
    void CopyMembers(const GTRModel& src);

    // pre-transposed dotproduct for GTR use speed-wise
    void GTRDotProduct(const DoubleVec2d& first, const DoubleVec2d& second,
                       DoubleVec2d& answer);

    // private validity checkers
    void BarfOnBadGTRRates(const DoubleVec1d& rts) const;

    // scratch pad matrix to avoid allocating lots of DoubleVec2ds
    DoubleVec2d scratch;

  protected:
    // Validation
    virtual bool IsValidDataModel()               const;

  public:
    GTRModel(long nmarkers,
             long numCategories,
             DoubleVec1d categoryRates,
             DoubleVec1d categoryProbabilities,
             double userAutoCorrelationValue,
             bool doNormalize,
             double relMuRate,
             double freqA,
             double freqC,
             double freqG,
             double freqT,
             double freqAC,
             double freqAG,
             double freqAT,
             double freqCG,
             double freqCT,
             double freqTG,
             double perBaseErrorRate);
    virtual ~GTRModel();
    virtual DataModel* Clone() const;
    GTRModel(const GTRModel& src);
    GTRModel& operator=(const GTRModel& src);

    virtual model_type GetModelType() const         { return GTR; };

    // Access
    void SetRates(const DoubleVec1d& rts);
    DoubleVec1d GetRates() const;

    // Data Likelihood Calculation
    virtual void Finalize();
    virtual void     RescaleLengths(double length1, double length2);
    virtual double** ComputeSiteDLs(double** siteDL1, double** siteDL2);
    virtual double   ComputeSubtreeDL(Cell& rootdls, double** startmarker,
                                      double** endmarker, long posn);
    virtual double   ComputeCatDL(long startmarker, long endmarker);

    // Factories
    virtual StringVec1d  CreateDataModelReport() const;
    virtual StringVec1d  ToXML(std::size_t nspaces) const;

    virtual DoubleVec1d  SimulateMarker(double branchlength,
                                        long whichmarker, const DoubleVec1d & state) const;
};

//------------------------------------------------------------------------------------
// Allele model

class AlleleModel : public DataModel
{
  private:
    AlleleModel();  // undefined
  protected:
    virtual ~AlleleModel() {};
    AlleleModel(
        long nmarkers,
        long numCategories,
        DoubleVec1d categoryRates,
        DoubleVec1d categoryProbabilities,
        double userAutoCorrelationValue,
        bool doNormalize,
        long nbins,
        double relMuRate);

    DoubleVec2d m_likes;   // likelihoods:  marker x category
    LongVec1d   m_bincounts;  // per marker, total number of usable bins

  public:

    virtual DoubleVec1d RescaleLength(double length);
    virtual double ComputeCatDL(long startmarker, long endmarker);
    virtual void ComputeSiteDLs(Cell_ptr child1, Cell_ptr child2,
                                Cell_ptr thiscell, const DoubleVec1d& vv1,
                                const DoubleVec1d& vv2, long marker) = 0;
    virtual double ComputeSubtreeDLs(Cell& rootdls,  double** startmarker,
                                     double** endmarker, long posn) = 0;
    virtual DoubleVec1d ChooseAncestralState(long marker);
    virtual string CellToData(Cell_ptr cell, long marker) const = 0;
};

//------------------------------------------------------------------------------------
// Step wise model

class StepwiseModel : public AlleleModel
{
  private:

    StepwiseModel();    // undefined

    LongVec1d m_offsets;    // per marker, difference of least allele from 0
    LongVec1d m_threshhold; //per marker, threshhold for allelic differences
    vector<DoubleVec2d> m_steps; // pre-computed constants for mutations of
    // various sizes: nbins x nbins.
    // Per marker, because it uses m_threshhold.
    long m_allowance;      // extra bins to each side of actual markers

    // helper functions for likelihood
    // 2003/11/07 erynes: Changing member variable "threshhold" from hardcoded
    //                    long to marker-specific LongVec1d required the passing
    //                    of the current marker to these functions.
    double Probability(double t, long diff, long marker) const;
    void CalculateSteps(long marker);
    void Initialize(const StringVec2d& uniqueAlleles);

  public:
    // accepting compiler constructed copy-ctor and operator=

    StepwiseModel(
        long nmarkers,
        const StringVec2d& uniqueAlleles,
        long numCategories,
        DoubleVec1d categoryRates,
        DoubleVec1d categoryProbabilities,
        double userAutoCorrelationValue,
        bool doNormalize,
        double relMuRate
        );
    virtual ~StepwiseModel() {};
    virtual DataModel* Clone() const;

    virtual model_type GetModelType() const         { return Stepwise; };

    virtual vector<double> DataToLikes(const string& datum, long marker = 0) const;

    // Data Likelihood Calculation
    virtual void ComputeSiteDLs(Cell_ptr child1, Cell_ptr child2, Cell_ptr thiscell,
                                const DoubleVec1d& vv1, const DoubleVec1d& vv2, long marker);
    virtual double ComputeSubtreeDLs(Cell& rootdls,
                                     double** startmarker, double** endmarker, long posn);

    // Factories
    virtual StringVec1d CreateDataModelReport() const;

    virtual DoubleVec1d SimulateMarker(double branchlength,
                                       long whichmarker, const DoubleVec1d & state) const;
    virtual string CellToData(Cell_ptr cell, long marker) const;
};

//------------------------------------------------------------------------------------
// Brownian model

class BrownianModel : public AlleleModel
{
  private:
    BrownianModel();    // private

  public:
    BrownianModel(
        long nmarkers,
        long numCategories,
        DoubleVec1d categoryRates,
        DoubleVec1d categoryProbabilities,
        double userAutoCorrelationValue,
        bool doNormalize,
        double relMuRate
        );
    virtual ~BrownianModel()   {};
    virtual DataModel* Clone() const;

    virtual bool ShouldNormalize() const {return false;};
    virtual void SetNormalize(bool norm);

    virtual model_type GetModelType() const         { return Brownian; };
    virtual vector<double> DataToLikes(const string& datum, long marker = 0) const;

    // Data Likelihood Calculation
    virtual void ComputeSiteDLs(Cell_ptr child1, Cell_ptr child2, Cell_ptr thiscell,
                                const DoubleVec1d& vv1, const DoubleVec1d& vv2, long marker);
    virtual double ComputeSubtreeDLs(Cell& rootdls,
                                     double** startmarker, double** endmarker, long posn);

    // Factories
    virtual StringVec1d CreateDataModelReport() const;

    virtual DoubleVec1d SimulateMarker(double branchlength,
                                       long whichmarker, const DoubleVec1d & state) const;
    virtual string CellToData(Cell_ptr cell, long marker) const;
};

//------------------------------------------------------------------------------------
// K-Allele model

class KAlleleModel : public AlleleModel
{
  private:
    KAlleleModel();   // undefined
    std::vector< std::map<string, long> > m_allelemaps; //one map per marker

    // helper functions for likelihood
    double probMathFunc( double ut, double coef);
    void Initialize(const StringVec2d& uniqueAlleles);

  public:
    KAlleleModel(
        long nmarkers,
        const StringVec2d& unknownHaplotypeAlleles,
        long numCategories,
        DoubleVec1d categoryRates,
        DoubleVec1d categoryProbabilities,
        double userAutoCorrelationValue,
        bool doNormalize,
        double relMuRate
        );

    virtual ~KAlleleModel() {};
    virtual DataModel* Clone() const;

    virtual bool       ShouldNormalize()    const { return m_normalize; };
    virtual model_type GetModelType()       const { return KAllele; };

    virtual vector<double> DataToLikes(const string& datum, long marker = 0) const;

    // Data Likelihood Calculation
    virtual void ComputeSiteDLs(Cell_ptr child1, Cell_ptr child2, Cell_ptr thiscell,
                                const DoubleVec1d& vv1, const DoubleVec1d& vv2, long marker);
    virtual double ComputeSubtreeDLs(Cell& rootdls,
                                     double** startmarker, double** endmarker, long posn);

    // Factories
    virtual StringVec1d CreateDataModelReport() const;
    virtual StringVec1d ToXML(std::size_t nspaces) const;

    virtual DoubleVec1d SimulateMarker(double branchlength,
                                       long whichmarker,
                                       const DoubleVec1d& state) const;
    virtual string CellToData(Cell_ptr cell, long marker) const;
};

//------------------------------------------------------------------------------------
// MixedKS model

class MixedKSModel : public AlleleModel
{
  private:
    MixedKSModel();   // undefined

    LongVec1d m_offsets;    // per marker, difference of least allele from 0
    LongVec1d m_threshhold; //per marker, threshhold for allelic differences
    vector<DoubleVec2d> m_steps; // pre-computed constants for mutations of
    // various sizes: nbins x nbins.
    // Per marker, because it uses m_threshhold.
    long m_allowance;      // extra bins to each side of actual markers

    // helper functions for likelihood
    // 2003/11/07 erynes: Changing member variable "threshhold" from hardcoded
    //                    long to marker-specific LongVec1d required the passing
    //                    of the current marker to these functions.
    double Probability(double t, long diff, long marker) const;
    void CalculateSteps(long marker);

    virtual void Initialize(const StringVec2d& uniqueAlleles);
  public:
    // accepting compiler constructed copy-ctor and operator=

    MixedKSModel(
        long nmarkers,
        const StringVec2d& uniqueAlleles,
        long numCategories,
        DoubleVec1d categoryRates,
        DoubleVec1d categoryProbabilities,
        double userAutoCorrelationValue,
        bool doNormalize,
        double relMuRate,
        double alphaVal,
        bool doOptimization);

    virtual ~MixedKSModel() {};
    virtual DataModel* Clone() const;

    virtual model_type GetModelType() const { return MixedKS; };

    virtual vector<double> DataToLikes(const string& datum, long marker = 0) const;

    // Data Likelihood Calculation
    virtual void ComputeSiteDLs(Cell_ptr child1,Cell_ptr child2,Cell_ptr thiscell,
                                const DoubleVec1d& vv1, const DoubleVec1d& vv2,
                                long marker);
    virtual double ComputeSubtreeDLs(Cell& rootdls,
                                     double** startmarker, double** endmarker,
                                     long posn);

    // Factories
    virtual StringVec1d CreateDataModelReport() const;
    virtual StringVec1d ToXML(std::size_t nspaces) const;

    double GetAlpha() const { return m_alpha; };

    void SetAlpha(double alphaVal);
    virtual void SetAlpha(double alphaVal, long rep, long chain);
    virtual void ResetAlpha();

    virtual DoubleVec1d SimulateMarker(double branchlength, long whichmarker,
                                       const DoubleVec1d& state) const;
    virtual bool OptimizeDataModel(Tree* tree, const Locus& locus);
    virtual void WriteAlpha(std::ofstream&, long loc, long rep, long chain);
    virtual string CellToData(Cell_ptr cell, long marker) const;

  private:

    // helper functions for likelihood used in KAllele
    double probMathFunc( double ut, double coef);
    double m_alpha;
    double m_origAlpha;
    double m_beta;        //beta = 1 - alpha;
    bool m_isOpt;         //whether to optimize
    DoubleVec2d m_alphaRpt;
};

#endif // DLMODEL_H

//____________________________________________________________________________________
