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


#ifndef ARRANGER_H
#define ARRANGER_H

#include <cassert>
#include <cmath>
#include <map>
#include <string>
#include <vector>

#include "vectorx.h"
#include "branchbuffer.h"
#include "arranger_types.h"
#include "prior.h"                      // EpochSizeArranger contains a vector<Prior>

// #include "event.h"    uses Event methods in .cpp
// #include "tree.h"
// #include "forceparam.h"
// #include "forcesummary.h"
// #include "random.h"

class Event;
class Tree;
class ForceParameters;
class Random;
class ForceSummary;
class ChainState;
class TimeList;
class Epoch;

/*******************************************************************
 The Arranger class transforms a tree and decides whether the new ree should be accepted or rejected.
 It provides a routine DenovoTree() to make an independent tree, and a routine Rearrange to
 make a tree based on a previous tree.

 Subclasses of Arranger carry out different types of transformation.

 For all Arrangers that modify the tree, as a postcondition the Range information in all Branches
 remaining in the tree must be correct before Tree::Prune is called.  The ResimulatingArranger as
 currently written achieves this by making the Ranges correct at the "active face" or current point
 of rearrangement.  This would be problematic in a recombinant tree if rearrangement stopped before
 the root was reached, as lower Branches would not be made correct.  However, the InactiveRecEvent
 will not allow rearrangement to terminate until the root is reached.

 Written by Jim Sloan, revised by Mary Kuhner

 BayesArranger added 11/26/03 Mary.  This Arranger changes the parameter values rather than the tree.

 DenovoArranger added 11/26/03 Mary.  This Arranger makes trees from scratch ("de novo") and replaces
 the Denovo member function which could not be implemented on most Arrangers.

 RecSiteArranger removed 3/29/06 Mary as it was a deteriorating pseudogene.

 StairArranger added 4/4/07 Mary.  This Arranger changes the stairway of doom (AKA wiggling stick).

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

class Arranger
{
  private:
    Arranger();                           // undefined
    Arranger& operator=(const Arranger&); // undefined

  protected:
    Arranger(const Arranger& src);

    double  m_timing, m_savetiming;

  public:
    Tree   *m_tree;             // public for speedy access by Event subclasses; also set by Chain::StartRegion()
    Random *randomSource;

    Arranger(double timing);
    virtual ~Arranger() {};

    // does not copy the tree pointer!
    virtual Arranger* Clone() const = 0;

    // Arrangement functions
    virtual void    SetParameters(ChainState & chainstate) = 0;
    virtual void    Rearrange(ChainState & chstate) = 0;
    virtual void    ScoreRearrangement(ChainState & chstate) = 0;
    virtual bool    AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree) = 0;

    // Getter/Setters
    Tree *  GetTree() const          { return m_tree; };
    double  GetTiming() const        { return m_timing; };
    virtual std::string  GetName() const     { return arrangerstrings::BASE; };

    void    SetTiming(double t);
    void    SetSaveTiming()             { m_savetiming = m_timing; };
    void    RestoreTiming()             { m_timing = m_savetiming; };
};

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

/*********************************************************************
 ResimArranger is an abstract base class for Arrangers which do any form of resimulation of lineages
 downward through the tree.  It is tightly coupled with the Event class, which provides expertise needed
 in resimulation.  The eventvec is a vector of all available types of Events (depending on forces and
 other strategy details) which ResimArranger polls to conduct its resimulation.
***********************************************************************/

class ResimArranger : public Arranger
{
  protected:
    std::vector<Event*> m_eventvec;    // code for various event types

    // helper functions
    ResimArranger(const ResimArranger& src);
    void           ClearEventVec();
    virtual void   DropAll(double eventT);
    virtual double EventTime(Event*& returnevent, double eventT, double nextT);
    virtual double Activate(Tree * oldtree);
    virtual void   Resimulate(double eventT, ChainState & chainstate);
    virtual bool   StopNow() const;
    double         NextInterval(double lastT, FC_Status & fcstatus);
    virtual void   CleanupAfterResimulate();

    bool m_hasLogisticSelection;
    double m_logsel_cutoff;

  public:
    // The following variables are helpers for use by the subclasses
    // of Event.  They are public because otherwise Event would have to be a
    // friend, and this would tie Arranger to the specific subclasses of Event.
    LongVec1d m_xactives;               // dim: cross partitions
    LongVec2d m_pactives;               // dim: part-force X partitions
    LongVec1d m_xinactives;             // dim: cross partitions
    LongVec2d m_pinactives;             // dim: part-force X partitions

