// $Id: plforces.h,v 1.41 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 */


// PLforces class --------------------------------------------------
//
// know how to calculate waiting times and point probabilities
// of the parameters
//

#ifndef PLFORCES_H
#define PLFORCES_H

#include <algorithm>
#include <cassert>
#include <cmath>
#include <functional>
#include <numeric>
#include <string>
#include <vector>

#include "treesum.h"
#include "summary.h"

class PLForces
{
  private:
    PLForces ();  // not implemented

  protected:
    PLForces(const PLForces& src) :
        m_nPop(src.m_nPop), m_pTreedata(src.m_pTreedata), m_start(src.m_start),
        m_end(src.m_end), m_growthstart(src.m_growthstart)
    { };

    long m_nPop;
    TreeSummary *m_pTreedata;    // non-owning pointer
    long m_start;        //offsets into the param vector that does not yet exist
    long m_end;

    long m_growthstart; //we need to know where the growth parameters are stored
    // we need this in all forces to handle the scaled times.

  public:
    PLForces (long thisnpop)
    {
        m_nPop = thisnpop;
        m_start = 0L;
        m_end = 0L;
        m_growthstart = 0L;
    };

    virtual ~ PLForces () {};

    virtual PLForces* Clone() const = 0;

    virtual double lnWait (const vector < double >&param,
                           const TreeSummary * treedata) = 0;

    virtual double lnPoint (const vector < double >&param,
                            const vector < double >&lparam,
                            const TreeSummary * treedata) = 0;

    virtual double DlnWait (const vector < double >&param,
                            const TreeSummary * treedata,
                            const long &whichparam) = 0;

    virtual double DlnPoint (const vector < double >&param,
                             const TreeSummary * treedata,
                             const long &whichparam) = 0;

    virtual long GetNparam() { return m_nPop; };
    long GetStart()  { return m_start; }; // used by Force::FindOrdinalPosition()

    virtual void SetStart(const long& start) { m_start = start; };
    virtual void SetEnd  (const long& end  ) { m_end   = end;   };

    virtual void SetGrowthStart(const long& growthstart)
    { m_growthstart = growthstart; };
};

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

class CoalescePL:public PLForces
{
  private:
    CoalescePL ();  // not implemented

  protected:
    // dim: lpforce X xpart
    vector<vector<DoubleVec1d::size_type> > m_whichlocalpart;

    // dim: pforce X xpart
    vector<vector<DoubleVec1d::size_type> > m_whichpart;

    // dim: lpforce X part X xpart-list
    vector<vector<vector<DoubleVec1d::size_type> > > m_whichlocalxparts;

    // dim: pforce X part X xpart-list
    vector<vector<vector<DoubleVec1d::size_type> > > m_whichxparts;

    ForceVec m_localpartforces;

    double CalculateScaledLPCounts(const DoubleVec1d& params,
                                   const LongVec2d& picks) const;

    CoalescePL(const CoalescePL& src) : PLForces(src) { };

  public:
    // coalescence related wait, and point functions
    // and derivatives of them
    ~CoalescePL () { };

    CoalescePL (long thisnpop) : PLForces(thisnpop) {};

    virtual PLForces* Clone() const
    {
        return new CoalescePL(*this);
    };

    void SetLocalPartForces(const ForceSummary& fs);

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata,
                    const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata,
                     const long &whichparam);
};

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

// subclass of CoalescePL that uses long form with growth
class CoalesceGrowPL:public CoalescePL
{
  private:
    CoalesceGrowPL();  // not implemented

  protected:
    TimeManager* m_timesize;

    CoalesceGrowPL(const CoalesceGrowPL& src) :
        CoalescePL(src)
    { };

    double lnWaitForTinyGrowth(const double theta, const double growth,
                               const unsigned long xpartition,
                               const std::list<Interval>& treesummary);

    double DlnWaitForTinyGrowth(const vector < double >&param,
                                const TreeSummary * treedata,
                                const long &whichparam);
  public:
    ~CoalesceGrowPL () { };

    CoalesceGrowPL (long thisnpop) : CoalescePL(thisnpop) {};

    virtual PLForces* Clone() const
    {
        return new CoalesceGrowPL(*this);
    };

