"""Minimal phylogenetics utilities (no external dependencies beyond numpy).

- parse a Newick/NEXUS tree with branch lengths
- prune to a set of tips
- Fitch parsimony for a binary trait (minimum number of changes; ACCTRAN/DELTRAN gain counts)
- Felsenstein independent contrasts for a continuous trait (Brownian-motion rate)
- ML ancestral states under Brownian motion (via rerooting-free two-pass algorithm)
"""
import re
import sys
import numpy as np

sys.setrecursionlimit(100000)


class Node:
    __slots__ = ("name", "length", "children", "parent", "data")

    def __init__(self, name=None, length=0.0):
        self.name = name
        self.length = length
        self.children = []
        self.parent = None
        self.data = {}

    def is_tip(self):
        return not self.children


def parse_newick(s):
    """Iterative Newick parser. Returns root Node."""
    s = s.strip()
    if s.endswith(";"):
        s = s[:-1]
    # strip NEXUS-style comments [&...]
    s = re.sub(r"\[[^\]]*\]", "", s)
    root = Node()
    cur = root
    i, n = 0, len(s)
    token = []

    def flush(node, tok):
        t = "".join(tok).strip()
        if not t:
            return
        if ":" in t:
            nm, ln = t.rsplit(":", 1)
            node.name = nm.strip("'") or node.name
            try:
                node.length = float(ln)
            except ValueError:
                node.length = 0.0
        else:
            node.name = t.strip("'")

    while i < n:
        c = s[i]
        if c == "(":
            child = Node()
            child.parent = cur
            cur.children.append(child)
            cur = child
            token = []
        elif c == ",":
            flush(cur, token)
            token = []
            sib = Node()
            sib.parent = cur.parent
            cur.parent.children.append(sib)
            cur = sib
        elif c == ")":
            flush(cur, token)
            token = []
            cur = cur.parent
        else:
            token.append(c)
        i += 1
    flush(cur, token)
    return root


def read_nexus_tree(path):
    """Read the first tree from a NEXUS file (handles a Translate block) or a plain Newick file."""
    txt = open(path, encoding="utf8", errors="replace").read()
    if "#NEXUS" not in txt[:200].upper():
        return parse_newick(txt)
    trans = {}
    m = re.search(r"translate(.*?);", txt, flags=re.S | re.I)
    if m:
        for item in m.group(1).split(","):
            parts = item.strip().split(None, 1)
            if len(parts) == 2:
                trans[parts[0]] = parts[1].strip().strip("'")
    m = re.search(r"^\s*tree\s+\S+\s*=\s*(?:\[[^\]]*\]\s*)?(.*?;)", txt, flags=re.S | re.I | re.M)
    root = parse_newick(m.group(1))
    if trans:
        for nd in iter_nodes(root):
            if nd.is_tip() and nd.name in trans:
                nd.name = trans[nd.name]
    return root


def iter_nodes(root):
    """Pre-order traversal (iterative)."""
    stack = [root]
    while stack:
        nd = stack.pop()
        yield nd
        stack.extend(reversed(nd.children))


def postorder(root):
    out = list(iter_nodes(root))
    out.reverse()  # reversed pre-order with children pushed reversed is a valid post-order
    return out


def tips(root):
    return [nd for nd in iter_nodes(root) if nd.is_tip()]


def prune(root, keep):
    """Return a new tree containing only tips whose name is in `keep`; unary nodes are collapsed
    (branch lengths summed)."""
    keep = set(keep)

    def rec(nd):
        if nd.is_tip():
            if nd.name in keep:
                return Node(nd.name, nd.length)
            return None
        kids = [k for k in (rec(c) for c in nd.children) if k is not None]
        if not kids:
            return None
        if len(kids) == 1:
            k = kids[0]
            k.length += nd.length
            return k
        new = Node(nd.name, nd.length)
        for k in kids:
            k.parent = new
            new.children.append(k)
        return new

    new_root = rec(root)
    if new_root is not None:
        new_root.parent = None
    return new_root