    BranchBuffer m_activelist;          // lineages currently active
    BranchBuffer m_inactivelist;        // inactive lineages in current interval

    ResimArranger(const ForceSummary & fs, double timing);
    virtual ~ResimArranger() = 0;       // "implemented pure virtual"

    // Arrangement functions
    virtual void   SetParameters(ChainState & chainstate);
    virtual void   Rearrange(ChainState & chstate);
    virtual void   ScoreRearrangement(ChainState & chstate);
    virtual bool   AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree);
    virtual double Hastings(ChainState & chstate);

    long   ActiveSize() const { return m_activelist.Size(); };

    // Getters/Setters
    virtual std::string GetName() const { return arrangerstrings::RESIM; };
};

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

class DropArranger: public ResimArranger
{
  protected:
    DropArranger(const DropArranger& src) :
        ResimArranger(src) {};

  public:
    DropArranger(const ForceSummary & fs, double timing)
        : ResimArranger(fs, timing) {};
    virtual ~DropArranger() {};
    virtual Arranger*  Clone() const;
    virtual std::string     GetName() const { return arrangerstrings::DROP; };

};

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

class DenovoArranger: public ResimArranger
{
  protected:
    DenovoArranger(const DenovoArranger& src) :
        ResimArranger(src) {};

  public:
    DenovoArranger(const ForceSummary & fs, double timing)
        : ResimArranger(fs, timing) {};
    virtual  ~DenovoArranger() {};
    virtual Arranger*  Clone() const;
    virtual void       Rearrange(ChainState & chstate);
    virtual bool       AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree);
    virtual std::string     GetName() const { return arrangerstrings::DENO; };

};

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

class BaseHapArranger : public Arranger
{
  protected:
    BaseHapArranger(const BaseHapArranger& src)
        : Arranger(src) {};

  public:
    BaseHapArranger(double timing) : Arranger(timing)   {};
    virtual void      SetParameters(ChainState & chainstate);
    virtual void      Rearrange(ChainState & chstate) = 0;
    virtual void      ScoreRearrangement(ChainState & chstate);
    virtual bool      AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree);
    virtual Arranger* Clone() const = 0;
    virtual std::string    GetName() const = 0;
};

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

class HapArranger : public BaseHapArranger
{
  protected:
    HapArranger(const HapArranger& src)
        : BaseHapArranger(src) {};

  public:
    HapArranger(double timing) : BaseHapArranger(timing) {};
    virtual void      Rearrange(ChainState & chstate);
    virtual Arranger* Clone() const;
    virtual std::string    GetName() const { return arrangerstrings::HAP; };
};

//------------------------------------------------------------------------------------
// The 'Probability Haplotype Arranger' (which I'm calling it here in the code
//  because that's what it is; we'll probably call it the 'Trait Haplotype
//  Arranger' for the user) is the arranger to use when you have haplotypes
//  that need to be swapped among a set instead of among just two alternating
//  options.
//  It's being added here for use for trait data with sets of possible
//  haplotype resolutions (i.e. HH, Hh, and hH).

class ProbHapArranger : public BaseHapArranger
{
  protected:
    ProbHapArranger(const ProbHapArranger& src)
        : BaseHapArranger(src) {};

  public:
    ProbHapArranger(double timing):BaseHapArranger(timing) {};
    virtual void      Rearrange(ChainState & chstate);
    virtual void      ScoreRearrangement(ChainState & chstate);
    virtual Arranger* Clone() const;
    virtual std::string    GetName() const { return arrangerstrings::PROBHAP; };
};

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

class TreeSizeArranger : public ResimArranger
{
  protected:
    TreeSizeArranger(const TreeSizeArranger& src) : ResimArranger(src) {};
    virtual double Activate(Tree * oldtree);
    virtual void   Resimulate(double eventT, ChainState & chainstate);
    virtual void   CleanupAfterResimulate();
    bool NodeChangesEpoch(Branch_ptr pBranch, double newtime, const DoubleVec1d & epochtimes) const;

  public:
    TreeSizeArranger(const ForceSummary & fs, double timing);
    virtual ~TreeSizeArranger() {};
    virtual double     Hastings(ChainState & chstate);
    virtual Arranger*  Clone() const;
    virtual std::string GetName() const { return arrangerstrings::TREESIZE; };
};

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

class EpochSizeArranger : public ResimArranger
{
  private:
    bool m_tree_changed; // did our rearrangement touch the tree?

  protected:
    const std::vector<Epoch> * m_epochs;  // non-owning pointer
    vector<Prior> m_priors;

