from treeswift import *
from sys import stdin, stdout, argv, exit, stderr
import argparse
from os import path
import re
import io
import argparse
import arviz as az
import pandas as pd
import plotly.graph_objects as go
import logging
from sys import stdout
import numpy
import cvxopt
from numpy import *
import logging
from sys import stdout


FastRoot_usage = '''
================================= FastRoot example commands =================================

TreeSAK FastRoot -m1 r1_mcmc.txt -m2 r1_mcmc.txt -o convergence_plot.png

# This script was modified based on the script FastRoot.py from Uyen Mai: 
https://github.com/uym2/MinVar-Rooting/blob/master/FastRoot.py

==============================================================================================
'''


def new_logger(myName,myLevel=logging.INFO,myStream=stdout):
    logger = logging.getLogger(myName)
    logger.setLevel(myLevel)
    handler = logging.StreamHandler(myStream)


class Tree_extend(object):
    def __init__(self, ddpTree=None, tree_file=None, schema="newick"):  # ,logger_id=1,logger_stream=sys.stderr):
        # self.logger = new_logger(__name__+ "_" + str(logger_id),myStream=logger_stream)
        if tree_file:
            self.ddpTree = read_tree(tree_file, schema)
        else:
            self.ddpTree = ddpTree

    def Bottomup_label(self):
        # assign each node a label so that we can later relate to it
        i = 0
        for node in self.ddpTree.traverse_postorder():
            if node.is_leaf():
                node.name = 'L' + str(i)
            else:
                node.name = 'I' + str(i)
            i += 1

    def Topdown_label(self, label_type="all"):
        # assign each node a label so that we can later relate to it
        i = 0

        for node in self.ddpTree.traverse_preorder():
            if node.is_leaf():
                if label_type == "all" or label_type == "leaves":
                    node.name = 'L' + str(i)
                else:
                    node.name = node.label
            else:
                if label_type == "all" or label_type == "internal":
                    node.name = 'I' + str(i)
                else:
                    node.name = node.label
            i += 1

    def Bottomup_update(self):
        for node in self.ddpTree.traverse_postorder():
            self.Node_init(node)
            self.bUp_update(node)

    def Topdown_update(self):
        for node in self.ddpTree.traverse_preorder():
            self.tDown_update(node, self.Opt_function)

    def compute_distances(self):
        D = {}

        def __compute_dRoot__(node, cumm_l):
            if node.is_leaf():
                D[node.name] = cumm_l
            else:
                for child in node.child_nodes():
                    __compute_dRoot__(child, cumm_l + child.edge_length)

        __compute_dRoot__(self.ddpTree.root, 0)
        return D

    def compute_ingroup_distances(self):
        D = []

        def __compute_dLeaf__(node, cumm_l):
            if node.is_leaf():
                D.append(cumm_l)
            else:
                for child in node.child_nodes():
                    __compute_dLeaf__(child, cumm_l + child.edge_length)

        children = self.ddpTree.root.child_nodes()
        crowded_child = None
        maxleaf = -1

        for node in children:
            if node.nleaf > maxleaf:
                maxleaf = node.nleaf
                crowded_child = node

        __compute_dLeaf__(children[1], 0)

        return D

    def filter_branch(self, threshold=None):
        # filter out abnormally long branches
        i = 1
        self.logger.info("Iteration: " + str(i))
        self.Reroot()
        while 1:
            check = self.filter_by_threshold(threshold=threshold)
            if (not check):
                self.logger.info("I could not remove anything more! I stop here!")
                break
            i += 1
            self.logger.info("Iteration: " + str(i))
            self.reset()
            self.Reroot()

    def filter_by_threshold(self, threshold=None, k=3.5):
        if threshold is None:
            threshold = self.compute_threshold(k=k)

        def __filter__(node, cumm_l):
            removed = False
            node.child_removed = False
            for child in node.child_nodes():
                check = __filter__(child, cumm_l + child.edge_length)
                removed = removed or check

            p = node.parent_node
            # if ( cumm_l > threshold ) or ( node.child_removed and len(node.child_nodes()) == 0 ):
            if (cumm_l > threshold) or (node.child_removed and node.num_children() == 0):
                # remove node
                p.remove_child(node)
                # update parent node
                p.child_removed = True
                removed = True
                try:
                    self.logger.info(node.label + " removed")
                except:
                    self.logger.info(node.name + " removed")
            # elif len(node.child_nodes()) == 1:
            elif node.num_child_nodes() == 1:
                # remove node and attach its only child to its parent
                e1 = node.edge_length
                child = node.child_nodes()[0]
                e2 = child.edge_length
                p.remove_child(node)
                node.remove_child(child)
                p.add_child(child)
                child.edge_length = e1 + e2
            return removed

        return __filter__(self.get_root(), 0)

    def compute_threhold(self, k=3.5):
        self.logger.warning("Abstract class! Should never be called")
        return 0

    def reset(self):
        self.logger.warning("Abstract class! Should never be called")

    def find_root(self):
        self.Topdown_label()  # temporarily included for debugging
        self.Bottomup_update()
        self.prepare_root()
        self.Topdown_update()

    def opt_score(self):
        self.logger.warning("Abstract class! Should never be called")

    def report_score(self):
        self.logger.warning("Abstract class! Should never be called")

    def Reroot(self):
        self.find_root()
        # self.report_score()
        # d2currRoot = 0
        # br2currRoot = 0
        if self.opt_root != self.ddpTree.root:
            # d2currRoot,br2currRoot = self.reroot_at_edge(self.opt_root.edge, self.opt_root.edge_length-self.opt_x, self.opt_x)
            self.reroot_at_edge(self.opt_root, self.opt_x)
            # self.ddpTree.reroot(self.opt_root,self.opt_x)

        # return head_id, tail_id, edge_length, self.opt_x
        # return d2currRoot,br2currRoot

    def Opt_function(self, node):
        self.logger.warning("Abstract method! Should never be called")

    def tree_as_newick(self, outstream=sys.stdout, label_by_name=False):
        # dendropy's method to write newick seems to have problem ...
        self.__write_newick(self.ddpTree.root, outstream, label_by_name=label_by_name)
        outstream.write(";\n")

    #            outstream.write(bytes(";\n", "ascii"))

    def __write_newick(self, node, outstream, label_by_name=False):
        if node.is_leaf():
            if label_by_name:
                outstream.write(str(node.name))
            #                    outstream.write(bytes(str(node.name), "ascii"))
            else:
                try:
                    outstream.write(node.label)
                #                        outstream.write(bytes(node.label, "ascii"))
                except:
                    outstream.write(node.label)
        #                        outstream.write(bytes(str(node.label), "ascii"))
        else:
            outstream.write('(')
            # outstream.write(bytes('(', "ascii"))
            is_first_child = True
            for child in node.child_nodes():
                if is_first_child:
                    is_first_child = False
                else:
                    outstream.write(',')
                #                        outstream.write(bytes(',', "ascii"))
                self.__write_newick(child, outstream, label_by_name=label_by_name)
            outstream.write(')')
        #                outstream.write(bytes(')', "ascii"))
        if not node.is_leaf():
            if label_by_name:
                outstream.write(str(node.name))
            #                    outstream.write(bytes(str(node.name), "ascii"))
            elif node.label is not None:
                outstream.write(str(node.label))
        #                    outstream.write(bytes(str(node.label), "ascii"))

        if not node.edge_length is None:
            outstream.write(":" + str(node.edge_length))

    #                outstream.write(bytes(":" + str(node.edge_length), "ascii"))

    def reroot_at_edge(self, node, length):
        # the method provided by dendropy DOESN'T seem to work ...
        # change edge to opt_root
        length1 = node.edge_length - length
        length2 = length
        if not node:
            return
        head = node  # opt_root = v = node
        tail = node.parent  # u parent of opt_root
        if not tail:
            return

        if (length2 == 0) and head.is_leaf():
            return 0, 0

        # new_root = self.ddpTree.node_factory()
        new_root = Node()

        tail.remove_child(head)

        new_root.add_child(head)
        head.edge_length = length2

        p = tail.parent
        l = tail.edge_length

        new_root.add_child(tail)
        tail.edge_length = length1

        br2currRoot = 0
        d2currRoot = length1

        #            if tail.label == self.ddpTree.root.label:
        if (tail is self.ddpTree.root):
            head = new_root

        while tail is not self.ddpTree.root:
            # MAD@ add
            # q = tail.parent #tail should have 2 parents right now: new_root and its old parent
            q = head.parent
            # End MAD@ add
            head = tail
            tail = p
            p = tail.parent

            br2currRoot += 1
            d2currRoot += l

            l1 = tail.edge_length
            tail.remove_child(head)
            # MAD@ add
            head.parent = q
            # End MAD@ add

            head.add_child(tail)
            tail.edge_length = l
            l = l1

        # out of while loop: tail IS now tree.root
        if tail.num_children() == 1:
            # merge the 2 branches of the old root and adjust the branch length
            # sis = [child for child in tail.child_nodes()][0]
            sis = tail.child_nodes()[0]
            l = sis.edge_length
            tail.remove_child(sis)
            head.add_child(sis)
            sis.edge_length = l + tail.edge_length
            head.remove_child(tail)
            # tail.remove_child(head)

        new_root.name = self.ddpTree.root.name
        self.ddpTree.root.name = "OLD"
        self.ddpTree.root = new_root

        ### MAD@ add
        #            for node in self.ddpTree.traverse_postorder():
        #                for child in node.child_nodes():
        #                    if child.parent_node is not node:
        #                        logger.info("Error found!")
        #                        child.parent_node = node
        ### MAD@ add

        return d2currRoot, br2currRoot

    def get_root(self):
        return self.ddpTree.root