def node_heights(root):
    """Distance from root for every node; returns dict id(node)->depth and max depth."""
    depth = {id(root): 0.0}
    for nd in iter_nodes(root):
        for c in nd.children:
            depth[id(c)] = depth[id(nd)] + c.length
    return depth, max(depth.values())


def to_newick(root):
    def rec(nd):
        if nd.is_tip():
            return f"{nd.name}:{nd.length:.6f}"
        return "(" + ",".join(rec(c) for c in nd.children) + f"):{nd.length:.6f}"
    return rec(root) + ";"


# ---------------------------------------------------------------- parsimony

def fitch(root, state):
    """Fitch parsimony for a binary tip trait. `state` maps tip name -> 0/1.
    Returns dict with: steps (minimum number of changes) and n_clades (number of maximal clades in
    which every tip has state 1 = number of origins if the trait is never lost)."""
    po = postorder(root)
    S = {}
    steps = 0
    for nd in po:
        if nd.is_tip():
            S[id(nd)] = frozenset([state[nd.name]])
        else:
            sets = [S[id(c)] for c in nd.children]
            inter = frozenset.intersection(*sets)
            if inter:
                S[id(nd)] = inter
            else:
                cnt = {0: 0, 1: 0}
                for st in sets:
                    for v in st:
                        cnt[v] += 1
                mx = max(cnt.values())
                S[id(nd)] = frozenset(v for v in cnt if cnt[v] == mx)
                steps += len(sets) - mx
    all1 = {}
    for nd in po:
        if nd.is_tip():
            all1[id(nd)] = state[nd.name] == 1
        else:
            all1[id(nd)] = all(all1[id(c)] for c in nd.children)
    n_clades = 0
    for nd in iter_nodes(root):
        if all1[id(nd)] and (nd.parent is None or not all1[id(nd.parent)]):
            n_clades += 1
    return dict(steps=steps, n_clades=n_clades)


def origins(root, state, eps=1e-3):
    """Number of independent origins (0->1 changes) of a binary tip trait, with the root in state 0.
    Returns dict: gains_min = fewest gains among most-parsimonious reconstructions (losses preferred
    over parallel gains when equally parsimonious); gains_max = most gains among most-parsimonious
    reconstructions; gains_noloss = origins if the trait is never lost (maximal all-1 clades);
    steps = parsimony length with the root fixed at 0."""
    tot1, g1, l1 = min_gains_irreversible_free(root, state, loss_cost=1.0, gain_cost=1.0 + eps)
    tot2, g2, l2 = min_gains_irreversible_free(root, state, loss_cost=1.0 + eps, gain_cost=1.0)
    f = fitch(root, state)
    return dict(gains_min=g1, losses_at_min=l1, gains_max=g2, losses_at_max=l2,
                gains_noloss=f["n_clades"], steps=g1 + l1)


def min_gains_irreversible_free(root, state, loss_cost=1.0, gain_cost=1.0):
    """Sankoff parsimony with separate gain/loss costs, root constrained to state 0.
    Returns (total cost, number of gains, number of losses) for one optimal reconstruction
    (ties resolved in favour of the parent state)."""
    po = postorder(root)
    INF = 1e18
    C = {}
    for nd in po:
        if nd.is_tip():
            s = state[nd.name]
            C[id(nd)] = (0.0 if s == 0 else INF, 0.0 if s == 1 else INF)
        else:
            c0 = c1 = 0.0
            for ch in nd.children:
                a, b = C[id(ch)]
                c0 += min(a, b + gain_cost)
                c1 += min(b, a + loss_cost)
            C[id(nd)] = (c0, c1)
    total = C[id(root)][0]
    A = {id(root): 0}
    gains = losses = 0
    for nd in iter_nodes(root):
        if nd.parent is None:
            continue
        p = A[id(nd.parent)]
        a, b = C[id(nd)]
        if p == 0:
            st = 0 if a <= b + gain_cost else 1
        else:
            st = 1 if b <= a + loss_cost else 0
        A[id(nd)] = st
        if p == 0 and st == 1:
            gains += 1
        if p == 1 and st == 0:
            losses += 1
    return total, gains, losses