    void SetTimeManager(TimeManager* newtimesize) {m_timesize = newtimesize;};

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata, const long &whichparam);

    //  DlnPoint is growth-independent, so just use CoalescePL::DlnPoint
    //  NB:  not any more, with growth+disease+recombination!
    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata, const long &whichparam);
};

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

// exponential growth coalescence related wait, and point functions
// and derivatives of them
class GrowPL:public PLForces
{
  private:
    GrowPL(); // not implemented

  protected:
    GrowPL(const GrowPL& src) :
        PLForces(src), m_thetastart(src.m_thetastart)
    { };

    double DlnWaitForTinyGrowth(const vector < double >&param,
                                const TreeSummary * treedata,
                                const long &whichparam);

    long m_thetastart;

  public:
    ~GrowPL () { };

    GrowPL (long thisnpop) : PLForces(thisnpop)
    {
        m_thetastart = 0;
    };

    virtual PLForces* Clone() const
    {
        return new GrowPL(*this);
    };

    virtual void SetThetaStart(const long& thetastart)
    { m_thetastart = thetastart; };

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam, const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata, const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata, const long &whichparam);
};

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

class DiseasePL:public PLForces
{
  private:
    DiseasePL(); // not implemented

  protected:
    DoubleVec1d m_msum;
    DiseasePL(const DiseasePL& src) : PLForces(src), m_msum(src.m_msum) { };

  public:
    // disease related wait, and point functions
    // and derivatives of them

    DiseasePL (long thisnstati) : PLForces(thisnstati), m_msum(thisnstati,0.0) {}

    ~DiseasePL () { };

    virtual PLForces* Clone() const
    {
        return new DiseasePL(*this);
    };

    double lnWait (const vector < double >&param, const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam, const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata, const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata, const long &whichparam);
};

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

class DiseaseLogisticSelectionPL:public DiseasePL
{
  private:
    DiseaseLogisticSelectionPL(); // not implemented

  protected:
    long m_s_is_here;

    DiseaseLogisticSelectionPL(const DiseaseLogisticSelectionPL& src) :
        DiseasePL(src)
    { };

    double DlnWaitForTinyLogisticSelectionCoefficient(const vector<double>& param,
                                                      const TreeSummary *treedata,
                                                      const long &whichparam);
  public:
    // disease related wait, and point functions
    // and derivatives of them
    DiseaseLogisticSelectionPL(long thisnstati) : DiseasePL(thisnstati) { };

    ~DiseaseLogisticSelectionPL() { };

    virtual PLForces* Clone() const
    {
        return new DiseaseLogisticSelectionPL(*this);
    };

    void SetSelectionCoefficientLocation(long paramvecindex)
    {
        m_s_is_here = paramvecindex;
    };

    double lnWait (const vector < double >&param, const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam, const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata, const long &whichparam);

    // DlnPoint is selection-independent, so just use DiseasePL::DlnPoint().
};

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

class MigratePL:public PLForces
{
  private:
    MigratePL (); // not implemented

  protected:
    DoubleVec1d m_msum;
    MigratePL(const MigratePL& src) : PLForces(src), m_msum(src.m_msum) { };

  public:
    // migration related wait, and point functions
    // and derivatives of them

    MigratePL (long thisnpop) : PLForces(thisnpop)
    {
        m_msum = CreateVec1d(m_nPop, static_cast<double>(0.0));
    };

    ~MigratePL () { };

    virtual PLForces* Clone() const
    {
        return new MigratePL(*this);
    };

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata,
                    const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata,
                     const long &whichparam);
};

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

class DivMigPL:public PLForces
{
  private:
    DivMigPL ();  // not implemented

  protected:
    DivMigPL(const DivMigPL& src) : PLForces(src) { };

  public:
    // migration related wait, and point functions
    // and derivatives of them

    DivMigPL (long thisnpop) : PLForces(thisnpop) {};

    ~DivMigPL () { };