def cvxopt_solve_qp(P, q, G=None, h=None, A=None, b=None, maxIter=1000):
    P = .5 * (P + P.T)  # make sure P is symmetric
    args = [cvxopt.matrix(P), cvxopt.matrix(q)]
    if G is not None:
        args.extend([cvxopt.matrix(G), cvxopt.matrix(h)])
        if A is not None:
            args.extend([cvxopt.matrix(A), cvxopt.matrix(b)])
    sol = cvxopt.solvers.qp(*args,options={'show_progress':False,'maxiters':maxIter})
    if 'optimal' not in sol['status']:
        if "unknown" in sol['status']:
            logger.warning("Couldn't find optimal solution on one branch. Perhaps due to maximum iterations exceeded. Consider increasing the maximum iterations via -x.")
        else:
            logger.warning("Couldn't find optimal solution on one branch. Solution status: " + sol['status'])
    #return None
    return numpy.array(sol['x']).reshape((P.shape[1],))

EPSILON = 1e-5

class RTT_Tree(Tree_extend):
    # supportive base class to implement RTT-reroot, hence the name
    def __init__(self, smplTimes, ddpTree=None, tree_file=None, schema="newick", logger_id=1, logger_stream=stderr,
                 maxIter=1000):
        super(RTT_Tree, self).__init__(ddpTree, tree_file, schema)
        self.logger = new_logger("RTT_Tree_" + str(logger_id), myStream=logger_stream)
        self.smplTimes = smplTimes
        self.reset()
        self.maxIter = maxIter

    def reset(self):
        self.RTT = None
        self.opt_root = self.ddpTree.root
        self.opt_y = 0
        self.opt_x = 0
        self.opt_mu = 0
        self.tmin = min(self.smplTimes.values())

    def Node_init(self, node, nleaf=1, SDI=0, SD=0, ST=0, SDT=0, SSD=0):
        node.SDI = SDI
        node.SD = SD
        node.nleaf = nleaf
        node.ST = ST
        node.SDT = SDT
        node.SSD = SSD

    def Opt_function(self, node, SST, deltaT, deltaD, SDT, SSD, ST, SD):
        n = self.total_leaves
        a, b, c, d, e, f = n, SST, (-2 * deltaT), (2 * deltaD), (-2 * SDT), SSD
        k, m, r = 2 * (n - 2 * node.nleaf), -2 * ST, 2 * SD

        tmin = self.tmin

        # use quadprog to compute mu_star, y_star, and x_star
        P = array([[a, k / 2, c / 2.], [k / 2, n, m / 2], [c / 2, m / 2, b]])
        q = array([d / 2., r / 2, e / 2])
        G = array([[-1., 0., 0.], [0., 0., -1.], [1., 0., 0.], [0., 1., -tmin]])
        h = array([0., EPSILON, node.edge_length, 0]).reshape((4,))
        solution = cvxopt_solve_qp(P, q, G, h, maxIter=self.maxIter)
        x_star = solution[0]
        y_star = solution[1]
        mu_star = solution[2]
        curr_RTT = a * x_star * x_star + b * mu_star * mu_star + c * x_star * mu_star + d * x_star + e * mu_star + f + n * y_star * y_star + k * x_star * y_star + m * mu_star * y_star + r * y_star

        if self.RTT is None or (curr_RTT - self.RTT < -EPSILON):
            self.RTT = curr_RTT
            self.opt_root = node
            self.opt_x = node.edge_length - x_star
            self.opt_y = y_star
            self.opt_mu = mu_star

    def bUp_update(self, node):
        if node.is_leaf():
            node.nleaf = 1
            node.SDI = 0
            node.ST = self.smplTimes[node.label]
        else:
            node.nleaf = 0
            node.SDI = 0
            node.ST = 0
            for child in node.child_nodes():
                node.nleaf += child.nleaf
                node.SDI += child.SDI + child.nleaf * child.edge_length
                node.ST += child.ST

    def Update_var(self, child, node, edge_length):
        SST = self.SST
        deltaT = self.ddpTree.root.ST - 2 * child.ST
        deltaD = -2 * child.nleaf * edge_length - 2 * child.SDI + node.SD
        SDT = node.SDT
        SSD = node.SSD
        ST = self.ST
        SD = node.SD
        return SST, deltaT, deltaD, SDT, SSD, ST, SD

    def tDown_update(self, node, opt_function):
        for child in node.child_nodes():
            child.SD = node.SD + (self.total_leaves - 2 * child.nleaf) * child.edge_length
            child.SDT = node.SDT + child.edge_length * (self.ddpTree.root.ST - 2 * child.ST)
            child.SSD = node.SSD + (self.total_leaves - 4 * child.nleaf) * (child.edge_length ** 2) + 2 * (
                        node.SD - 2 * child.SDI) * child.edge_length
            SST, deltaT, deltaD, SDT, SSD, ST, SD = self.Update_var(child, node, child.edge_length)
            opt_function(child, SST, deltaT, deltaD, SDT, SSD, ST, SD)

    def prepare_root(self):
        root = self.get_root()
        root.SD = root.SDI
        self.total_leaves = root.nleaf
        self.ST = root.ST
        self.ddpTree.root.droot = 0
        self.ddpTree.root.troot = 0
        root.SD, root.SSD, root.SDT, self.SST = 0, 0, 0, 0
        for v in self.ddpTree.traverse_preorder():
            if not v.is_root():
                # must have defined edge lengths
                v.droot = v.parent.droot + v.edge_length
                if v.is_leaf():
                    root.SSD += (v.droot ** 2)
                    self.SST += (self.smplTimes[v.label] ** 2)
                    root.SD += v.droot
                    root.SDT += (v.droot * self.smplTimes[v.label])

    def opt_score(self):
        return self.RTT

    def return_values(self):
        # returns RTT score, mu, t0
        return (self.opt_score() / self.total_leaves), (self.opt_mu), (self.opt_y / self.opt_mu)

    def report_score(self):
        return "RTT=" + str(self.opt_score() / self.total_leaves) + "\tmu=" + str(self.opt_mu) + "\tt0=" + str(
            self.opt_y / self.opt_mu)


