// $Id: likelihood.h,v 1.60 2018/01/03 21:33:02 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 */


// likelihood class --------------------------------------------------
//
// PostLike
//     --> SinglePostLike      [single region, single chain]
//     --> ReplicatePostLike   [single region, multiple replicate]
//     --> RegionPostLike      [multiple region, multiple replicate]
//     --> GammaRegionPostLike [mult. reg., mult. rep., with background
//                              mutation rates gamma-distributed among regions]

#ifndef LIKELIHOOD_H
#define LIKELIHOOD_H

#include <functional>
#include <vector>
#include "vectorx.h"
#include "defaults.h"
#include "paramstat.h"

// now necessary to handle the evil kludge for stick rearrangement
// the commented out forward declaration is all that is otherwise
// necessary
#include "plforces.h"
// class PLForces

class ForceParameters;
class Maximizer;
class TreeSummary;
class TreeCollector;
class GammaRegionPostLike;
class ForceSummary;

class PostLike
{
  protected:
    unsigned long m_nRegions;
    unsigned long m_nReplicate;
    unsigned long m_nParam;
    unsigned long m_nForces;
    GammaRegionPostLike *m_pGammaRegionPostLike;
    vector <force_type> m_forcestags;
    vector <ParamStatus> m_default_pstatusguides; // set by Maximizer, used in derivatives
    vector <ParamStatus> m_working_pstatusguides; // ditto

    // PLforces objects
    // non-owning pointers (objects are owned by Force class)
    vector < PLForces * >m_forces;

    LongVec1d m_parameter_types; //specifies to which force a parameter belongs
    // is a simple offset list from forces.begin()
    // parameter_type is filled in FillForces()

    long m_growthstart;
    long m_growthend;
    long m_s_is_here;

    vector<double> m_minvalues;
    vector<double> m_maxvalues;

    long FindForce(long whichparam) { return m_parameter_types[whichparam];} ;

    bool Calc_sumG_ProbGP_over_ProbGPo(const DoubleVec1d& param,
                                       const DoubleVec1d& lparam,
                                       const DoubleVec1d& ln_ProbGPo,
                                       DoubleVec1d& ln_ProbGP_over_ProbGPo,
                                       double& sumG_ProbGP_over_ProbGPo,
                                       const vector<TreeSummary*> *data);

    double DCalc_sumG_BasicNumerator(const vector<double>& param,
                                     const vector<double>& ln_ProbGP_over_ProbGPo,
                                     const vector<TreeSummary*> *pGenealogies,
                                     const long& whichparam);

    // Changes the many-vector parameter list into a 1-D vector.
    void Linearize(const ForceParameters * pForceParameters,
                   DoubleVec1d & param0);

    virtual likelihoodtype GetTag() = 0; // enum

  public:
    PostLike(const ForceSummary &thisforces,
             const long thisNregion,
             const long thisNreplicate, const long thisNParam);
    virtual ~PostLike (void);

    // Sets up data connection
    virtual long GetNparam(void)
    {
        return m_nParam;
    };
    virtual DoubleVec1d GetMeanParam0() = 0;
    virtual void SetParam0() = 0;
    bool isLinearParam(long whichparam);
    bool isLogisticSelection(long whichparam) { return 0 != m_s_is_here && whichparam == m_s_is_here; };

    friend class Maximizer; // for speed, maximizer writes to postlike's gradient guide

    // Interface for calculating the log-likelihood.
    virtual bool Calculate(const DoubleVec1d& param,
                           const DoubleVec1d& lparam,
                           double& lnL) = 0;

    // Calculates the log of the probability of a genealogy given the parameters,
    // i.e., log(Prob(G|P)).
    // This method is public for use by Bayesian estimation.
    double Calc_lnProbGP(const DoubleVec1d& param,
                         const DoubleVec1d& lparam,
                         const TreeSummary *treedata);

    // calculates derivates
    virtual bool DCalculate(const DoubleVec1d& param,
                            DoubleVec1d& gradient) = 0;

