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


/*******************************************************************
 Class Tree represents a genealogy (not necessarily a "tree" in the recombinant cases).  It has two subtypes,
 a PlainTree (no recombination) and a RecTree (with recombination) because the recombination machinery is too
 expensive to carry if you don't need it.

 Normally all Trees used during rearrangement come from either copying an existing tree
 or copying the prototype tree in Registry.

 DEBUG:  This class is a monster.  Anything to make it simpler would be good.

 Written by Jim Sloan, heavily revised by Mary Kuhner

 02/08/05 Pulled out locus specific stuff into LocusLike class -- Mary
 2004/9/15 Killed the LocusLike class, merging into Locus
 2005/3/1 Moved alias info, reluctantly, from DLCalc to Tree (because
    DLCalc is a shared object and the alias is not shareable!) -- Mary
 2006/02/13 Begun to add jointed stick
*******************************************************************/

#ifndef TREE_H
#define TREE_H

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

#include "constants.h"
#include "definitions.h"
#include "individual.h"                 // for IndVec member
#include "locus.h"                      // for DataModel
#include "range.h"
#include "rangex.h"
#include "timelist.h"                   // for TimeList member
#include "vectorx.h"

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

// typedef also used by Event and Arranger
typedef std::pair<Branch_ptr, Branch_ptr> branchpair;