class OGR_Tree(Tree_extend):
    # supportive class to implement outgroup-reroot (OGR = outgroup reroot, hence the name)
    # this rooting method solve the difficulty in finding the root when there are mulitple outgroups
    # and they are not monophyletic. It seeks for the rooting place that maximizes the triplet score
    # of the specified outgroups.
    def __init__(self, outgroups, ddpTree=None, tree_file=None, schema="newick", logger_id=1, logger_stream=sys.stderr):
        super(OGR_Tree, self).__init__(ddpTree, tree_file, schema)
        self.logger = new_logger("OGR_Tree_" + str(logger_id), myStream=logger_stream)
        # L = self.ddpTree.leaf_nodes()
        L = []
        for leaf in self.ddpTree.traverse_leaves():
            L.append(leaf)
        self.OGs = set([x.label for x in L if x.label in set(outgroups)])
        self.nOGs = len(self.OGs)
        self.nIGs = len(L) - self.nOGs
        self.max_nTrpls = self.nIGs * self.nOGs * (self.nOGs - 1) / 2 + self.nOGs * self.nIGs * (self.nIGs - 1) / 2
        self.reset()

    def reset(self):
        self.opt_root = self.ddpTree.root
        self.opt_nTrpls = 0

    def Node_init(self, node, nTrpl_in=0, nTrpl_out=0, nOGs=0, nIGs=0):
        node.nTrpl_in = nTrpl_in
        node.nTrpl_out = nTrpl_out
        node.nOGs = nOGs
        node.nIGs = nIGs

    def Opt_function(self, node):
        curr_nTrpls = node.nTrpl_in + node.nTrpl_out
        if curr_nTrpls > self.opt_nTrpls:
            self.opt_nTrpls = curr_nTrpls
            self.opt_root = node
            self.opt_x = node.edge_length / 2  # NOTE: this method does not consider branch length, the *middle point* of the edge is just arbitrarily chosen

    def bUp_update(self, node):
        if node.is_leaf():
            node.nOGs = 1 if node.label in self.OGs else 0
            node.nIGs = 1 if node.nOGs == 0 else 0
        else:
            C = node.child_nodes()

            node.nOGs = sum([c.nOGs for c in C])
            node.nIGs = sum([c.nIGs for c in C])

            node.nTrpl_in = sum([c.nTrpl_in for c in C])

            for i, c1 in enumerate(C):
                for c2 in C[i + 1:]:
                    IG_trpls = c1.nIGs * c2.nIGs * (self.nOGs - node.nOGs)
                    OG_trpls = c1.nOGs * c2.nOGs * (self.nIGs - node.nIGs)
                    node.nTrpl_in += IG_trpls + OG_trpls

    def tDown_update(self, node, opt_function):
        C = node.child_nodes()

        for child in C:
            C1 = [c for c in C if c is not child]
            child.nTrpl_out = node.nTrpl_out

            for i, c1 in enumerate(C1):
                child.nTrpl_out += c1.nTrpl_in
                child.nTrpl_out += (self.nIGs - node.nIGs) * c1.nIGs * child.nOGs
                child.nTrpl_out += (self.nOGs - node.nOGs) * c1.nOGs * child.nIGs

                for c2 in C1[i + 1:]:
                    IG_trpls = c1.nIGs * c2.nIGs * child.nOGs
                    OG_trpls = c1.nOGs * c2.nOGs * child.nIGs

                    child.nTrpl_out += IG_trpls + OG_trpls

            opt_function(child)

    def prepare_root(self):
        pass

    def opt_score(self):
        return self.opt_nTrpls / float(self.max_nTrpls) if self.max_nTrpls != 0 else None

    def report_score(self):
        myScore = self.opt_score()
        if myScore is None:
            self.logger.warning("OG rooting failed because the tree has no outgroup")
        return "Triplet score: " + str(self.opt_score())