    virtual PLForces* Clone() const
    {
        return new DivMigPL(*this);
    };

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);
    // NB:  Divergence is currently incompatible with likelihood
    // evaluation, so the following two functions are not used.
    // If they are ever needed, implementation is very similar to the
    // MigPL versions.
    double DlnWait (const vector < double >&,
                    const TreeSummary *,
                    const long &) { assert(false); return 0;};

    double DlnPoint (const vector < double >&,
                     const TreeSummary *,
                     const long &) { assert(false); return 0;};
};

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

class RecombinePL:public PLForces
{
  private:
    RecombinePL ();  // not implemented

  protected:
    RecombinePL(const RecombinePL& src) : PLForces(src) { };

  public:
    // recombination related wait, and point functions
    // and derivatives of them
    RecombinePL(long thisnpop) : PLForces(thisnpop) {};

    ~RecombinePL () { };

    virtual PLForces* Clone() const
    {
        return new RecombinePL(*this);
    };

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata,
                    const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata,
                     const long &whichparam);
};

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

class DivPL: public PLForces
// Yes, this class does NOTHING!  (except for keeping track of its
// parameters, which makes the maximizer happy)
{
  private:
    DivPL(); // not implemented

  protected:
    DivPL(const DivPL& src) : PLForces(src) {};

  public:
    DivPL(long nparam) : PLForces(nparam) {};

    ~DivPL() {};

    virtual PLForces* Clone() const
    { return new DivPL(*this); };

    double lnWait (const vector < double >&,
                   const TreeSummary * ) { return 1.0; };

    double lnPoint (const vector < double >&,
                    const vector < double >&,
                    const TreeSummary *) { return 1.0; };

    double DlnWait (const vector < double >&,
                    const TreeSummary *,
                    const long &) { return 1.0; };

    double DlnPoint (const vector < double >&,
                     const TreeSummary *,
                     const long &) { return 1.0; };
};

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

class SelectPL:public PLForces
{
  private:
    SelectPL(); // not implemented

  protected:
    SelectPL(const SelectPL& src) : PLForces(src) { };

  public:
    // selection related wait, and point functions
    // and derivatives of them
    virtual PLForces* Clone() const
    {
        return new SelectPL(*this);
    };

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata,
                    const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata,
                     const long &whichparam);
};

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

class CoalesceLogisticSelectionPL:public CoalescePL
{
  private:
    CoalesceLogisticSelectionPL(); // not implemented

  protected:
    CoalesceLogisticSelectionPL(const CoalesceLogisticSelectionPL& src) : CoalescePL(src),
                                                                          m_s_is_here(src.m_s_is_here) { };

    long m_s_is_here;

    double lnWaitForTinyLogisticSelectionCoefficient(const vector<double>& param,
                                                     const TreeSummary *treedata);

    double DlnWaitForTinyLogisticSelectionCoefficient(const vector<double>& param,
                                                      const TreeSummary *treedata,
                                                      const long &whichparam);

  public:
    ~CoalesceLogisticSelectionPL () { };

    // perhaps move to the .cpp file?
    CoalesceLogisticSelectionPL (long thisnpop) : CoalescePL(thisnpop)
    {
        if (2 != thisnpop)
        {
            string msg = "Attempted to create a CoalesceLogisticSelectionPL object ";
            msg += "with " + ToString(thisnpop) + " populations.  This object can only ";
            msg += "be used with two populations, one for the favored allele and ";
            msg += "one for the disfavored allele.";
            throw implementation_error(msg);
        }
    };

    virtual PLForces* Clone() const
    {
        return new CoalesceLogisticSelectionPL(*this);
    };

    void SetSelectionCoefficientLocation(long paramvecindex)
    {
        m_s_is_here = paramvecindex;
    };

    double lnWait(const vector<double>& param,
                  const TreeSummary *treedata);

    double lnPoint(const vector<double>& param,
                   const vector<double>& lparam,
                   const TreeSummary *treedata);

    double DlnWait(const vector<double>& param,
                   const TreeSummary *treedata, const long &whichparam);

    double DlnPoint(const vector<double>& param,
                    const TreeSummary *treedata, const long &whichparam);
};

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

class LogisticSelectionPL:public PLForces
{
  private:
    LogisticSelectionPL(); // not implemented