    EpochSizeArranger(const EpochSizeArranger& src);
    // we disavow the following functions, ugh, but it's necessary
    virtual double Activate(Tree *) { assert(false); return 0.0; };
    virtual void Resimulate(double, ChainState &) { assert(false); };
    virtual void CleanupAfterResimulate() { assert(false); };

  public:
    EpochSizeArranger(const ForceSummary & fs, double timing);
    virtual ~EpochSizeArranger() {};
    virtual double     Hastings(ChainState & chstate);
    virtual Arranger*  Clone() const;
    virtual std::string GetName() const { return arrangerstrings::EPOCHSIZE; };

    // Arrangement functions.
    // We accept ResimArranger::SetParameters().
    virtual void    Rearrange(ChainState & chstate);
    virtual void    ScoreRearrangement(ChainState & chstate);
    virtual bool    AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree);
};

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

class EpochNudgeArranger : public ResimArranger
{
  private:
    bool m_tree_changed; // did our rearrangement touch the tree?

  protected:
    const std::vector<Epoch> * m_epochs;  // non-owning pointer
    vector<Prior> m_priors;

    EpochNudgeArranger(const EpochNudgeArranger& src);
    // we disavow the following functions, ugh, but it's necessary
    virtual double Activate(Tree *) { assert(false); return 0.0; };
    virtual void Resimulate(double, ChainState &) { assert(false); };
    virtual void CleanupAfterResimulate() { assert(false); };

  public:
    EpochNudgeArranger(const ForceSummary & fs, double timing);
    virtual ~EpochNudgeArranger() {};
    virtual double     Hastings(ChainState & chstate);
    virtual Arranger*  Clone() const;
    virtual std::string GetName() const { return arrangerstrings::EPOCHSIZE; };

    // Arrangement functions.
    // We accept ResimArranger::SetParameters().
    virtual void    Rearrange(ChainState & chstate);
    virtual void    ScoreRearrangement(ChainState & chstate);
    virtual bool    AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree);
};

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

class BayesArranger : public Arranger
{
  private:
    DoubleVec1d m_oldlogs;
    DoubleVec1d m_newlogs;

  protected:
    BayesArranger(const BayesArranger& src);

  public:
    BayesArranger(double timing) : Arranger(timing) {};
    virtual void      SetParameters(ChainState & chainstate);
    virtual void      Rearrange(ChainState & chstate);
    virtual void      ScoreRearrangement(ChainState & chstate);
    virtual bool      AcceptAndSynchronize(ChainState &, double temperature, bool badtree);
    virtual Arranger* Clone() const;
    virtual std::string GetName() const { return arrangerstrings::BAYES; };
};

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

class LocusArranger : public Arranger
{
  protected:
    // LocusArranger(const LocusArranger& src);  //we accept the default.

  public:
    LocusArranger(double timing) : Arranger(timing) {};
    virtual void      SetParameters(ChainState & chainstate);
    virtual void      Rearrange(ChainState & chstate);
    virtual void      ScoreRearrangement(ChainState & chstate);
    virtual bool      AcceptAndSynchronize(ChainState &, double temperature, bool badtree);
    virtual Arranger* Clone() const;
    virtual std::string GetName() const { return arrangerstrings::LOCUS; };
};

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

class ZilchArranger: public Arranger
{
  public:
    ZilchArranger(double timing) : Arranger(timing) {};
    virtual      ~ZilchArranger() {};
    virtual void SetParameters(ChainState &) {};
    virtual void Rearrange(ChainState &) {};
    virtual void ScoreRearrangement(ChainState &) {};
    virtual bool AcceptAndSynchronize(ChainState &, double temperature, bool badtree);
    virtual Arranger* Clone() const;
    virtual std::string GetName() const { return arrangerstrings::ZILCH; };
};

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

class StairArranger: public Arranger
{
  protected:
    double LnPBrownStickParams(const DoubleVec1d & param, const TreeSummary * partialtreedata) const;

  public:
    StairArranger(double timing) : Arranger(timing) {};
    virtual ~StairArranger() {};
    virtual Arranger* Clone() const;
    virtual std::string GetName() const { return arrangerstrings::STICK; };

    // Arrangement functions
    virtual void SetParameters(ChainState & chainstate);
    virtual void Rearrange(ChainState & chstate);
    virtual void ScoreRearrangement(ChainState & chstate);
    virtual bool AcceptAndSynchronize(ChainState & chstate, double temperature, bool badtree);
};

#endif // ARRANGER_H

//____________________________________________________________________________________