class MPR_Tree(Tree_extend):
    # supportive class to implement midpoint-reroot (mpr = mid point reroot, hence the name)G
    def __init__(self, ddpTree=None, tree_file=None, schema="newick", logger_id=1, logger_stream=sys.stderr):
        super(MPR_Tree, self).__init__(ddpTree, tree_file, schema)
        self.logger = new_logger("MPR_Tree_" + str(logger_id), myStream=logger_stream)
        self.reset()

    def reset(self):
        self.max_distance = -1
        self.opt_root = self.ddpTree.root
        self.opt_x = 0

    def Node_init(self, node, max_in=None, max_out=-1):
        node.max_in = max_in if max_in else [0, 0]
        node.max_out = max_out

    def Opt_function(self, node):
        m = max(node.max_in)
        curr_max_distance = m + node.max_out
        x = (node.max_out - m) / 2
        if curr_max_distance > self.max_distance and x >= 0 and x <= node.edge_length:
            self.max_distance = curr_max_distance
            self.opt_x = x
            self.opt_root = node

    def bUp_update(self, node):
        if not node.is_leaf():
            node.max_in = []
            for child in node.child_nodes():
                node.max_in.append(max(child.max_in) + child.edge_length)

    def tDown_update(self, node, opt_function):
        child_idx = 0
        for child in node.child_nodes():
            child.max_out = max([node.max_out] + [node.max_in[k] for k in range(len(node.max_in))
                                                  if k != child_idx]) + child.edge_length
            opt_function(child)
            child_idx += 1

    def prepare_root(self):
        pass

    def compute_threhold(self, k=3.5):
        self.logger.warning("Trying to compute threshold for MPR_Tree, which is not supported.")
        return 0

    def opt_score(self):
        return self.max_distance / 2

    def report_score(self):
        return "Tree height: " + str(self.opt_score())