# ---------------------------------------------------------------- Brownian motion

def contrasts(root, x, tip_extra=0.0):
    """Felsenstein (1985) standardized independent contrasts for tip values x (dict name->value).
    Tree must be bifurcating at nodes used (polytomies are resolved arbitrarily with zero-length
    branches). Returns list of dicts (contrast, sd (sqrt of summed branch length), raw difference,
    node) and the BM rate estimate sigma2 = mean(contrast^2).
    tip_extra (in branch-length units) is added to every terminal branch; with tip_extra =
    tau2 / sigma2 this is Brownian motion plus independent tip (measurement) error of variance tau2."""
    val = {}
    extra = {}
    out = []
    for nd in postorder(root):
        if nd.is_tip():
            val[id(nd)] = x[nd.name]
            extra[id(nd)] = float(tip_extra)
            continue
        kids = list(nd.children)
        # resolve polytomy sequentially
        a = kids[0]
        va, la = val[id(a)], a.length + extra[id(a)]
        for b in kids[1:]:
            vb, lb = val[id(b)], b.length + extra[id(b)]
            tot = la + lb
            if tot <= 0:
                tot = 1e-9
            out.append(dict(contrast=(va - vb) / np.sqrt(tot), raw=va - vb, v=tot, node=nd,
                            la=la, lb=lb))
            if la <= 0 and lb <= 0:
                vn = 0.5 * (va + vb)
                ln = 0.0
            else:
                la_ = max(la, 1e-12)
                lb_ = max(lb, 1e-12)
                vn = (va / la_ + vb / lb_) / (1 / la_ + 1 / lb_)
                ln = la_ * lb_ / (la_ + lb_)
            va, la = vn, ln
        val[id(nd)] = va
        extra[id(nd)] = la  # extra length added to the branch above nd
    c = np.array([o["contrast"] for o in out])
    return out, float(np.mean(c ** 2)), val


def bm_ancestral_states(root, x):
    """ML ancestral states under Brownian motion for all internal nodes (two-pass algorithm).
    Returns dict id(node)->estimate."""
    # upward pass: conditional mean m_u and precision-equivalent variance v_u of subtree at top of node
    m, v = {}, {}
    for nd in postorder(root):
        if nd.is_tip():
            m[id(nd)] = x[nd.name]
            v[id(nd)] = 0.0
        else:
            w = 0.0
            s = 0.0
            for c in nd.children:
                vv = max(v[id(c)] + c.length, 1e-12)
                w += 1.0 / vv
                s += m[id(c)] / vv
            m[id(nd)] = s / w
            v[id(nd)] = 1.0 / w
    est = {}
    # downward pass: combine subtree info with info from the rest of the tree
    md, vd = {}, {}  # message from parent side arriving at node (at the node itself)
    for nd in iter_nodes(root):
        if nd.parent is None:
            est[id(nd)] = m[id(nd)]
            md[id(nd)] = None
            continue
        p = nd.parent
        # info at parent excluding this child
        w = 0.0
        s = 0.0
        for c in p.children:
            if c is nd:
                continue
            vv = max(v[id(c)] + c.length, 1e-12)
            w += 1.0 / vv
            s += m[id(c)] / vv
        if md[id(p)] is not None:
            vv = max(vd[id(p)], 1e-12)
            w += 1.0 / vv
            s += md[id(p)] / vv
        mp = s / w
        vp = 1.0 / w + nd.length
        md[id(nd)], vd[id(nd)] = mp, vp
        if nd.is_tip():
            est[id(nd)] = x[nd.name]
        else:
            w2 = 1.0 / max(v[id(nd)], 1e-12) + 1.0 / max(vp, 1e-12)
            est[id(nd)] = (m[id(nd)] / max(v[id(nd)], 1e-12) + mp / max(vp, 1e-12)) / w2
    return est