class TreeSummary;                      // return type of SummarizeTree()
class TipData;
class Random;
class ForceParameters;
class TimeManager;

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

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

    vector<LocusCell> m_protoCells;     // cache of locuscells for all loci

    vector<Branch_ptr> FindAllBranchesAtTime(double eventT);

  protected:
    Tree();
    Tree(const Tree & tree, bool makestump);
    IndVec          m_individuals;        // associates tips from same individual
    double          m_overallDL;          // data log-likelihood
    LongVec2d       m_aliases;            // dim: loci x markers
    Random *        m_randomSource;       // non-owning pointer
    TimeList        m_timeList;           // all lineages
    vector<Locus> * m_pLocusVec;          // not owning.  Also, not const because of simulation
    long int        m_totalSites;         // span of entire region, including gaps
    TimeManager *   m_timeManager;        // our timemanager....
    bool            m_hasSnpPanel;        // the region that created this tree includes SNP panel data

    // A "timelist" iterator to the the branch after the cut during rearrangement. Set by ActivateBranch().
    // Used during mid-rearrangement so that CopyPartialBody() & Rectree::Prune() can be done quicker.
    // This member is not copied!
    Branchiter    m_firstInvalid;

    // Protected rearrangement primitives.
    virtual void    Break(Branch_ptr pBranch) = 0;
    void    SetFirstInvalidFrom(Branchconstiter & target);
    void    SetFirstInvalid() { m_firstInvalid = m_timeList.BeginBranch(); };
    void    MarkForDLRecalc(Branch_ptr br);

    // Protected Branch-management primitives.
    const vector<LocusCell> & CollectCells(); // Get locuscells for all loci.
    vector<Branch_ptr>  GetTips(StringVec1d & names)  const;

  public:
    // Creation and destruction.
    virtual           ~Tree();
    virtual Tree     *Clone()                 const = 0;
    virtual Tree     *MakeStump()             const = 0;
    virtual void      Clear();
    virtual void      CopyTips(const Tree * tree);
    virtual void      CopyBody(const Tree * tree);
    virtual void      CopyStick(const Tree * tree);
    virtual void      CopyPartialBody(const Tree * tree);
    virtual TBranch_ptr CreateTip(const TipData & tipdata, const vector<LocusCell> & cells,
                                  const vector<LocusCell> & movingcells, const rangeset & diseasesites);
    virtual TBranch_ptr CreateTip(const TipData & tipdata, const vector<LocusCell> & cells,
                                  const vector<LocusCell> & movingcells, const rangeset & diseasesites,
                                  const vector<Locus> & loci);
    void      SetTreeTimeManager(TimeManager * tm);

    // Getters.
    TimeList & GetTimeList()                 { return m_timeList; };
    const TimeList & GetTimeList()           const { return m_timeList; };
    virtual rangevector GetLocusSubtrees(rangepair span) const;
    Branch_ptr     GetTip(const std::string & name) const;
    TimeManager * GetTimeManager()           { return m_timeManager; };
    const TimeManager * GetTimeManager()     const { return m_timeManager; };
    double      RootTime()                  const { return m_timeList.RootTime(); };
    bool        NoPhaseUnknownSites()       const;
    long int    GetNsites()                 const;
    long int    GetTotalSites()             { return m_totalSites; };

    bool GroomForGrowth(const DoubleVec1d & thetas, const DoubleVec1d & growths, double temperature);
    bool GroomForLogisticSelection(const DoubleVec1d & thetas,
                                   double s, // logistic selection coefficient or zero
                                   double temperature);
    bool GetSnpPanelFlag() { return m_hasSnpPanel; };

    // Setters.
    void SetIndividualsWithTips(const vector<Individual> & indvec);
    void SetLocusVec(vector<Locus> * loc);
    void SetSnpPanelFlag(bool flag) { m_hasSnpPanel = flag; };

    // Likelihood manipulation.
    virtual void      CalculateDataLikes() = 0;
    virtual double    CalculateDataLikesForFixedLoci();
    double    GetDLValue()            const { return m_overallDL; };
    const LongVec1d & GetAliases(long int loc) const { return m_aliases[loc]; };
    void      SetupAliases(const std::vector<Locus> & loci);

    // Rearrangement primitives.
    virtual vector<Branch_ptr> ActivateTips(Tree * othertree);
    virtual Branch_ptr   ActivateBranch(Tree * othertree);
    virtual Branch_ptr   ActivateRoot(FC_Status & fcstatus);
    long CountNodesBetween(double tipwards_time, double rootwards_time) const;

    // ChoosePreferentiallyTowardsRoot does not break the tree, but does set m_firstInvalid.
    // to the interval containing the return value.  So does ChooseFirstBranchInEpoch.
    Branch_ptr   ChoosePreferentiallyTowardsRoot(Tree * othertree);
    Branch_ptr   ChooseFirstBranchInEpoch(double targettime, Tree * othertree);

    virtual void      AttachBase(Branch_ptr newroot);
    virtual vector<Branch_ptr> FindBranchesImmediatelyTipwardOf(Branchiter start);
    vector<Branch_ptr> FindBranchesBetween(double startinterval, double stopinterval);
    virtual vector<Branch_ptr> FirstInterval(double eventT);
    virtual void      NextInterval(Branch_ptr branch );
    virtual void      Prune();
    void      TrimStickToRoot();
    void      SwapSiteDLs();
    virtual void      PickNewSiteDLs();
    virtual void      ReassignDLsFor(std::string lname, long int marker, long int ind);
    void      SetNewTimesFrom(Branchiter start, const DoubleVec1d & newtimes);
    Branch_ptr TransitionEpoch(double eventT, long int newpop, long int maxevents, Branch_ptr pActive);
    vector<Branch_ptr> FindBranchesStartingOnOpenInterval(double starttime, double endtime);
    vector<Branch_ptr> FindEpochBranchesAt(double time);
    vector<Branch_ptr> FindBranchesStartingRootwardOf(double time);

    // Force-specific rearrangement primitives.
    Branch_ptr   Coalesce(Branch_ptr child1, Branch_ptr child2, double tevent, const rangeset & fcsites);

    virtual Branch_ptr   CoalesceActive(double eventT, Branch_ptr active1,
                                        Branch_ptr active2, const rangeset & fcsites);
    virtual Branch_ptr   CoalesceInactive(double eventT, Branch_ptr active,
                                          Branch_ptr inactive, const rangeset & fcsites);
    virtual Branch_ptr   Migrate(double eventT, long int topop, long int maxEvents, Branch_ptr active);
    virtual Branch_ptr   DiseaseMutate(double eventT, long int endstatus, long int maxEvents, Branch_ptr active);

    // TreeSummaryFactory.
    TreeSummary * SummarizeTree() const;

    // Invariant checking.
    bool      IsValidTree()               const;  // check invariants (debugging function)
    bool      ConsistentWithParameters(const ForceParameters& fp) const;  // debugging function
    bool      operator==(const Tree & src) const; // compare trees
    bool      operator!=(const Tree & src) const { return !(*this == src); };

    // TimeManager call throughs (we provide a public front for TimeManager).  CopyStick() is also in this category.
    void DestroyStick();
    void SetStickParams(const ForceParameters & fp);
    bool UsingStick() const;
    void ScoreStick(TreeSummary & treesum) const;
    DoubleVec1d XpartThetasAtT(double time, const ForceParameters & fp) const;
    DoubleVec1d PartitionThetasAtT(double time, force_type force, const ForceParameters & fp) const;

    // TreeSizeArranger.
    virtual void  SetCurTargetLinkweightFrom(const BranchBuffer &) { }; // no-op
    virtual void  ClearCurTargetLinkweight() { };                       // no-op

    // Simulation.
    virtual bool     SimulateDataIfNeeded();
    virtual long int NumberOfRecombinations() = 0;

    // Debugging functions.
    // MakeCoalescent strips all migration nodes from the tree and forces all remaining
    // coalescences to their expectation times.  DO NOT use with migration!
    void MakeCoalescent(double theta) { m_timeList.MakeCoalescent(theta); };
    //
    // DLCheck returns an error message with the first marker at which each branch pair differs;
    // unmentioned branches don't differ.
    void DLCheck(const Tree & other) const;
    //
    // Call through debugging code.
    void PrintStickThetasToFile(std::ofstream & of) const;
    void PrintStickFreqsToFile(std::ofstream & of) const;
    void PrintStickFreqsToFileAtTime(std::ofstream & of, double time) const;
    void PrintStickThetasToFileForJoint300(std::ofstream & of) const;
    void PrintStickToFile(std::ofstream & of) const;
    void PrintDirectionalMutationEventCountsToFile(std::ofstream & of) const;
    void PrintTimeTilFirstEventToFile(std::ofstream & of) const;
    void PrintTraitPhenotypeAtLastCoalescence(std::ofstream & of) const;
};

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

