// ============================================================= //
//                                                               //
//   File      : CT_dtree.cxx                                    //
//   Purpose   :                                                 //
//                                                               //
//   Institute of Microbiology (Technical University Munich)     //
//   http://www.arb-home.de/                                     //
//                                                               //
// ============================================================= //

#include "CT_hash.hxx"
#include "CT_ctree.hxx"

inline void inc_insert_progress(arb_progress *insertProgress) { if (insertProgress) ++(*insertProgress); }

PART *DeconstructedTree::deconstruct_full_subtree(const TreeNode *tree, const GBT_LEN& len, const double& weight, arb_progress *insertProgress) {
    /*! deconstruct TreeNode and register found partitions.
     * @param len length of branch towards subtree 'tree'
     * @param weight summarized for indentical branches (from different trees)
     */

    arb_assert(tree->father);

    PART *ptree = NULp;
    if (tree->is_leaf()) {
        ptree = specSpace.create_tree_PART(tree, weight);
    }
    else {
        if (!(insertProgress && insertProgress->aborted())) {
            PART *pl = deconstruct_full_subtree(tree->get_leftson(), tree->leftlen, weight, insertProgress);
            if (pl) {
                PART *pr = deconstruct_full_subtree(tree->get_rightson(), tree->rightlen, weight, insertProgress);
                if (!pr) {
                    arb_assert(insertProgress && insertProgress->aborted());
                    delete pl;
                }
                else {
                    arb_assert(pl->disjunct_from(pr));

                    ptree = pl->clone();
                    ptree->add_members_from(pr);

                    registry->put_part_from_complete_tree(pl, tree->get_leftson());
                    registry->put_part_from_complete_tree(pr, tree->get_rightson());
                }
            }
        }
    }
    if (ptree) ptree->set_len(len);
    inc_insert_progress(insertProgress);

    return ptree;
}

void DeconstructedTree::deconstruct_full_rootnode(const TreeNode *tree, const double& weight, arb_progress *insertProgress) {
    /*! deconstruct TreeNode and register found partitions.
     * @param weight summarized for indentical branches (from different trees)
     */

    arb_assert(!tree->father);
    arb_assert(!tree->is_leaf());

    double root_length = (tree->leftlen + tree->rightlen);

    PART *pl = deconstruct_full_subtree(tree->get_leftson(), root_length, weight, insertProgress);
    if (!pl) {
        arb_assert(insertProgress && insertProgress->aborted());
    }
    else {
        PART *pr = deconstruct_full_subtree(tree->get_rightson(), root_length, weight, insertProgress);
        if (!pr) {
            arb_assert(insertProgress && insertProgress->aborted());
            delete pl;
        }
        else {
            arb_assert(pl->disjunct_from(pr));

            // add only one of the partition pl and pr (they both represent the root-edge)
            registry->put_part_from_complete_tree(pl, tree->get_leftson());
            delete pr;

            inc_insert_progress(insertProgress);
        }
    }
}

PART *DeconstructedTree::deconstruct_partial_subtree(const TreeNode *tree, const GBT_LEN& len, const double& weight, const PART *partialTree, arb_progress *insertProgress) {
    /*! deconstruct partial TreeNode
     *
     * similar to deconstruct_full_subtree(),
     * but the set of missing species is added at each branch.
     */

    arb_assert(tree->father);

    PART *ptree = NULp;
    if (tree->is_leaf()) {
        ptree = specSpace.create_tree_PART(tree, weight);
    }
    else {
        if (!(insertProgress && insertProgress->aborted())) {
            PART *pl = deconstruct_partial_subtree(tree->get_leftson(), tree->leftlen, weight, partialTree, insertProgress);
            if (pl) {
                PART *pr = deconstruct_partial_subtree(tree->get_rightson(), tree->rightlen, weight, partialTree, insertProgress);
                if (!pr) {
                    arb_assert(insertProgress && insertProgress->aborted());
                    delete pl;
                }
                else {
                    arb_assert(pl->disjunct_from(pr));

                    ptree = pl->clone();
                    ptree->add_members_from(pr);

                    registry->put_part_from_partial_tree(pl, partialTree, tree->get_leftson());
                    registry->put_part_from_partial_tree(pr, partialTree, tree->get_rightson());
                }
            }
        }
    }
    if (ptree) ptree->set_len(len);
    inc_insert_progress(insertProgress);

    return ptree;
}

void DeconstructedTree::deconstruct_partial_rootnode(const TreeNode *tree, const double& weight, const PART *partialTree, arb_progress *insertProgress) {
    /*! deconstruct partial TreeNode
     *
     * similar to deconstruct_full_rootnode(),
     * but the set of missing species is added at each branch.
     */

    arb_assert(!tree->father);
    arb_assert(!tree->is_leaf());

    double root_length = (tree->leftlen + tree->rightlen);

    PART *pl = deconstruct_partial_subtree(tree->get_leftson(), root_length, weight, partialTree, insertProgress);
    if (!pl) {
        arb_assert(insertProgress && insertProgress->aborted());
    }
    else {
        PART *pr = deconstruct_partial_subtree(tree->get_rightson(), root_length, weight, partialTree, insertProgress);
        if (!pr) {
            arb_assert(insertProgress && insertProgress->aborted());
            delete pl;
        }
        else {
            arb_assert(pl->disjunct_from(pr));

            pr->add_members_from(pl); // in 'pr' we collect the whole partialTree
            registry->put_part_from_partial_tree(pl, partialTree, tree->get_leftson()); // because we are at root edge, we only insert one partition ('pl')

            arb_assert(pr->equals(partialTree));

            arb_assert(is_similar(pr->get_weight(), weight, 0.01));
            registry->put_artificial_part(pr);
            inc_insert_progress(insertProgress);
        }
    }
}

void SpeciesSpace::add_tree_to_PART(const TreeNode *tree, PART& part) const {
    if (tree->is_leaf()) {
        part.setbit(get_species_index(tree->name));
    }
    else {
        add_tree_to_PART(tree->get_leftson(), part);
        add_tree_to_PART(tree->get_rightson(), part);
    }
}

PART *SpeciesSpace::create_tree_PART(const TreeNode *tree, const double& weight) const {
    PART *part = new PART(get_PartitionSize(), weight);
    if (part) add_tree_to_PART(tree, *part);
    return part;
}