#ifndef MARGINAL_TREE_H
#define MARGINAL_TREE_H


#include <stack>
#include <vector>
#include <map>
#include <set>
using namespace std;

#include "Utils.h"
#include "Utils2.h"
#include "Utils3.h"




///////////////////////////////////////////////////////////////////////////////////////////////
// Define a simple coalescent tree. My experience shows that such a data structure can 
// be quite useful


// yet another structure to represent marginal tree

class MarginalTree
{
public:
    MarginalTree();
    void Clear();
    void Binarize();
    void Consolidate();
    void BuildDescendantInfo();
    void BuildDescendantInfo2();
    void InitDefaultEdgeLen();
    void InitUnitEdgelen();
    double GetDefaultEdgeLen( int child );
    void SetParent( int child, int par, bool fAdjLen=true );
    int GetParent( int child ) const;
    int GetLeftDescendant( int node ) const;
    void SetLeftDescendant(int node, int nodeLeftDesc);
    int GetRightDescendant( int node ) const;
    void SetRightDescendant(int node, int nodeRightDesc);
    double GetEdgeLen( int childNodeIndex ) const;
    double GetTotEdgeLen() const;
    int GetTotNodesNum() const  {return listNodeLabels.size();}
    int GetNumLeaves() const {return numLeaves;}
    void SetNumLeaves(int nl) { numLeaves = nl;}
    void ConsDecedentInfo( vector<vector<int> > &descNodes ) const;
    void ConsAllDecedentInfo( vector<set<int> > &descNodes, bool fIncSelf = true ) const;
    void ConsDecedentLeavesInfo( vector<set<int> > &descNodes ) const;
    void ConsDecedentLeavesInfo2( map<set<int>, int> &mapSTToSTs ) const;
    void ConsDecedentLeavesInfo3( vector<set<int> > &descNodes ) const;
    void ConsDecedentLeavesInfo4( map<set<int>, int> &mapSTToSTs ) const;
    void ConsDecedentLeavesInfoLabels( vector<set<int> > &leafNodeLabels ) const;
    void ConsAllDescendantsInfo( vector<set<int> > &descNodes ) const;
    void ConsHeightsInfo( vector<int> &nodesHt ) const;
    void Dump() const;
    int GetLabel(int r) const 
    {  YW_ASSERT_INFO( r >= 0 && r <(int)listNodeLabels.size(), "wrong3");  return listNodeLabels[r]; }
	void SetLabel(int node, int lbl)
	{
		YW_ASSERT_INFO( node >= 0 && node <(int)listNodeLabels.size(), "wrong4");  listNodeLabels[node] = lbl;
	}
    int GetPosForLabel(int lbl) const;
    void GetlabelsFor(const set<int> &setPos, set<int> &setLbls) const;
    bool IsLeaf( int node ) const { return node >= 0 && node < numLeaves; }
    bool IsToplogicSame( const MarginalTree &tree ) const;
    int GetMRCA(int v1, int v2) const;
    int GetFirstNonselfAnces(int v, const set<int> &setAnces) const;
    void GetChildren(int node, set<int> &listChildren) const;
    int GetMaxHt() const;
    void RemoveLeafNodeFromBinaryTree( int lfn );
    bool AreTwoPathsDisjoint( int sn1, int en1, int sn2, int en2 ) const;
    int GetPath( int sn, int en, set<int> &edgesOnPath ) const;
    double GetPathLen(int sn, int en);
    void CollectNumEdgesToAnces(std::map<int, std::map<int, int> > &mapNodeDistToAnces) const;
    void OutputGML( const char *fileName ) const;
	string GetNewick() const;
    string GetNewickSorted(bool fLen) const;
	string GetNewickNoBrLen() const;
    string GetNewickNoBrLen2() const;
    void GetLeavesUnder( int nn, set<int> &leavesUnder) const;
    void GetLeafSetsForCuts( const vector<int> &listCuts, vector< set<int> > &listLeafSets )  const;
    int GetMRCAForNodes( const set<int> & listNodes ) const;
    bool IsNodeUnder( int nn, int ancesNode) const;
    void RandPermuateLeaves();
	int GetTriple(int a, int b, int c) const;
	int GetSibling(int a) const;
    int GetSiblingFromDescInfo(int a, const vector<vector<int> > &descNodeInfo) const;
    bool AreNodesSibling(int a, int b) const;
	void SetBranchLen(int b, double len)
	{
		YW_ASSERT_INFO(b<(int)listEdgeDist.size(), "Branch wrong");
		listEdgeDist[b] = len;
	}
	void SetLabelList(const vector<int> &listLbls) {listNodeLabels = listLbls;}
    void GetLabelList( vector<int> &listLbls) const { listLbls = listNodeLabels; }
    void GetLabelListForLeaf( vector<int> &listLbls) const;
	void SetParList(const vector<int> &listPars) {listParentNodePos = listPars;}
	void SetBranchLenList(const vector<double> &listLens) {listEdgeDist = listLens;}
	void SortByLeafId();
	void FixDupIds();
	double GetHeight() const;
	int GetRoot() const { return GetTotNodesNum()-1; }
	void SwapBranches(int nodeBranch1, int nodeBranch2);
	void RearrangeParIncOrder();
    void ResetIncLabel();
    void IncLabels();
    void GetTreeEdgeLen(vector<double> &listEdgeDistOut ) const { listEdgeDistOut = this->listEdgeDist; }
    void MapLeafLblConsecutiveOrder( vector<int> &listLeafLblsOld );
    void RemapLeafLabels(const map<int,int> &mapLeafLblsToNew);
    void FindAllSplits( vector<set<int> > &listSplits ) const;
    void FindSibLeafPairs( vector<pair<int,int> > &listSibPairs ) const;
    void MakeLeafSubtreeOfTwo( int posLeaf, int lblChild1, int lblChild2, double len1, double len2 );
    void FindDiffSubtreesFrom( const MarginalTree &mtreeRef, set<int> &setDiffBrs, set<int> &setDiffBrsOrigOnly ) const;
    bool IsOutgroup(int lvid) const;
    void PerformSPR(int nodeSPRSrc, int nodeSPRDest, bool fReOrdering = true);
    void PerformSPRWithDescInfo( int nodeSPRSrc, int nodeSPRDest, const vector<vector<int> > &descNodesInfo );
    void PerformDisjointSPRs(const vector<pair<int,int> > &listSPRs);
    void GetAllNodesBottomUpFrom(int nodeRoot, vector<int> &listNodes) const;
    void RerootSubtree(int nodeST, int nodeSTNewroot);

public:
    int  CalcNormHeight( int node);
	void GetParPosInfo( vector<int> &parPosList ) {parPosList = listParentNodePos;}
	double GetHeightOfNode(int node) const;
    string GetNewickAt(int node, bool fSort = false, bool fLen=true) const;