class minVAR_Base_Tree(Tree_extend):
    # supportive base class to implement VAR-reroot, hence the name
    def __init__(self, ddpTree=None, tree_file=None, schema="newick",logger_id=1,logger_stream=sys.stderr):
        super(minVAR_Base_Tree, self).__init__(ddpTree, tree_file, schema)
        self.logger = new_logger("MinVar_Tree_" + str(logger_id),myStream=logger_stream)
        self.reset()

    def reset(self):
        self.minVAR = None
        self.opt_root = self.ddpTree.root
        self.opt_x = 0

    def Node_init(self, node, nleaf=1, sum_in=0, sum_total=0, var=-1):
        node.sum_in = sum_in
        node.sum_total = sum_total
        node.nleaf = nleaf
        node.var = var

    def Opt_function(self, node, a, b, c):
        self.logger.info("Abstract method! Should never be called")

    def compute_dRoot_VAR(self):
        cumm = {'ssq': 0, 'sum': 0}

        def compute_dRoot(node, cumm_l):
            if node.is_leaf():
                cumm['ssq'] += cumm_l ** 2
                cumm['sum'] += cumm_l
            else:
                for child in node.child_nodes():
                    compute_dRoot(child, cumm_l + child.edge_length)

        compute_dRoot(self.get_root(), 0)
        N = self.get_root().nleaf
        root_var = cumm['ssq'] / N - (cumm['sum'] / N) ** 2
        self.get_root().var = root_var

    def bUp_update(self, node):
        if node.is_leaf():
            node.nleaf = 1
            node.sum_in = 0
        else:
            node.nleaf = 0
            node.sum_in = 0
            for child in node.child_nodes():
                node.nleaf += child.nleaf
                node.sum_in += child.sum_in + child.nleaf * child.edge_length

    def Update_var(self, child, node, edge_length):
        alpha = 2 * (node.sum_total - 2 * (child.sum_in + child.nleaf * edge_length)) / self.total_leaves
        beta = 1 - 2 * float(child.nleaf) / self.total_leaves
        a = 1 - beta * beta
        b = alpha - 2 * node.sum_total * beta / self.total_leaves
        c = node.var
        child.var = a * edge_length * edge_length + b * edge_length + c
        return a, b, c

    def tDown_update(self, node, opt_function):
        for child in node.child_nodes():
            child.sum_total = node.sum_total + (self.total_leaves - 2 * child.nleaf) * child.edge_length
            a, b, c = self.Update_var(child, node, child.edge_length)
            opt_function(child, a, b, c)

    def prepare_root(self):
        root = self.get_root()
        root.sum_total = root.sum_in
        self.compute_dRoot_VAR()
        self.total_leaves = root.nleaf

    def opt_score(self):
        return self.minVAR

    def report_score(self):
        return "MinVar score: " + str(self.opt_score())


class MVDF_Tree(minVAR_Base_Tree):
    # supportive class to implement VAR-reroot + deepest node + factorization
    def __init__(self, ddpTree=None, tree_file=None, schema="newick"):
        super(MVDF_Tree, self).__init__(ddpTree, tree_file, schema)
        self.deep_node = None

    def reset(self):
        super(MVDF_Tree, self).reset()
        self.deep_node = None

    def Opt_function(self, node, a, b, c):
        x = -b / (2 * a)
        if x >= 0 and x <= node.edge_length:
            #                curr_minVAR = a*x*x + b*x + c
            factor = float(node.nleaf) / self.total_leaves
            factor = factor * (1 - factor)
            curr_minVAR = (a * x * x + b * x + c) / factor

            if node.var < node.parent_node.var:
                deep_node = node
            else:
                deep_node = node.parent_node

            updateNeed = False
            if (self.deep_node is None) or (deep_node.var < self.deep_node.var):
                self.deep_node = deep_node
                self.minVAR = curr_minVAR
                updateNeed = True
            elif (self.deep_node is deep_node) and (curr_minVAR < self.minVAR):
                self.minVAR = curr_minVAR
                updateNeed = True

            if updateNeed:
                self.opt_root = node
                self.opt_x = node.edge_length - x

    #                self.logger.info(str(curr_minVAR) + "\t" + node.label
    #                      + "\t" + str(node.edge_length-x) + "\t" + str(self.Tree_records[node.idx].var)
    #                      + "\t" + (str(node.parent_node.label) if node.parent_node else "None")
    #                      + "\t" + str(self.Tree_records[node.parent_node.idx].var))

    def compute_threshold(self, k=3.5):
        # should be called only AFTER the MV root was found
        mean = (self.opt_root.sum_total - self.opt_x *
                (self.total_leaves - 2 * self.opt_root.nleaf)) / self.total_leaves
        factor = float(self.opt_root.nleaf) / self.total_leaves
        factor = factor * (1 - factor)
        rootVar = self.minVAR * factor
        self.logger.info(mean)
        self.logger.info(rootVar)
        std = math.sqrt(rootVar)
        return mean + k * std