  protected:
    LogisticSelectionPL(const LogisticSelectionPL& src) : PLForces(src), m_s_is_here(src.m_s_is_here) { };
    long m_s_is_here;

    double DlnWaitForTinyLogisticSelectionCoefficient(const vector<double>& param,
                                                      const TreeSummary *treedata,
                                                      const long &whichparam);

  public:
    ~LogisticSelectionPL () { };

    // MDEBUG the 0 is arbitrary--apparently this class doesn't use npop
    LogisticSelectionPL(long paramvecindex) : PLForces(0)
    {
        m_s_is_here = paramvecindex;
    };

    virtual PLForces* Clone() const
    {
        return new LogisticSelectionPL(*this);
    };

    double lnWait(const vector<double>& param,
                  const TreeSummary *treedata);

    double lnPoint(const vector<double>& param,
                   const vector<double>& lparam,
                   const TreeSummary *treedata);

    double DlnWait(const vector<double>& param,
                   const TreeSummary *treedata, const long &whichparam);

    double DlnPoint(const vector<double>& param,
                    const TreeSummary *treedata, const long &whichparam);
};

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

class StickSelectPL:public PLForces
{
  private:
    StickSelectPL(); // not implemented

  public:
    // stair selection related wait, and point functions
    // and derivatives of them
    StickSelectPL(const ForceSummary& fs);

    virtual ~StickSelectPL();

    // we accept the default copy ctor and operator=

    virtual PLForces* Clone() const
    {
        return new StickSelectPL(*this);
    };

    double lnPTreeStick (const vector < double >&param,
                         const vector < double >&lparam,
                         const TreeSummary * treedata);

    double lnWait (const vector < double >&param,
                   const TreeSummary * treedata);

    double lnPoint (const vector < double >&param,
                    const vector < double >&lparam,
                    const TreeSummary * treedata);

    double DlnWait (const vector < double >&param,
                    const TreeSummary * treedata,
                    const long &whichparam);

    double DlnPoint (const vector < double >&param,
                     const TreeSummary * treedata,
                     const long &whichparam);

    void SetThetastart(long start) {m_thetastart = start;};
    void SetThetaend(long end) {m_thetaend = end;};
    void SetToSmallAIndex(long index);
    void SetToBigAIndex(long index);
    void SetSelCoeffIndex(long index) {m_s_here = index;};
    void SetRecRateIndex(long index) {m_r_here = index;};
    void SetLocalPartForces(const ForceSummary& fs);

    // these are exposed as public so that Arranger.PStickParams() can find
    // them -- the only alternative considered at this time was making them
    // free functions, which Jon rejects for now.
    double StickMean(double freqA, double s, double toA, double toa) const;
    double StickVar(double freqA, double tipfreqA, double thetaA) const;

  protected:
    long m_thetastart, m_thetaend;
    long m_toBigA_here;
    long m_toSmallA_here;
    long m_s_here;
    long m_r_here;

    // helper(s) used for stick selection
    double m_minuslnsqrt2pi;

    // helper(s) used for disease
    long m_disstart;
    long m_ndis;
    long m_disindex;

    // helpers(s) used for coalescence
    long m_nxparts;
    LongVec1d m_partindex;
    ForceVec m_localpartforces;

    // dim: lpforce X xpart
    vector<vector<DoubleVec1d::size_type> > m_whichlocalpart;

    // dim: pforce X xpart
    vector<vector<DoubleVec1d::size_type> > m_whichpart;

    // dim: lpforce X part X xpart-list
    vector<vector<vector<DoubleVec1d::size_type> > > m_whichlocalxparts;

    // dim: pforce X part X xpart-list
    vector<vector<vector<DoubleVec1d::size_type> > > m_whichxparts;

    // dim: xpart
    vector<DoubleVec1d::size_type> m_whichmigpart;

    double CommonFunctional(double freqBigA, double prevfreqBigA, double s,
                            double toBigA, double toSmallA, double length) const;

    // the ln point term covering P(Tree|Stick,params)
    double lnPointTreeStick (const vector < double >&param,
                             const vector < double >&lparam,
                             const TreeSummary * treedata);
};

#endif // PLFORCES_H

//____________________________________________________________________________________