    // Use an array to store  leaves
    int numLeaves;
    // assume the first numLeaves nodes are leaves
    vector<int> listNodeLabels;
    vector<int> listParentNodePos;
    vector< double > listEdgeDist;
    vector<int> listLeftDescs;
    vector<int> listRightDescs;

private:
    void MapLeafLblConsecutiveOrderAt(int rootST, int &idNext, vector<int> &listLeafLblsOld );
};


///////////////////////////////////////////////////////////////////////////////////////////////
// Global Utilities
class TaxaMapper;


bool ReadinMarginalTrees( ifstream &inFile, vector<MarginalTree> &treeList);
bool ReadinMarginalTreesNewick( ifstream &inFile, int numLeaves, vector<MarginalTree> &treeList, TaxaMapper *pTMapper = NULL, bool fDup=false);
bool ReadinMarginalTreesNewickWLen( ifstream &inFile, int numLeaves, vector<MarginalTree> &treeList, TaxaMapper *pTMapper = NULL);
void AddRootAsLeafToTree( MarginalTree &tree1, bool fIdNonNeg = false );
void GenRandBinaryTree( int numLeaves, MarginalTree &tree1);
void GenRandBinaryTreeClock( int numLeaves, double totHt, MarginalTree &tree1);
// vector<int>: list of leaves in the order from top down, int = top node of chain
void FindChainsInTree( const MarginalTree &tree1,  map< vector<int>, int > &foundChains  );
void InitMarginalTree(MarginalTree &mTree,  int numLeaves, const vector<int> &listLabels, const vector<int> &listParentNodePos );
bool ReadinMarginalTreesNewickWLenString( const string &strNewick, int numLeaves, MarginalTree &treeOut, bool fStartFromZero = true, TaxaMapper *pTMapper = NULL);
void ReadinMarginalTreesNewickWLenString2( const string &strNewick, int numLeaves, MarginalTree &treeOut);
void CollapseEquivTrees(const vector<MarginalTree> &listOrigTrees, vector<MarginalTree> &listUniqTrees, vector<int> &listMultiplicity  );
void FindOneNNIMTreesFrom(MarginalTree &mTreeSrc, vector<MarginalTree> &listNNITrees, vector< pair<int,int> > *pListPairEdgesSwapped=NULL);
void CreateSubtreeFromLeaves( MarginalTree &mTreeOrig, const set<int> &setLeafLabels, MarginalTree &mTreeSub, map<int,int> &mapNewNodeToOldNode );
void UpdateBranchLenInSubtree( MarginalTree &mTreeOrig, map<int,int> &mapNewNodeToOldNode, MarginalTree &mTreeSub );
void RemapLeafIntLabelsTaxaMap( MarginalTree &mtree, map<string,string> &mapper );
void RemapMargTree( MarginalTree &mtree, TaxaMapper &refTMapper );
void FindMatchedSubtrees( MarginalTree &mtreeNew, MarginalTree &mtreeRef, map<int,int> &mapSTNewToRef );

#endif  // MARGINAL_TREE_H