class MVD0_Tree(minVAR_Base_Tree):
    # supportive class to implement VAR-reroot + deepest node + no factorization
    def __init__(self, ddpTree=None, tree_file=None, schema="newick"):
        super(MVD0_Tree, self).__init__(ddpTree, tree_file, schema)
        self.deep_node = None

    def reset(self):
        super(MVD0_Tree, self).reset()
        self.deep_node = None

    def Opt_function(self, node, a, b, c):
        x = -b / (2 * a)
        if x >= 0 and x <= node.edge_length:
            curr_minVAR = a * x * x + b * x + c

            if node.var < node.parent_node.var:
                deep_node = node
            else:
                deep_node = node.parent_node

            updateNeed = False
            if (self.deep_node is None) or (deep_node.var < self.deep_node.var):
                self.deep_node = deep_node
                self.minVAR = curr_minVAR
                updateNeed = True
            elif (self.deep_node is deep_node) and (curr_minVAR < self.minVAR):
                self.minVAR = curr_minVAR
                updateNeed = True

            if updateNeed:
                self.opt_root = node
                self.opt_x = node.edge_length - x

    #                self.logger.info(str(curr_minVAR) + "\t" + node.label
    #                      + "\t" + str(node.edge_length-x) + "\t" + str(self.Tree_records[node.idx].var)
    #                      + "\t" + (str(node.parent_node.label) if node.parent_node else "None")
    #                      + "\t" + str(self.Tree_records[node.parent_node.idx].var))

    def compute_threshold(self, k=3.5):
        # should be called only AFTER the MV root was found
        mean = (self.opt_root.sum_total - self.opt_x *
                (self.total_leaves - 2 * self.opt_root.nleaf)) / self.total_leaves
        self.logger.info(mean)
        self.logger.info(self.minVAR)
        std = math.sqrt(self.minVAR)
        return mean + k * std


class MV0F_Tree(minVAR_Base_Tree):
    # supportive class to implement VAR-reroot + no deepest node + factorization
    #        def __init__(self, ddpTree = None, tree_file = None, schema = "newick"):
    #            super().__init__(ddpTree, tree_file, schema)

    def Opt_function(self, node, a, b, c):
        x = -b / (2 * a)
        if x >= 0 and x <= node.edge_length:
            #                curr_minVAR = a*x*x + b*x + c
            factor = float(node.nleaf) / self.total_leaves
            factor = factor * (1 - factor)
            curr_minVAR = (a * x * x + b * x + c) / factor
            if self.minVAR is None or curr_minVAR < self.minVAR:
                self.minVAR = curr_minVAR
                self.opt_root = node
                self.opt_x = node.edge_length - x

    #                self.logger.info(str(curr_minVAR) + "\t" + node.label
    #                      + "\t" + str(node.edge_length-x) + "\t" + str(self.Tree_records[node.idx].var)
    #                      + "\t" + (str(node.parent_node.label) if node.parent_node else "None")
    #                      + "\t" + str(self.Tree_records[node.parent_node.idx].var))

    def compute_threshold(self, k=3.5):
        # should be called only AFTER the MV root was found
        mean = (self.opt_root.sum_total - self.opt_x *
                (self.total_leaves - 2 * self.opt_root.nleaf)) / self.total_leaves
        factor = float(self.opt_root.nleaf) / self.total_leaves
        factor = factor * (1 - factor)
        rootVar = self.minVAR * factor
        self.logger.info(mean)
        self.logger.info(rootVar)
        std = math.sqrt(rootVar)
        return mean + k * std


class MV00_Tree(minVAR_Base_Tree):
    # supportive class to implement VAR-reroot + no deepest node + no factorization
    #        def __init__(self, ddpTree = None, tree_file = None, schema = "newick"):
    #            super().__init__(ddpTree, tree_file, schema)

    def Opt_function(self, node, a, b, c):
        x = -b / (2 * a)
        if x >= 0 and x <= node.edge_length:
            curr_minVAR = a * x * x + b * x + c
            if self.minVAR is None or curr_minVAR < self.minVAR:
                self.minVAR = curr_minVAR
                self.opt_root = node
                self.opt_x = node.edge_length - x

    def compute_threshold(self, k=3.5):
        # should be called only AFTER the MV root was found
        mean = (self.opt_root.sum_total - self.opt_x *
                (self.total_leaves - 2 * self.opt_root.nleaf)) / self.total_leaves
        self.logger.info(mean)
        self.logger.info(self.minVAR)
        std = math.sqrt(self.minVAR)
        return mean + k * std