def bm_with_tip_error(root, x, rho_grid=None):
    """REML fit of Brownian motion (rate sigma2 per unit branch length) plus independent tip error
    (variance tau2), by profiling over rho = tau2 / sigma2 (Felsenstein-contrast likelihood; cf. Ives,
    Midford & Garland 2007 Syst Biol 56:252).  Returns dict with the estimates, the log-likelihood
    gain over rho = 0 and a likelihood-ratio 95 % interval (profile over rho)."""
    if rho_grid is None:
        rho_grid = np.concatenate([[0.0], np.logspace(-2, 3, 151)])

    def prof(rho):
        cons, _, _ = contrasts(root, x, tip_extra=rho)
        d = np.array([c["raw"] for c in cons])
        v = np.array([c["v"] for c in cons])
        s2 = float(np.mean(d * d / v))
        ll = float(-0.5 * np.sum(np.log(2 * np.pi * s2 * v) + 1.0))
        return ll, s2

    res = [prof(r) for r in rho_grid]
    ll = np.array([r[0] for r in res])
    s2 = np.array([r[1] for r in res])
    i = int(np.argmax(ll))
    ok = ll >= ll[i] - 1.92
    return dict(sigma2=float(s2[i]), tau2=float(s2[i] * rho_grid[i]), rho=float(rho_grid[i]),
                tip_error_sd=float(np.sqrt(s2[i] * rho_grid[i])),
                sigma2_lo95=float(s2[ok].min()), sigma2_hi95=float(s2[ok].max()),
                tau2_lo95=float((s2[ok] * rho_grid[ok]).min()), tau2_hi95=float((s2[ok] * rho_grid[ok]).max()),
                sigma2_no_error=float(s2[0]), dloglik_vs_no_error=float(ll[i] - ll[0]),
                n_contrasts=int(len(contrasts(root, x)[0])))


def bm_rate_given_tip_error(root, x, tau2):
    """ML Brownian rate for a subtree when the tip-error variance tau2 is fixed (taken from a larger fit)."""
    grid = np.logspace(-5, 0, 201)
    best = (-np.inf, np.nan)
    lls = []
    for s2 in grid:
        cons, _, _ = contrasts(root, x, tip_extra=tau2 / s2)
        d = np.array([c["raw"] for c in cons])
        v = np.array([c["v"] for c in cons])
        ll = float(-0.5 * np.sum(np.log(2 * np.pi * s2 * v) + d * d / (s2 * v)))
        lls.append(ll)
        if ll > best[0]:
            best = (ll, s2)
    lls = np.array(lls)
    ok = lls >= best[0] - 1.92
    return dict(sigma2=float(best[1]), sigma2_lo95=float(grid[ok].min()), sigma2_hi95=float(grid[ok].max()))


def sister_pairs(root, x):
    """All cherries (pairs of sister tips): returns arrays (total branch length separating the pair,
    difference in x)."""
    T, D = [], []
    for nd in iter_nodes(root):
        if len(nd.children) == 2 and all(c.is_tip() for c in nd.children):
            a, b = nd.children
            T.append(a.length + b.length)
            D.append(x[a.name] - x[b.name])
    return np.array(T), np.array(D)


def slice_clades(root, t_cut):
    """Cut an ultrametric tree at t_cut (time before present) and return the list of clades (lists of tip
    names) that descend from each lineage crossing the cut."""
    depth, mx = node_heights(root)
    out = []
    stack = [root]
    while stack:
        nd = stack.pop()
        age = mx - depth[id(nd)]
        if nd.is_tip() or age <= t_cut:
            out.append([t.name for t in tips(nd)])
        else:
            stack.extend(nd.children)
    return out