"""Executable evidence for the graceful tree conjecture record set.

A graceful labeling of a tree T with m edges is an injection
l : V(T) -> {0, 1, ..., m} whose induced edge labels
{ |l(u) - l(v)| : uv in E(T) } are exactly {1, 2, ..., m}.

This script does four things with exact integer arithmetic and no third-party
packages:

  1. Generates every free tree on n vertices up to isomorphism by centroid
     decomposition, and checks the counts against OEIS A000055.
  2. Searches for a graceful labeling of each tree by backtracking over vertex
     labels in a depth-first order, with isomorphic sibling subtrees forced
     into increasing label order.
  3. Searches for an alpha-labeling (a graceful labeling that respects the
     bipartition) of each tree, and reports the trees that have none.
  4. Records search cost, so the scaling of the exhaustive route is visible.

Usage:

    python3 research/graceful-trees/checks.py --max-n 18
    python3 research/graceful-trees/checks.py --max-n 20 --skip-alpha
    python3 research/graceful-trees/checks.py --max-n 14 --json
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from functools import lru_cache

# Number of trees on n unlabeled nodes, OEIS A000055, n = 0, 1, 2, ...
# An independent check on the generator, never an input to it.
A000055 = [
    1, 1, 1, 1, 2, 3, 6, 11, 23, 47, 106, 235, 551, 1301, 3159, 7741,
    19320, 48629, 123867, 317955, 823065, 2144505, 5623756, 14828074,
]


# ---------------------------------------------------------------------------
# Rooted trees up to isomorphism
#
# A rooted tree is the sorted tuple of its children's canonical forms. Sorting
# makes the form canonical: two rooted trees are isomorphic exactly when their
# tuples are equal.
# ---------------------------------------------------------------------------


@lru_cache(maxsize=None)
def rooted_trees(size: int) -> tuple:
    """Every rooted tree on `size` nodes, canonical, as nested tuples."""
    if size <= 0:
        return ()
    if size == 1:
        return ((),)
    return tuple(_multisets(size - 1, size - 1))


def _multisets(remaining: int, max_size: int) -> list:
    """Child multisets with `remaining` nodes, no child larger than `max_size`.

    Children come out in non-increasing canonical order, so each multiset is
    produced exactly once.
    """
    if remaining == 0:
        return [()]
    if max_size <= 0:
        return []
    out = []
    for size in range(min(remaining, max_size), 0, -1):
        for child in rooted_trees(size):
            for rest in _multisets_after(remaining - size, size, child):
                out.append((child,) + rest)
    return out


def _multisets_after(remaining: int, size: int, child: tuple) -> list:
    """Continue a multiset whose last child was `child`, of `size` nodes."""
    if remaining == 0:
        return [()]
    out = []
    if size <= remaining:
        for c in rooted_trees(size):
            if c > child:
                continue
            for rest in _multisets_after(remaining - size, size, c):
                out.append((c,) + rest)
    for smaller in range(min(remaining, size - 1), 0, -1):
        for c in rooted_trees(smaller):
            for rest in _multisets_after(remaining - smaller, smaller, c):
                out.append((c,) + rest)
    return out


# ---------------------------------------------------------------------------
# Free trees by centroid decomposition
#
# Every tree has one centroid, or two adjacent centroids. Rooting at a centroid
# bounds every branch by floor(n/2) nodes, so splitting on the two cases
# enumerates each free tree once and needs only small rooted trees.
# ---------------------------------------------------------------------------


def free_trees(n: int):
    """Yield every free tree on n vertices as an edge list, up to isomorphism."""
    if n <= 0:
        return
    if n == 1:
        yield []
        return

    # One centroid: every branch holds at most (n-1)//2 nodes.
    for children in _multisets(n - 1, (n - 1) // 2):
        yield _edges_from_rooted(children)

    # Two adjacent centroids: two rooted halves of n/2 nodes joined by an edge.
    if n % 2 == 0:
        halves = rooted_trees(n // 2)
        for i, left in enumerate(halves):
            for right in halves[: i + 1]:
                yield _edges_from_two_rooted(left, right)


def _edges_from_rooted(children: tuple) -> list:
    """Edge list for a root whose child subtrees are `children`."""
    edges: list[tuple[int, int]] = []
    counter = [1]

    def attach(parent: int, subtree: tuple) -> None:
        node = counter[0]
        counter[0] += 1
        edges.append((parent, node))
        for grandchild in subtree:
            attach(node, grandchild)

    for child in children:
        attach(0, child)
    return edges


def _edges_from_two_rooted(left: tuple, right: tuple) -> list:
    """Edge list for two rooted trees joined by an edge between their roots."""
    edges = _edges_from_rooted(left)
    offset = _node_count(left)
    edges.append((0, offset))
    edges.extend((u + offset, v + offset) for u, v in _edges_from_rooted(right))
    return edges


def _node_count(subtree: tuple) -> int:
    return 1 + sum(_node_count(child) for child in subtree)


# ---------------------------------------------------------------------------
# Labeling search
# ---------------------------------------------------------------------------


class Timeout(Exception):
    """The node budget ran out before the search settled the question."""


class Search:
    """Backtracking search for a graceful or alpha-labeling of one tree.

    Vertices are labeled in a depth-first order rooted at a high-degree vertex,
    so every step after the first fixes exactly one edge label and can be
    rejected the moment that label repeats.

    Two symmetry breaks keep the search from wandering through relabelings of
    the same solution:

      * The root takes only labels up to m/2, using the complement symmetry
        l -> m - l.
      * Sibling subtrees that are isomorphic are forced into increasing root
        label order. Swapping them is an automorphism, so a labeling exists
        exactly when one exists in that order. Without this, a vertex with k
        leaf children costs k! redundant branches.
    """

    def __init__(
        self,
        n: int,
        edges: list,
        node_budget: int = 2_000_000,
        symmetry: bool = True,
        edge_order: bool = True,
    ):
        self.n = n
        self.m = n - 1
        self.edges = edges
        self.adj = [[] for _ in range(n)]
        for u, v in edges:
            self.adj[u].append(v)
            self.adj[v].append(u)
        self.nodes = 0
        self.budget = node_budget
        self.symmetry = symmetry
        self.edge_order = edge_order
        self._plan = None

    # -- search order -------------------------------------------------------

    def _shape(self, vertex: int, parent: int) -> tuple:
        return tuple(
            sorted(self._shape(w, vertex) for w in self.adj[vertex] if w != parent)
        )

    def plan(self) -> list:
        """Order of (vertex, parent, twin) triples for the search.

        `twin` is the previously placed isomorphic sibling, or -1. The search
        requires a vertex to take a larger label than its twin.
        """
        if self._plan is not None:
            return self._plan
        root = max(range(self.n), key=lambda v: len(self.adj[v]))
        order = [(root, -1, -1)]

        def descend(vertex: int, parent: int) -> None:
            children = [w for w in self.adj[vertex] if w != parent]
            shaped = sorted(
                ((self._shape(w, vertex), w) for w in children),
                key=lambda pair: (len(str(pair[0])), pair[0], pair[1]),
            )
            previous_shape, previous_vertex = None, -1
            for shape, child in shaped:
                twin = previous_vertex if shape == previous_shape else -1
                order.append((child, vertex, twin))
                previous_shape, previous_vertex = shape, child
                descend(child, vertex)

        descend(root, -1)
        self._plan = order
        return order

    # -- graceful labeling --------------------------------------------------

    def graceful(self):
        """A graceful labeling as a list of vertex labels, or None."""
        if self.n == 1:
            return [0]
        order = self.plan()
        root = order[0][0]
        label = [-1] * self.n
        used_label = [False] * (self.m + 1)
        used_edge = [False] * (self.m + 1)
        for first in range(self.m // 2 + 1):
            label[root] = first
            used_label[first] = True
            if self._extend(order, 1, label, used_label, used_edge, None):
                return list(label)
            used_label[first] = False
            label[root] = -1
        return None

    def alpha(self) -> bool:
        """True when the tree has an alpha-labeling.

        An alpha-labeling is a graceful labeling with a threshold k such that
        every edge has one end at most k and the other above k. The two
        bipartition classes then occupy the blocks {0..k} and {k+1..m}, so the
        class sizes fix k. One orientation suffices: the complement symmetry
        l -> m - l maps an alpha-labeling of one orientation to the other.
        """
        if self.n <= 2:
            return True
        side = self._bipartition()
        block = sum(1 for s in side if s == 0)
        allowed = [
            (0, block - 1) if s == 0 else (block, self.m) for s in side
        ]
        order = self.plan()
        root = order[0][0]
        label = [-1] * self.n
        used_label = [False] * (self.m + 1)
        used_edge = [False] * (self.m + 1)
        low, high = allowed[root]
        for first in range(low, high + 1):
            label[root] = first
            used_label[first] = True
            if self._extend(order, 1, label, used_label, used_edge, allowed):
                return True
            used_label[first] = False
            label[root] = -1
        return False

    def _extend(self, order, index, label, used_label, used_edge, allowed) -> bool:
        if index == len(order):
            return True
        self.nodes += 1
        if self.nodes > self.budget:
            raise Timeout()
        vertex, parent, twin = order[index]
        anchor = label[parent]
        low, high = allowed[vertex] if allowed else (0, self.m)
        if twin >= 0 and self.symmetry:
            low = max(low, label[twin] + 1)
        # Large edge labels are the scarce resource: only a few vertex pairs can
        # ever produce them, and every one of 1..m has to appear. Trying the
        # candidate that consumes the largest still-free edge label first turns
        # the path family from the worst case into an easy one.
        candidates = range(low, high + 1)
        if self.edge_order:
            candidates = sorted(candidates, key=lambda v: -abs(v - anchor))
        for value in candidates:
            if used_label[value]:
                continue
            edge = value - anchor if value > anchor else anchor - value
            if used_edge[edge]:
                continue
            used_label[value] = True
            used_edge[edge] = True
            label[vertex] = value
            if self._extend(order, index + 1, label, used_label, used_edge, allowed):
                return True
            label[vertex] = -1
            used_label[value] = False
            used_edge[edge] = False
        return False

    # -- tree shape ---------------------------------------------------------

    def _bipartition(self) -> list:
        side = [-1] * self.n
        side[0] = 0
        stack = [0]
        while stack:
            v = stack.pop()
            for w in self.adj[v]:
                if side[w] < 0:
                    side[w] = 1 - side[v]
                    stack.append(w)
        return side

    def is_caterpillar(self) -> bool:
        """True when deleting every leaf leaves a path or nothing."""
        if self.n <= 3:
            return True
        degree = [len(a) for a in self.adj]
        spine = [v for v in range(self.n) if degree[v] > 1]
        if not spine:
            return True
        inner = {v: sum(1 for w in self.adj[v] if degree[w] > 1) for v in spine}
        return all(d <= 2 for d in inner.values())

    def diameter(self) -> int:
        far, _ = self._farthest(0)
        _, distance = self._farthest(far)
        return distance

    def _farthest(self, start: int):
        seen = [-1] * self.n
        seen[start] = 0
        queue = [start]
        best, best_distance = start, 0
        while queue:
            v = queue.pop(0)
            for w in self.adj[v]:
                if seen[w] < 0:
                    seen[w] = seen[v] + 1
                    if seen[w] > best_distance:
                        best, best_distance = w, seen[w]
                    queue.append(w)
        return best, best_distance

    def max_degree(self) -> int:
        return max(len(a) for a in self.adj)


def is_graceful_labeling(n: int, edges: list, label: list) -> bool:
    """Independent check that a labeling really is graceful."""
    m = n - 1
    if sorted(label) != list(range(m + 1)):
        return False
    induced = sorted(abs(label[u] - label[v]) for u, v in edges)
    return induced == list(range(1, m + 1))


# ---------------------------------------------------------------------------
# Verification sweep
# ---------------------------------------------------------------------------


def verify(
    max_n: int,
    skip_alpha: bool = False,
    node_budget: int = 2_000_000,
    symmetry: bool = True,
    edge_order: bool = True,
) -> dict:
    report = {
        "max_n": max_n,
        "symmetry_breaking": symmetry,
        "edge_label_ordering": edge_order,
        "node_budget": node_budget,
        "sizes": [],
        "total_trees": 0,
        "total_graceful": 0,
        "counts_match_a000055": True,
        "not_graceful": [],
        "unresolved": [],
        "alpha_failures": [],
        "alpha_failure_totals": {},
        "hardest": None,
    }
    started = time.time()
    hardest_nodes = -1

    for n in range(1, max_n + 1):
        size_started = time.time()
        trees = graceful = alpha_missing = size_nodes = 0
        for edges in free_trees(n):
            trees += 1
            search = Search(
                n, edges, node_budget=node_budget,
                symmetry=symmetry, edge_order=edge_order,
            )
            try:
                labeling = search.graceful()
            except Timeout:
                report["unresolved"].append({"n": n, "edges": edges})
                labeling = None
            else:
                if labeling is None:
                    report["not_graceful"].append({"n": n, "edges": edges})
                else:
                    assert is_graceful_labeling(n, edges, labeling), (n, edges, labeling)
                    graceful += 1
            size_nodes += search.nodes
            if search.nodes > hardest_nodes:
                hardest_nodes = search.nodes
                report["hardest"] = {
                    "n": n,
                    "edges": edges,
                    "search_nodes": search.nodes,
                    "diameter": search.diameter(),
                    "max_degree": search.max_degree(),
                    "caterpillar": search.is_caterpillar(),
                }
            if not skip_alpha and n >= 2:
                alpha_search = Search(
                    n, edges, node_budget=node_budget,
                    symmetry=symmetry, edge_order=edge_order,
                )
                try:
                    has_alpha = alpha_search.alpha()
                except Timeout:
                    has_alpha = True  # counted as unresolved, not as a failure
                    report["unresolved"].append({"n": n, "edges": edges, "kind": "alpha"})
                if not has_alpha:
                    alpha_missing += 1
                    if len(report["alpha_failures"]) < 25:
                        report["alpha_failures"].append(
                            {
                                "n": n,
                                "edges": edges,
                                "diameter": alpha_search.diameter(),
                                "max_degree": alpha_search.max_degree(),
                                "caterpillar": alpha_search.is_caterpillar(),
                            }
                        )
        expected = A000055[n] if n < len(A000055) else None
        if expected is not None and expected != trees:
            report["counts_match_a000055"] = False
        if not skip_alpha:
            report["alpha_failure_totals"][str(n)] = alpha_missing
        report["sizes"].append(
            {
                "n": n,
                "trees": trees,
                "expected_trees": expected,
                "graceful": graceful,
                "alpha_missing": None if skip_alpha else alpha_missing,
                "search_nodes": size_nodes,
                "seconds": round(time.time() - size_started, 3),
            }
        )
        report["total_trees"] += trees
        report["total_graceful"] += graceful

    report["elapsed_seconds"] = round(time.time() - started, 3)
    report["all_graceful"] = (
        report["total_trees"] == report["total_graceful"] and not report["unresolved"]
    )
    return report


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--max-n", type=int, default=16)
    parser.add_argument("--skip-alpha", action="store_true")
    parser.add_argument("--node-budget", type=int, default=2_000_000)
    parser.add_argument(
        "--no-symmetry",
        action="store_true",
        help="drop the isomorphic-sibling ordering constraint, to measure its effect",
    )
    parser.add_argument(
        "--no-edge-order",
        action="store_true",
        help="try vertex labels in ascending order instead of largest-edge-label first",
    )
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    sys.setrecursionlimit(20000)
    report = verify(
        args.max_n,
        skip_alpha=args.skip_alpha,
        node_budget=args.node_budget,
        symmetry=not args.no_symmetry,
        edge_order=not args.no_edge_order,
    )

    if args.json:
        print(json.dumps(report, indent=2))
        return 0 if report["all_graceful"] else 1

    print(f"free trees and graceful labelings through n = {report['max_n']}")
    header = (
        f"{'n':>3}  {'trees':>7}  {'A000055':>8}  {'graceful':>9}  "
        f"{'no alpha':>9}  {'nodes':>11}  {'sec':>7}"
    )
    print(header)
    print("-" * len(header))
    for row in report["sizes"]:
        alpha = "-" if row["alpha_missing"] is None else row["alpha_missing"]
        print(
            f"{row['n']:>3}  {row['trees']:>7}  {str(row['expected_trees']):>8}  "
            f"{row['graceful']:>9}  {str(alpha):>9}  {row['search_nodes']:>11}  "
            f"{row['seconds']:>7.2f}"
        )
    print()
    print(f"trees checked:        {report['total_trees']}")
    print(f"graceful:             {report['total_graceful']}")
    print(f"counts match A000055: {report['counts_match_a000055']}")
    print(f"every tree graceful:  {report['all_graceful']}")
    if report["hardest"]:
        h = report["hardest"]
        print(
            f"hardest search:       n={h['n']}, {h['search_nodes']} nodes, "
            f"diameter {h['diameter']}, max degree {h['max_degree']}, "
            f"caterpillar={h['caterpillar']}"
        )
    if report["alpha_failures"]:
        first = report["alpha_failures"][0]
        print(f"smallest tree with no alpha-labeling: n={first['n']}, edges={first['edges']}")
        totals = {k: v for k, v in report["alpha_failure_totals"].items() if v}
        print(f"trees with no alpha-labeling by n:    {totals}")
    if report["unresolved"]:
        print(f"unresolved within node budget:        {len(report['unresolved'])}")
    print(f"elapsed:              {report['elapsed_seconds']}s")
    return 0 if report["all_graceful"] else 1


if __name__ == "__main__":
    raise SystemExit(main())