class MBR_Tree(Tree_extend):
    # supportive class to implement midpoint balance root
    def __init__(self, ddpTree=None, tree_file=None, schema="newick"):
        super(MBR_Tree, self).__init__(ddpTree, tree_file, schema)

        self.BPs = []  # BPs : balance points
        self.opt_root = self.ddpTree.root
        self.opt_x = 0

    def Node_init(self, node, nleaf=1, sum_in=0, sum_out=-1):
        self.nleaf = nleaf
        self.sum_in = sum_in
        self.sum_out = sum_out

    def Opt_function(self, node):
        nleaf = node.nleaf
        mean_in = node.sum_in / nleaf
        mean_out = node.sum_out / (self.total_leaves - nleaf)
        x = (mean_out - mean_in) / 2
        if x >= 0 and x <= node.edge_length:
            self.BPs.append((node, x, mean_in + x))
            node.x = x
            node.mean = mean_in + x
        else:
            node.x = None
            node.mean = None

    def bUp_update(self, node):
        node.sum_in = 0
        if node.is_leaf():
            node.nleaf = 1
        else:
            node.nleaf = 0
            for child in node.child_nodes():
                node.nleaf += child.nleaf
                node.sum_in += child.sum_in + child.nleaf * child.edge_length

    def tDown_update(self, node, opt_function):
        child_idx = 0
        for child in node.child_nodes():
            child.sum_out = (node.sum_out + node.sum_in + child.edge_length *
                             (self.total_leaves - 2 * child.nleaf) - child.sum_in)
            opt_function(child)
            child_idx += 1

    def prepare_root(self):
        root = self.get_root()
        root.sum_out = 0
        self.total_leaves = root.nleaf
        root.x = None
        root.mean = None

    def list_balance_points(self):
        self.Topdown_label()
        self.Bottomup_update()
        self.prepare_root()
        self.Topdown_update()

        for (node, x, mean) in self.BPs:
            if node.is_leaf():
                #                  self.logger.info(node.label + "\t" + str(x) + "\t" + str(mean))
                self.logger.info(node.label + "\t" + str(x) + "\t" + str(mean))
            else:
                self.logger.info(node.label + "\t" + str(x) + "\t" + str(mean))

    def build_balance_tree(self):
        self.Topdown_label()  # keep this step for now for debugging purpose
        self.Bottomup_update()
        self.prepare_root()
        self.Topdown_update()

        # self.list_balance_points()

        self.balance_tree = self.ddpTree.extract_tree()

        # bottom up pruning
        for node in self.balance_tree.traverse_postorder():
            node.type = "real"
            node.BPbelow = False

            '''if node.is_leaf():
                self.logger.info("parent: " + node.label)# + "\t" + str(node.extraction_source.x))
            else:
                self.logger.info("parent: " + node.label)#+ "\t" + str(node.extraction_source.x))'''

            for ch in node.child_nodes():
                '''try:
                    self.logger.info("child: " + ch.label)# + "\t" + str(ch.extraction_source.x))
                except:
                    self.logger.info("child: " + ch.label) #+ "\t" + str(ch.extraction_source.x))'''

                if ch.BPbelow or (ch.extraction_source.x is not None):
                    node.BPbelow = True
                # node.BPbelow = node.BPbelow or ch.BPbelow or (ch.extraction_source.x is not None)

                if not ch.BPbelow:
                    # remove the whole clade under ch
                    # for ch1 in ch.child_nodes():
                    #    ch.remove_child(ch1)
                    edgelen = ch.edge_length
                    node.remove_child(ch)

                    if ch.extraction_source.x is not None:
                        # add a new node p at the balance point
                        # set p to be a child of node (edge length ch.edge_length - x)
                        # add a new node ch1 to be another child of p (edge length ch.mean)
                        edgelen = ch.edge_length

                        # p = self.ddpTree.node_factory()
                        # ch1 = self.ddpTree.node_factory()
                        p = Node()
                        ch1 = Node()

                        p.type = "bp"  # bp: balance-point
                        p.ref_child = ch.extraction_source  # link p to the original tree (for later use after finding midpoint)
                        ch1.type = "dm"  # dm: dummy

                        # node.remove_child(ch)
                        node.add_child(p)
                        p.add_child(ch1)

                        p.edge_length = edgelen - ch.extraction_source.x
                        ch1.edge_length = ch.extraction_source.mean

                elif ch.extraction_source.x is not None:
                    # add a new node p at the balance point
                    # set p to be a child of node (edge length ch.edge_length - x)
                    # set ch to be a child of p (edge length x)
                    # add a new node ch1 to be another child of p (edge length ch.mean)

                    edgelen = ch.edge_length

                    # p = self.ddpTree.node_factory()
                    p = Node()
                    # ch1 = self.ddpTree.node_factory()
                    ch1 = Node()

                    p.type = "bp"
                    p.ref_child = ch.extraction_source  # link p to the original tree (for later use after finding midpoint)
                    ch1.type = "dm"

                    node.remove_child(ch)
                    node.add_child(p)
                    p.add_child(ch)
                    p.add_child(ch1)

                    ch.edge_length = ch.extraction_source.x
                    p.edge_length = edgelen - ch.extraction_source.x
                    ch1.edge_length = ch.extraction_source.mean

                    # topdown pruning
        node = self.balance_tree.root
        nchild = len(node.child_nodes())
        while nchild > 0 and nchild < 2:
            # node has less than 2 children
            temp = node
            node = node.child_nodes()[0]
            temp.remove_child(node)
            if node.type == "dm":
                node = temp
                break
            nchild = len(node.child_nodes())

        self.balance_tree.root = node
        self.balance_tree.root.edge_length = None
        # balance_tree.root = None

        # mptre = MPR_Tree(ddpTree=balance_tree)
        # mptre.tree_as_newick()

        # return balance_tree

    def find_root(self):
        self.build_balance_tree()
        mptre = MPR_Tree(ddpTree=self.balance_tree)
        mptre.tree_as_newick()
        mptre.find_root()

        self.logger.info(mptre.opt_root.type)

        if mptre.opt_root.type == "bp":
            self.opt_root = mptre.opt_root.ref_child
            self.opt_x = mptre.opt_root.ref_child.x + mptre.opt_x
        elif mptre.opt_root.type == "dm":
            self.logger.info("Hmm... Is it possible that a dummy was found as the opt_root?")
        else:
            self.opt_root = mptre.opt_root.extraction_source
            self.opt_x = mptre.opt_x

        self.logger.info(self.opt_root.label)
        self.logger.info(self.opt_x)


