"""
trees.py -- subtree counts of all trees with at most 18 vertices.

Purpose: recompute the list of 34 exceptional values in Theorem 1 of
Czabarka, Szekely and Wagner (2009), "The inverse problem for certain tree
parameters".  A tree on n >= 19 vertices has at least C(n+1,2) >= 190
subtrees (each pair of vertices, possibly equal, spans a path), so the
values below 190 are exactly the subtree counts below 190 of trees with at
most 18 vertices.

Two independent generators of the free trees on n vertices:
  A. networkx.nonisomorphic_trees (Wright-Richmond-Odlyzko-McKay algorithm);
  B. our own: add a leaf to every tree on n-1 vertices in every possible way
     and reduce by a canonical form (AHU encoding rooted at the centre, or
     the smaller of the two encodings for a bicentral tree).
The two lists are compared through their canonical forms, and the counts
are checked against OEIS A000055.

Two independent subtree counters:
  1. rooted product formula   c(T) = sum_v f(v),  f(v) = prod_children (1+f(w));
  2. vertex-minus-edge formula: every subtree S satisfies |V(S)|-|E(S)| = 1, so
        c(T) = sum_v s(v) - sum_{uv in E} s(uv),
     where s(v) = #subtrees containing v and s(uv) = #subtrees containing uv,
     both computed from branch counts by rerooting.
For trees with at most 12 vertices we also compare with the ESU enumerator.
"""

import os
import sys
from collections import defaultdict

import networkx as nx

from cg import adj_from_edges, c_esu, c_tree_dp

DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data")
A000055 = [1, 1, 1, 1, 2, 3, 6, 11, 23, 47, 106, 235, 551, 1301, 3159, 7741,
           19320, 48629, 123867]          # index n = 0..18
CSW_EXCEPTIONS = [2, 4, 5, 7, 8, 9, 12, 13, 14, 16, 18, 19, 22, 23, 26, 27, 29,
                  31, 33, 35, 38, 39, 42, 43, 46, 50, 52, 54, 60, 65, 68, 72, 77, 89]
NMAX = 18


# ------------------------------------------------------------------ helpers

def nbrs_from_edges(n, edges):
    nb = [[] for _ in range(n)]
    for u, v in edges:
        nb[u].append(v)
        nb[v].append(u)
    return nb


def centres(nb):
    n = len(nb)
    if n <= 2:
        return list(range(n))
    deg = [len(x) for x in nb]
    layer = [v for v in range(n) if deg[v] == 1]
    remaining = n
    while remaining > 2:            # peel leaves layer by layer
        remaining -= len(layer)
        new = []
        for v in layer:
            for w in nb[v]:
                deg[w] -= 1
                if deg[w] == 1:
                    new.append(w)
        layer = new
    return layer


def ahu(nb, root):
    """AHU canonical string of the tree rooted at `root` (iterative)."""
    parent = {root: None}
    order = [root]
    for v in order:
        for w in nb[v]:
            if w != parent[v]:
                parent[w] = v
                order.append(w)
    code = {}
    for v in reversed(order):
        kids = sorted(code[w] for w in nb[v] if w != parent[v])
        code[v] = "(" + "".join(kids) + ")"
    return code[root]


def canon(nb):
    return min(ahu(nb, c) for c in centres(nb))


# ---------------------------------------------------------------- generators

def trees_networkx(n):
    if n == 1:
        return [[]]
    return [list(T.edges()) for T in nx.nonisomorphic_trees(n)]


def trees_by_leaf_addition(nmax):
    levels = {1: {"()": []}}
    for n in range(2, nmax + 1):
        nxt = {}
        for edges in levels[n - 1].values():
            for v in range(n - 1):
                e2 = edges + [(v, n - 1)]
                key = canon(nbrs_from_edges(n, e2))
                if key not in nxt:
                    nxt[key] = e2
        levels[n] = nxt
        print(f"  leaf-addition: n={n}: {len(nxt)} trees", file=sys.stderr)
    return levels