class PlainTree : public Tree
{
  public:
    // Creation and destruction.
    PlainTree() : Tree()  {};
    PlainTree(const Tree & tree, bool makestump) :
        Tree(tree, makestump)       {};
    virtual          ~PlainTree()        {};
    virtual Tree     *Clone()                 const;
    virtual Tree     *MakeStump()             const;
    virtual void      CalculateDataLikes();
    virtual void      Break(Branch_ptr pBranch);
    virtual long int  NumberOfRecombinations() { return 0L; };

    // Yes, everything else is the same as Tree.
};

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

// To keep track of final coalescences:
typedef std::map<long int, rangeset> rangesetcount;
typedef std::map<long int, rangeset>::iterator RSCIter;
typedef std::map<long int, rangeset>::const_iterator RSCcIter;
typedef std::list<std::pair<double, rangesetcount> > sitecountlist;

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

class RecTree : public Tree
{
  private:
    Linkweight m_curTargetLinkweight;
    Linkweight m_newTargetLinkweight;
    std::set<long int> GetIntervalTreeStartSites() const;

  protected:
    vector<Locus> *    m_pMovingLocusVec; // Not owning.  But we can change 'em!
    vector<LocusCell> m_protoMovingCells;

    virtual  void     Break(Branch_ptr branch);
    const vector<LocusCell> & CollectMovingCells();

  public:
    // Creation and destruction.
    RecTree();
    RecTree(const RecTree & tree, bool makestump);
    ~RecTree() {};

    virtual  Tree    *Clone()     const;
    virtual  Tree    *MakeStump() const;
    virtual  void     Clear();
    virtual  void     CopyTips(const Tree * tree);
    virtual  void     CopyBody(const Tree * tree);
    virtual  void     CopyPartialBody(const Tree * tree);

    // Getters.
    virtual  rangevector GetLocusSubtrees(rangepair span) const;
    Linkweight GetCurTargetLinkweight() const { return m_curTargetLinkweight; };
    Linkweight GetNewTargetLinkweight() const { return m_newTargetLinkweight; };
    unsigned long int GetNumMovingLoci() { return m_pMovingLocusVec->size(); };