    // calculates the of log of the probability of a genealogy given a
    // stick.
    // This method is public for use by StairArranger.
    double Calc_lnProbGS(const DoubleVec1d& param,
                         const DoubleVec1d& lparam,
                         const TreeSummary* treedata);

    const DoubleVec1d& GetMinValues() const { return m_minvalues; };
    const DoubleVec1d& GetMaxValues() const { return m_maxvalues; };

    // evil kludge to give stair rearrangement access to the Mean and
    // Variance
    const StickSelectPL& GetStickSelectPL() const;

  private:
    // fills the forces classes into the forces vector
    void FillForces (const ForceSummary&thisforces);
    PostLike();
};

//------------------------------------------------------------------------------------

class SinglePostLike : public PostLike
{
  private:
    TreeCollector *m_data;  // main hook into the treesummary data
    DoubleVec1d m_param0;
    DoubleVec1d m_lparam0;
    DoubleVec1d m_lnProbGPo;      //this remains constant through maximization
    DoubleVec1d m_ln_ProbGP_over_ProbGPo; //this does NOT! It is constant for a
    // specific parameter.
    // "G" stands for "genealogy."
    double m_last_lnL;
    double m_totalNumTrees;
    SinglePostLike();

  protected:
    virtual likelihoodtype GetTag() { return ltype_ssingle; }; // enum

  public:
    SinglePostLike(const ForceSummary &thisforces,
                   const long thisNregion,
                   const long thisNreplicate, const long thisNParam);
    ~SinglePostLike();
    void SetParam0();
    DoubleVec1d GetParam0() { return m_param0; };
    virtual void Setup(TreeCollector * treedata);
    bool Calculate(const DoubleVec1d& param,
                   const DoubleVec1d& lparam,
                   double& lnL);
    bool DCalculate(const DoubleVec1d& param,
                    DoubleVec1d& gradient);
    DoubleVec1d GetMeanParam0 ();
};

//------------------------------------------------------------------------------------

class ReplicatePostLike : public PostLike
{
  private:
    DoubleVec2d m_param0;
    DoubleVec2d m_lparam0;
    vector<TreeCollector*> m_data;
    DoubleVec2d m_lnProbGPo;
    DoubleVec2d m_ln_ProbGP_over_ProbGPo;
    double m_last_lnL;
    DoubleVec3d m_pg0; // rep X tree X rep, filled in by Setup(),
    // used by FillGeyerWeights
    DoubleVec1d m_logGeyerWeights;

    bool FillGeyerWeights();
    bool GeyerLike(const DoubleVec1d& param, const DoubleVec1d& lparam,
                   double &newlike, const DoubleVec1d& oldlike);
    bool GeyerStart(DoubleVec1d &lnL);
    void TransformProbG0With(const DoubleVec1d& weights);
    ReplicatePostLike();

  protected:
    virtual likelihoodtype GetTag() { return ltype_replicate; }; // enum

  public:
    ReplicatePostLike(const ForceSummary &thisforces,
                      const long thisNregion,
                      const long thisNreplicate, const long thisNParam);
    //~ReplicatePostLike() {};
    void SetParam0();
    void CreateProbG0();
    virtual bool Setup(const vector<TreeCollector*>& treedata);
    bool Calculate(const DoubleVec1d& param,
                   const DoubleVec1d& lparam,
                   double& lnL);
    bool DCalculate(const DoubleVec1d& param,
                    DoubleVec1d& gradient);
    DoubleVec1d GetMeanParam0();
    DoubleVec1d GetGeyerWeights() {return m_logGeyerWeights;};
};

//------------------------------------------------------------------------------------

