#ifndef RBT_H
#define RBT_H

//
#include <iostream>
#include <fstream>
#include <vector>
using namespace std;

#include "Utils.h"
#include "Utils2.h"
#include "BinaryMatrix.h"

// define a leaf-labeled rooted binary tree
// note that we do not store the leaf label explicitly
// rather, we just assume leafs are not the same
// and leafs are simply labeled as 0, 1, 2, 3, 4 etc
// thus, it is enough to ORDER leaves
// IMPORTANT: we assume left node is always smaller (in
// the sense of its minimum left leave is always smaller
// than minimum right leaves

// a class for tree node
class RBTNode
{
public:
    // create a leaf node
    RBTNode(int lvid) : pLeft(NULL), pRight(NULL), pParent(NULL), lvid(lvid), height(0.0) {}
    // create an internal node with two
    RBTNode(RBTNode *pLeft, RBTNode *pRight);
    ~RBTNode() { Clear(); }

    // operation
    void SetHeight(double ht) { height = ht; }
    double GetHeight() const { return height; }
    RBTNode *CopySubTree();
    void AddToLeftEdge(int lvid);
    void AddToRightEdge(int lvid);
    RBTNode *AddSibling(int lvid);
    void AddSiblingToLeaf(int lvid);
    RBTNode *FindLeaf(int lvid, int &ponid); // IMPORTANT, in traversal,
                                             // assume post-order search, and return the how many nodes visited so far
    bool RemoveLeafSelf();                   // only remove self if it is a leaf
    void DetachSubtree();                    // detach this node from the rest of the tree
    RBTNode *AttachSubtree(RBTNode *pSib);

    // access
    RBTNode *GetLeftChild() const { return pLeft; }
    RBTNode *GetRightChild() const { return pRight; }
    RBTNode *GetParent() { return pParent; }
    RBTNode *GetSibling();
    void SetLeftChild(RBTNode *pLeft) { this->pLeft = pLeft; }
    void SetRightChild(RBTNode *pRight) { this->pRight = pRight; }
    void SetParent(RBTNode *pParent) { this->pParent = pParent; }
    int GetLeafId() { return lvid; }
    void SetLeafId(int idNew) { this->lvid = idNew; }
    RBTNode *GetLeftMostChild();
    int GetMinLeaveId();
    void GetLeaves(set<int> &lvs);
    bool IsLeaf() const;
    int GetNumLeavesUnder();
    bool IsLeftChild();
    bool IsRoot() { return this->pParent == NULL; }
    void Dump() const;
    string GetNewick() const;
    void OutputNodeGML(ofstream &ofs);
    void OutputEdgeGML(ofstream &ofs);

    // memory. free recursively
    void Clear();

private:
    void AdjustLRChildUpwards();

    // two children
    RBTNode *pLeft;
    RBTNode *pRight;
    RBTNode *pParent;
    int lvid;
    double height; // useful in some situations, normalized to between 0-1

    // utility
    static int idNodeNextToUse;
};

// define triplets
// Triplets are important for rooted tree, since the set of triplets
// uniquely define a RBT
typedef struct
{
    // note by convention, a < b. But c is on the other side of partition (a,b), c
    int a;
    int b;
    int c;
} TripleLeaves;

// define for traversal
typedef struct
{
    RBTNode *pCurNode;
} TraversRecord;

// sometimes, we want to an ID for the tree
// for the use say HMM states. We just simply
// use integer. Note that we need a way of interpreting this state
// and convert back and force between the real tree and its ID
typedef int RBT_ID;

// main class
class RBT
{
public:
    // different ways of initializing a tree
    // it can be by a supplied id
    RBT(int numLeaves, RBT_ID tid);
    RBT(const RBT &rhs);
    // interop with simple representation
    RBT(int numLeaves, const vector<int> &listNodeLabels, const vector<int> &listParentNodePos,
        const vector<double> &listEdgeDist);
    RBT &operator=(const RBT &rhs);
    //bool operator == (const RBT &rhs) { return IsSame(rhs); }
    ~RBT();

