//
//  TreeBuilder.h
//  
//
//  Created by Yufeng Wu on 9/11/15.
//
//

#ifndef ____TreeBuilder__
#define ____TreeBuilder__

#include <iostream>
#include <map>
#include <set>
#include <vector>
#include <string>
using namespace std;

//***********************************************************************
void TestNJ();

//***********************************************************************
// implement various methods to build a phylogenetic tree

// define distances between taxa
class PhyloDistance
{
public:
    void SetDistance(int node1, int node2, double dist);
    double GetDistance(int node1, int node2) const;
    void GetAllNodes( set<int> &nodesAll) const;
    double GetDistanceNonNeg( int node1, int node2 ) const;
    double CalcAveDistBtwClusters( const set< set<int> > &setClusters ) const;
    void PrepareVecDists(int maxTaxaSize);     // store distance in vectors s.t. access faster
    void Dump() const;
    
private:
    map< pair<int,int>, double> mapDists;
    vector<vector<double> > vecDists;
};

// distance based tree builder
class DistanceTreeBuilder
{
public:
    DistanceTreeBuilder(PhyloDistance &distPairwiseTaxaIn);
    string NJ();
    string ConstrainedNJ( const set<set<int> > &setClustersMustHave, int totElem );  // enforce NJ that must have the passed-in cluster
    string ConstrainedUPGMA( const set<set<int> > &setClustersMustHave, const set<set<int> > &setClustersForbidden, map<set<int>, double> &mapSTHts, int numTotElem = -1 );
    string ConstrainedUPGMA( const set<set<int> > &setClustersMustHave, const set<set<int> > &setClustersDesired, int numTopCandidates, const set<set<int> > &setClustersForbidden, map<set<int>, double> &mapSTHts, int numTotElem );
    void SetTaxonName(int id, const string &tname) { mapIndexToName[id] = tname; }
    void SetOutgroup(int og) { taxonOutgroup = og; }
    void SetNumThreads(int t) { numThreads = t; }
    
    double NJCalcAveDist(int nodecur, const set<int> &nodesToSearch);
    bool IsClusterIncompatibleWithSetofClus(const set<int> &clus1, const set<set<int> > &setClus, int numTotElem = -1) const;
    
private:
    void NJFindNgbrs( int nodeIdNew, set<int> &nodesToSearch, int &ngbr1, int &ngbr2, const set<set<int> > *psetClustersMustHave = NULL, map<int,set<int> > *pmapLinClades = NULL, int *ptotElem = NULL  );
    void NJFindNgbrsMT( int nodeIdNew, set<int> &nodesToSearch, int &ngbr1, int &ngbr2, const set<set<int> > *psetClustersMustHave = NULL, map<int,set<int> > *pmapLinClades = NULL, int *ptotElem = NULL );
    void NJFindNgbrsMT2( int nodeIdNew, set<int> &nodesToSearch, int &ngbr1, int &ngbr2, const set<set<int> > *psetClustersMustHave = NULL, map<int,set<int> > *pmapLinClades = NULL, int *ptotElem = NULL );
    bool IsClusterIncompatible(const set<int> &clus1, const set<int> &clus2, int numTotElem = -1) const;
    void UpdateDistUPGMA( const pair<set<int>, set<int> > &pairClus, const map<set<int>, pair<string,double> > &mapSubtree, map< pair<set<int>, set<int> >, double> &distMapCur );
    string GetTaxonNameFor(int index) const;
    int GetNumCompatCladesIn( const set<int> &clus1, const set<set<int> > &setCladesTest, int numTotElem ) const;
    
    PhyloDistance &distPairwiseTaxa;
    map<int, string> mapIndexToName;
    int taxonOutgroup;
    int numThreads;
};

//***********************************************************************
// tool for building UPGMA tree

class ConstrainedUPGMATreeBuilder
{
public:
    ConstrainedUPGMATreeBuilder( PhyloDistance &distPairwiseTaxaIn, const set<set<int> > &setClustersMustHave, const set<set<int> > &setClustersForbidden, int numTotElemIn=-1 );
    ConstrainedUPGMATreeBuilder(const ConstrainedUPGMATreeBuilder &rhs);
    string GetTree() const;
    string GetPartialConsTree() const;
    double GetMinCoalSubtrees( set<int> &st1, set<int> &st2 ) const;
    void GetCoalSubtreesHtBound(double htBound, set< pair<pair<set<int>, set<int> >, double> > &setCandidates ) const;
    void MergeSubtrees(const set<int> &st1, const set<int> &st2, double htMergedST);
    void GetMergeCandidates(  map< pair<set<int>,set<int> >, double > &setCandidates ) const;
    double GetCurDistForTwoClusters( const set<int> &clus1, const set<int> &clus2 ) const;
    void SetDistForTwoClusters( const set<int> &clus1, const set<int> &clus2, double dist );
    void AddLeafCluster(int leafId, double ht);
    int GetNumSubtrees() const;
    void GetAllSubtrees( map<set<int>, string > &mapSTs  ) const;
    void GetActiveSubtrees(set< set<int> > &setActiveSTs  ) const;
    bool IsDone() const;
    void Dump() const;
    
private:
    void Init();
    bool IsClusterIncompatible(const set<int> &clus1, const set<int> &clus2) const;
    bool IsClusterIncompatibleWithSetofClus(const set<int> &clus1, const set<set<int> > &setClus) const;
    void UpdateDistUPGMA( const set<int> &st1, const set<int> &st2 );
    
    PhyloDistance &distPairwiseTaxa;
    const set<set<int> > &setClustersMustHave;
    const set<set<int> > &setClustersForbidden;
    int numTotElem;
    map< pair<set<int>, set<int> >, double> distMapActivePair;
    map<set<int>, pair<string, double> > mapClusSubtree;
    vector< pair< set<int>,set<int> > > histSTMerge;
};

//***********************************************************************
// tool for building near-optimal UPGMA tree

class ConstrainedNearUPGMATreesBuilder
{
public:
    ConstrainedNearUPGMATreesBuilder( PhyloDistance &distPairwiseTaxaIn, const set<set<int> > &setClustersMustHave, const set<set<int> > &setClustersForbidden, int numTotElem );
    void Construct(int maxNumTrees, double thresMaxDistRatio);
    void GetTrees(set<string> &setConsTrees) const { setConsTrees = setTreeCons; }
    
private:
    PhyloDistance &distPairwiseTaxa;
    const set<set<int> > &setClustersMustHave;
    const set<set<int> > &setClustersForbidden;
    int numTotElem;
    set<string> setTreeCons;
};


#endif /* defined(____TreeBuilder__) */