def FastRoot(args):

    stream = args.infofile if args.infofile else stderr
    logger = new_logger(__name__, myStream=stream)
    logger.info("Running " + fastroot.PROGRAM_NAME + " version " + fastroot.PROGRAM_VERSION)
    new_argv = re.sub(' +', ' ', " ".join(argv).replace("\n", " "))
    logger.info(fastroot.PROGRAM_NAME + " was called as follows: " + new_argv)

    METHOD2FUNC = {'MP': MPR_Tree, 'MV': MV00_Tree, 'OG': OGR_Tree, 'RTT': RTT_Tree}
    METHOD2DESC = {'MP': "Midpoint", 'MV': "MinVar", 'OG': "Outgroup", 'RTT': "Root-to-Tip"}
    method = args.method.upper()

    # reading outgroups
    if args.outgroups:
        if method != 'OG':
            method = 'OG'
            logger.warning("The rooting method is set to outgroup rooting (OG) due to the presence of outgroups")
        if path.exists(args.outgroups):
            OGs = []
            for line in open(args.outgroups, 'r'):
                OGs.append(line.strip())
        else:
            OGs = args.outgroups.split()
    else:
        OGs = None

    # reading sampling times
    if args.smplTimes:
        smplTimes = {}
        for line in args.smplTimes:
            sp, t = line.strip().split()
            smplTimes[sp] = float(t)
        if method != 'RTT':
            method = 'RTT'
            logger.warning(
                "The rooting method is set to root-to-tip rooting (RTT) due to the presence of sampling times")

    if method == 'RTT' and args.smplTimes is None:
        logger.error("Need sampling times for root-to-tip rooting")
        exit()
    elif method == 'OG' and args.outgroups is None:
        logger.error("Need outgroups for outgroup rooting")
        exit()

    assert method in METHOD2FUNC, "Invalid method! Valid options: MP for midpoint, MV for minVAR, OG for outgroups, RTT for root-to-tip"
    logger.info("Rooting Method: " + METHOD2DESC[method] + " Rooting")

    if method == 'RTT':
        if args.maxIter and args.maxIter < 1000:
            logger.warning(
                "Invalid number of maximum iterations (-x). Must be at least 1000. Set back to 1000 by default.")
        maxIter = max(1000, args.maxIter) if args.maxIter else 1000
        logger.info("Maximum iterations: " + str(maxIter))
    elif args.maxIter is not None:
        logger.warning("The maximum number of iterations (-x) is only used with root-to-tip rooting (RTT)")

    # read and root each tree
    for i, line in enumerate(args.input):
        tree = read_tree(line, schema=args.schema.lower())
        if method == 'OG':
            a_tree = OGR_Tree(OGs, ddpTree=tree, logger_id=i + 1, logger_stream=stream)
        elif method == 'RTT':
            a_tree = RTT_Tree(smplTimes, ddpTree=tree, logger_id=i + 1, logger_stream=stream, maxIter=maxIter)
        else:
            a_tree = METHOD2FUNC[method](ddpTree=tree, logger_id=i + 1, logger_stream=stream)

        a_tree.Reroot()
        logger.info("Tree " + str(i + 1) + " " + a_tree.report_score())
        a_tree.tree_as_newick(outstream=args.outfile)


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input',        required=False, type=argparse.FileType('r'), default=stdin, help="Input File (default is STDIN)")
    parser.add_argument('-m', '--method',       required=False, type=str, default="MV", help="Method (MP for midpoint, MV for minVAR, OG for outgroup, RTT for root-to-tip) (default is MV)")
    parser.add_argument('-g', '--outgroups',    required=False, type=str, help="Specify the outgroups. If specifying a list of outgroups, put them between quotes (i.e. \"). Otherwise, specifying a file which containts all the outgroups. Can only be used with -m OG")
    parser.add_argument('-t', '--smplTimes',    required=False, type=argparse.FileType('r'), help="The file containing the sampling times at leaves; to be used with -m RTT")
    parser.add_argument('-o', '--outfile',      required=False, type=argparse.FileType('w'), default=stdout, help="Output File (default is STDOUT)")
    parser.add_argument('-s', '--schema',       required=False, type=str, default="newick", help="Schema of your input treefile (default is newick)")
    parser.add_argument('-f', '--infofile',     required=False, type=argparse.FileType('w'), default=None, help="Save all the logging to this file. Default: print to stderr")
    parser.add_argument("-x", "--maxIter",      required=False, type=int, default=None, help="Maximum number of iterations to run cvxopt")
    args = vars(parser.parse_args())
    FastRoot(args)