    // ID functions
    RBT_ID GetId();
    RBT_ID MapToId();
    bool IsSame(const RBT &tr) const;

    // splits functions
    bool IsSplitContained(const set<int> &split); // test whether a split is in the tree
    void GetAllSplits(vector<set<int>> &listSplits);

    // SPR function
    void FindSPRDistOneNgbrs(set<int> &ngbrIds);
    void FindSPRDistOneNgbrs(vector<RBT *> &ngbrTrees);
    void FindSPRDistOneNgbrsRestricted(vector<RBT *> &ngbrTrees, const vector<RBT *> &ConstraintTrees);
    bool IsOneSPRAway(const RBT &rbt) const; // testing whether it is one or two SPR away
    bool IsTwoSPRAway(const RBT &rbt) const;
    static void Consolidate(RBT &treeOpt, RBT &treeCmp);

    // editing
    bool RemoveLeaf(int lvid);
    void ReconstructById(RBT_ID tid);
    bool ReconstructNewick(const string &strNewick);
    void PruneLargeIdNodes(int idThres);
    void DeleteLeaves(set<int> &lvids);                                // delete leaves designated
    void RealignLeaves();                                              // sometimes, say after leave is deleted, leaves are no longer contiguous, this op sets it back to contiguous
    void AugamentDupRows(const vector<REMOVED_ROWS_INFO> &rmLvsStage); // restore the leaves removed during matrix preprocessing
    void SetRoot(RBTNode *pRootNew);
    RBTNode *GetRoot() { return pRoot; }

    // dynamic functions: allow adding new nodes
    bool AddLeaf(int pos);

    // access
    void GetLeaves(set<int> &lvs);
    void Dump() const;
    void OutputGML(const char *fileName);

    // Int-op with another format
    void RetrievePlainDesc(int &numLeaves, vector<int> &listNodeLabels, vector<int> &listParentNodePos,
                           vector<double> &listEdgeDist);
    int GetNodesNum() { return 2 * numLeaves - 1; }
    string GetNewick() const;
    int GetLeafNum() { return numLeaves; }
    bool IsEmpty() const { return pRoot == NULL && numLeaves == 0; }

    // compare
    int Compare(RBT &rhs);
    bool IsSameUnrootedTree(RBT &rhs);
    void CollectTips();
    RBTNode *GetTip(int id);
    void GetAllTips(vector<RBTNode *> &tips);

private:
    RBT() {}     // do not allow default construction
    void Init(); // common initialization
    //void ConsTripleMap();   // save all the triples
    // support traversal
    bool InitPostorderTranvers(TraversRecord &tr);
    bool NextPostorderTranvers(TraversRecord &tr);
    void RetrieveSplits();
    RBTNode *FindLeaf(int lvid, int &ponid);
    RBTNode *ReconstructNewickInternal(const string &strNewick);
    bool InternalAddleaf(int lvid, int pos);
    bool ReconstructByPlainDesc(const vector<int> &listNodeLabels, const vector<int> &listParentNodePos,
                                const vector<double> &listEdgeDist);
    void SetLvids(const vector<int> &mapLvids); // configure the name for the leaves

    // save a dynamic root node
    RBTNode *pRoot;

    // we also save the splits
    map<set<int>, bool> mapSplitsInTree;
    map<int, RBTNode *> mapTipPtrs;

    // note we do not normally allow morhping the tree
    // EXCEPT during initialtion. Since convert to id
    // can be slow, we cache it
    int numLeaves;
    RBT_ID tid;

    // collect of triples
    //map< TripleLeaves, bool > mapTriples;
};

///////////////////////////////////////////////////////////////////////////////////////
// useful functions

int GetNumRBT(int nlv);

#endif // RBT_H