class RegionPostLike : public PostLike
{
  protected:
    DoubleVec3d m_param0; // the "driving values" for each rep and reg
    DoubleVec3d m_lparam0; // the natural logarithms of the above
    vector<vector<TreeCollector *> > m_data; // the trees for each rep and reg
    DoubleVec3d m_lnProbGPo; // Dimensions: numtrees x reps x regions.  May include Geyer.
    DoubleVec3d m_ln_ProbGP_over_ProbGPo;  // Same dimensions as Prob(G|Po).
    DoubleVec1d m_sumG_ProbGP_over_ProbGPo; // One value per region.  May include Geyer.
    DoubleVec2d m_paramscalars; // Scale the eff. pop. size when we have nuc DNA and mtDNA, etc.
    DoubleVec2d m_lparamscalars; // Natural logarithms of the above scaling factors.
    DoubleVec1d m_totalNumTrees; // One value per region.
    double m_last_lnL; // Store the total lnL for later use.
    DoubleVec4d m_pg0; // reg X rep X tree X rep, filled in by Setup
    // used by FillGeyerWeights

    void TransformProbG0With(const DoubleVec1d& weights, long region);

    virtual likelihoodtype GetTag() { return ltype_region; }; // enum

  public:
    RegionPostLike(const ForceSummary &thisforces,
                   const long thisNregion,
                   const long thisNreplicate, const long thisNParam,
                   DoubleVec2d paramscalars);
    //~RegionPostLike() {};
    void SetParam0();
    virtual void Setup(const vector<vector<TreeCollector*> >& treedata,
                       const DoubleVec2d& logGeyerWeights);
    bool Calculate(const DoubleVec1d& param,
                   const DoubleVec1d& lparam,
                   double& lnL);
    bool DCalculate(const DoubleVec1d& param,
                    DoubleVec1d& gradient);
    DoubleVec1d GetMeanParam0();

  private:
    RegionPostLike();
};

//------------------------------------------------------------------------------------

class GammaRegionPostLike : public RegionPostLike
{
  protected:
    // Note:  All of the following variables of type DoubleVec3d
    // have the dimensions of numtrees x reps x regions.
    DoubleVec3d m_nevents_G; // The total number of events in tree G.
    DoubleVec3d m_lnPoint; // log(PointProb(G|P)).
    DoubleVec3d m_lnWait;  // log(WaitProb(G|P)).
    DoubleVec3d m_C;      // NOTE:  This takes the place of m_ln_ProbGP_over_ProbGPo.
    DoubleVec1d m_sumG_C; // NOTE:  This takes the place of m_sumG_ProbGP_over_ProbGPo.

    DoubleVec3d m_K_alpha_minus_nevents;
    // K(alpha - nevents, 2.0*sqrt(-alpha*m_lnWait)),
    // where K(v,x) is the modified Bessel fn. of the 2nd kind
    // of order v evaluated at x.
    DoubleVec3d m_K_alpha_minus_nevents_plus_1;
    // K(alpha - nevents + 1, 2.0*sqrt(-alpha*m_lnWait)).

    virtual likelihoodtype GetTag() { return ltype_gammaregion; }; // enum

    // Iterators over m_lnPoint and m_lnWait.
    vector<vector<vector<double> > >::iterator m_lnPoint_reg_it;
    vector<vector<double> >::iterator m_lnPoint_rep_it;
    vector<double>::iterator m_lnPoint_genealogy_it;
    vector<vector<vector<double> > >::iterator m_lnWait_reg_it;
    vector<vector<double> >::iterator m_lnWait_rep_it;
    vector<double>::iterator m_lnWait_genealogy_it;

  public:
    GammaRegionPostLike(const ForceSummary &theseforces,
                        const long thisNregion,
                        const long thisNreplicate, const long thisNParam,
                        DoubleVec2d paramscalars);
    //~GammaRegionPostLike() {};

    void Store_Current_lnPoint(const double& value);
    void Store_Current_lnWait(const double& value);
    virtual void Setup(const vector<vector<TreeCollector*> >& treedata,
                       const DoubleVec2d& logGeyerWeights);
    bool Calculate(const DoubleVec1d& param,
                   const DoubleVec1d& lparam,
                   double& lnL);
    bool DCalculate(const DoubleVec1d& param,
                    DoubleVec1d& gradient);

  private:
    GammaRegionPostLike();
};

#endif // LIKELIHOOD_H

//____________________________________________________________________________________