# ------------------------------------------------------------------ counters

def c_vertex_minus_edge(n, edges):
    if n == 1:
        return 1
    nb = nbrs_from_edges(n, edges)
    # down[v] = subtrees containing v inside the subtree of v (rooted at 0)
    parent = [-1] * n
    order = [0]
    parent[0] = 0
    for v in order:
        for w in nb[v]:
            if parent[w] == -1:
                parent[w] = v
                order.append(w)
    down = [1] * n
    for v in reversed(order):
        for w in nb[v]:
            if parent[w] == v and w != v:
                down[v] *= 1 + down[w]
    # up[v] = subtrees containing parent(v) inside the component of T - v
    # that contains parent(v)  (rerooting)
    up = [0] * n
    for v in order:
        for w in nb[v]:
            if parent[w] == v and w != v:
                prod = 1 if v == 0 else (1 + up[v])
                for u in nb[v]:
                    if parent[u] == v and u != v and u != w:
                        prod *= 1 + down[u]
                up[w] = prod
    s_vertex = 0
    for v in range(n):
        s = down[v] * (1 if v == 0 else (1 + up[v]))
        s_vertex += s
    s_edge = 0
    for v in range(1, n):
        s_edge += down[v] * up[v]          # edge v-parent(v)
    return s_vertex - s_edge


# ---------------------------------------------------------------------- main

def main():
    print("generator B (leaf addition) ...", file=sys.stderr)
    levels = trees_by_leaf_addition(NMAX)
    values = defaultdict(set)              # n -> subtree counts
    best = {}                              # value -> (n, edges)
    summary = []
    for n in range(1, NMAX + 1):
        listA = trees_networkx(n)
        canonA = {canon(nbrs_from_edges(n, e)) for e in listA}
        canonB = set(levels[n].keys())
        ok = len(listA) == len(canonA) == len(canonB) == A000055[n] and canonA == canonB
        summary.append(f"n={n}: networkx {len(listA)}, leaf-addition {len(canonB)}, "
                       f"OEIS A000055 {A000055[n]}, same isomorphism classes: {canonA == canonB}")
        assert ok, summary[-1]
        for edges in levels[n].values():
            adj = adj_from_edges(n, edges)
            c1 = c_tree_dp(adj)
            c2 = c_vertex_minus_edge(n, edges)
            assert c1 == c2, (n, edges, c1, c2)
            if n <= 12:
                assert c1 == c_esu(adj), (n, edges)
            values[n].add(c1)
            if c1 not in best or n < best[c1][0]:
                best[c1] = (n, edges)
        print(f"  n={n} done", file=sys.stderr)

    attained = set().union(*values.values())
    below190 = sorted(set(range(1, 190)) - attained)
    summary.append(f"values in [1,189] that are not subtree counts of trees with <= 18 vertices "
                   f"({len(below190)}): {below190}")
    summary.append(f"agrees with the 34 exceptions of Czabarka-Szekely-Wagner: {below190 == CSW_EXCEPTIONS}")
    minc = {n: min(values[n]) for n in values}
    summary.append("minimum subtree count by n (should be n(n+1)/2): " +
                   ", ".join(f"{n}:{minc[n]}" for n in sorted(minc)))
    assert all(minc[n] == n * (n + 1) // 2 for n in minc)
    with open(os.path.join(DATA, "trees_summary.txt"), "w") as f:
        f.write("\n".join(summary) + "\n")
    with open(os.path.join(DATA, "tree_values_below_190.csv"), "w") as f:
        f.write("k,n,edge_list\n")
        for k in range(1, 190):
            if k in best:
                n, edges = best[k]
                f.write(f"{k},{n}," + " ".join(f"{u}-{v}" for u, v in edges) + "\n")
    print("\n".join(summary))


if __name__ == "__main__":
    main()