    // Testers.
    virtual  bool DoesThisLocusJump(long int mloc) const;
    virtual  bool AnyRelativeHaplotypes() const;

    // Setters.
    void      SetMovingLocusVec(vector<Locus> * locs);
    void      SetMovingMapPosition(long int mloc, long int site);

    virtual TBranch_ptr CreateTip(const TipData & tipdata, const vector<LocusCell> & cells,
                                  const vector<LocusCell> & movingcells, const rangeset & diseasesites);

    virtual TBranch_ptr CreateTip(const TipData & tipdata, const vector<LocusCell> & cells,
                                  const vector<LocusCell> & movingcells, const rangeset & diseasesites,
                                  const vector<Locus> & loci);

    // Likelihood manipulation.
    virtual void        CalculateDataLikes();
    virtual double      CalculateDataLikesForMovingLocus(long int mloc);
    virtual DoubleVec1d CalculateDataLikesForFloatingLocus(long int mloc);
    virtual DoubleVec1d CalculateDataLikesWithRandomHaplotypesForFloatingLocus(long int mloc);
    virtual void CalculateDataLikesForAllHaplotypesForFloatingLocus(long int mloc, DoubleVec1d & mlikes);
    virtual bool UpdateDataLikesForIndividualsFrom(long int ind, long int mloc, DoubleVec1d & mlikes);

    // Rearrangement primitives.
    virtual  vector<Branch_ptr> ActivateTips(Tree * othertree);
    virtual  Branch_ptr  ActivateBranch(Tree * othertree);
    virtual  Branch_ptr  ActivateRoot(FC_Status & fcstatus);
    virtual  void     AttachBase(Branch_ptr newroot);
    virtual  vector<Branch_ptr> FirstInterval(double eventT);
    virtual  void     NextInterval(Branch_ptr branch);
    virtual  void     Prune();
    virtual  void     ReassignDLsFor(std::string lname, long int marker, long int ind);

    // Force-specific rearrangement primitives.
    virtual Branch_ptr   CoalesceActive(double eventT, Branch_ptr active1,
                                        Branch_ptr active2, const rangeset & fcsites);
    virtual Branch_ptr   CoalesceInactive(double eventT, Branch_ptr active,
                                          Branch_ptr inactive, const rangeset & fcsites);
    virtual Branch_ptr   Migrate(double eventT, long int topop, long int maxEvents, Branch_ptr active);
    virtual Branch_ptr   DiseaseMutate (double eventT, long int endstatus, long int maxEvents, Branch_ptr active);
    branchpair RecombineActive(double eventT, long int maxEvents, FPartMap fparts,
                               Branch_ptr active, long int recpoint, const rangeset & fcsites, bool lowSitesOnLeft);
    branchpair RecombineInactive(double eventT, long int maxEvents, FPartMap fparts,
                                 Branch_ptr active, long int recpoint, const rangeset & fcsites);

    // Map Summary.
    DoubleVec2d  GetMapSummary();

    // "startsites" is a set of Interval tree start SITEs, each just to the right of a
    // Littlelink recombination recpoint, or it is an endmarker (first or last site).
    std::set<long int> IgnoreDisallowedSubTrees(std::set<long int> startsites, rangeset allowedranges);

    DoubleVec1d ZeroDisallowedSites(DoubleVec1d datalikes, rangeset allowedranges);
    void RandomizeMovingHaplotypes(long int mlocus);

    // Simulation.
    virtual bool SimulateDataIfNeeded();
    virtual long int NumberOfRecombinations();

    // TreeSizeArranger.
    virtual void SetCurTargetLinkweightFrom(const BranchBuffer & brbuffer);
    virtual void ClearCurTargetLinkweight() { m_curTargetLinkweight = ZERO; };

    // Debugging functions.
    virtual void SetNewTargetLinkweightFrom(const BranchBuffer & brbuffer);
    virtual void PrintTipData(long int mloc, long int marker);
    virtual void PrintRangeSetCount(const rangesetcount & rsc);
    virtual rangesetcount RemoveEmpty(const rangesetcount & rsc);
};

#endif // TREE_H

//____________________________________________________________________________________
