"""
caterpillars.py -- independent check of the computer-assisted parts of
Czabarka, Szekely and Wagner (2009) for the number of subtrees.

The caterpillar C(a_0,...,a_k) (CSW, Section 3) is built from one vertex:
apply operation 1 (attach a new root to the old root) a_0 times, attach a
pendant leaf at the root, apply operation 1 a_1 times, attach a leaf, ...,
apply operation 1 a_k times.  With x_0 = a_0+1 and x_j = 2x_{j-1}+a_j,

    c(C) = T(x_k) - sum_{j<k} D(x_j),   T(x) = x(x+1)/2,  D(x) = (3x^2-x-2)/2,

where the x_j satisfy x_0 >= 1 and x_j >= 2 x_{j-1} (CSW, Lemma 2 and the
substitution after it).  We do NOT rely on this formula: every witness is
built as an explicit tree and counted by the rooted DP (cg.c_tree_dp) and,
when the count is at most 2*10^6, by the ESU enumerator (cg.c_esu).

Checks performed:
  (1) the formula agrees with direct counting on many random caterpillars;
  (2) base case of CSW Proposition 3: for 160 <= l <= 319, every integer in
          I_l = [11034, 11033 + D(l+1)]
      is a sum  sum_{j<k} D(x_j)  with x_0 >= 1, x_j >= 2x_{j-1}, x_{k-1} <= l;
  (3) the set of caterpillar subtree counts in [1, 1527] (exhaustive), to
      compare with CSW's report that caterpillars miss exactly the 34
      exceptions and 48, 61, 64, 79, 85, 93, 96, 123, 157, 183;
  (4) explicit witnesses for every k in [90, 1527] other than 89's list:
      a caterpillar when one exists, otherwise a tree from trees.py;
  (5) beyond 1527: every k in [1527, 300000] is a caterpillar count, and the
      recursive construction of the proof of Proposition 3 is run on random
      k up to 10^10, each result being built and counted.
"""

import os
import random
import sys

import numpy as np
from math import isqrt

from cg import adj_from_edges, c_esu, c_tree_dp

DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data")


def T(x):
    return x * (x + 1) // 2


def D(x):
    return (3 * x * x - x - 2) // 2


# ----------------------------------------------------------- caterpillars

def caterpillar_edges(a):
    """Edges of C(a_0,...,a_k) following the CSW construction."""
    edges = []
    root = 0
    nv = 1
    for j, aj in enumerate(a):
        for _ in range(aj):                 # operation 1
            edges.append((root, nv))
            root = nv
            nv += 1
        if j < len(a) - 1:                  # attach a pendant leaf
            edges.append((root, nv))
            nv += 1
    return nv, edges


def caterpillar_formula(a):
    x = a[0] + 1
    s = 0
    for aj in a[1:]:
        s += D(x)
        x = 2 * x + aj
    return T(x) - s


def a_from_chain(xs, y):
    """xs = (x_0,...,x_{k-1}) with x_0>=1, x_j>=2x_{j-1}; final x_k = y >= 2 x_{k-1}."""
    if not xs:
        return [y - 1]
    a = [xs[0] - 1]
    for j in range(1, len(xs)):
        a.append(xs[j] - 2 * xs[j - 1])
    a.append(y - 2 * xs[-1])
    assert min(a) >= 0
    return a


def subtree_count_lists(n, edges):
    """rooted product formula on adjacency lists (for large trees)."""
    nb = [[] for _ in range(n)]
    for u, v in edges:
        nb[u].append(v)
        nb[v].append(u)
    parent = [-1] * n
    parent[0] = 0
    order = [0]
    for v in order:
        for w in nb[v]:
            if parent[w] == -1:
                parent[w] = v
                order.append(w)
    assert len(order) == n
    f = [1] * n
    for v in reversed(order[1:]):
        f[parent[v]] *= 1 + f[v]
    return sum(f)


def count_tree(n, edges, esu_limit=2_000_000):
    adj = adj_from_edges(n, edges) if n <= 4000 else None
    c1 = c_tree_dp(adj)
    c2 = c_esu(adj, limit=esu_limit) if c1 <= esu_limit else None
    if c2 is not None:
        assert c1 == c2
    return c1, c2


# -------------------------------------------------------- minimal last term

def last_term_table(smax):
    """m[s] = least possible x_{k-1} over representations of s as a chain sum
    (m[0] = 0 for the empty sum; INF if s is not representable)."""
    INF = np.iinfo(np.int32).max
    m = np.full(smax + 1, INF, dtype=np.int32)
    m[0] = 0
    x = 2                                   # D(1) = 0 adds nothing
    while D(x) <= smax:
        d = D(x)
        src = m[: smax + 1 - d]
        tgt = m[d:]
        mask = (src <= x // 2) & (tgt == INF)
        tgt[mask] = x
        x += 1
    return m, INF


def chain_for(m, s):
    xs = []
    while s > 0:
        x = int(m[s])
        xs.append(x)
        s -= D(x)
        assert s >= 0 and m[s] <= x // 2
    return list(reversed(xs))


def main():
    rng = random.Random(20261023)
    summary = []

    # (1) formula vs direct count
    for _ in range(3000):
        k = rng.randint(0, 5)
        a = [rng.randint(0, 6) for _ in range(k + 1)]
        n, edges = caterpillar_edges(a)
        c1, c2 = count_tree(n, edges)
        assert c1 == caterpillar_formula(a) and (c2 is None or c2 == c1), a
    summary.append("(1) CSW caterpillar formula agrees with direct counting on 3000 random caterpillars")

    # table up to T(1527) (enough for exhaustive search of counts <= 1527)
    VMAX = 1527
    SMAX = T(VMAX)
    m, INF = last_term_table(max(SMAX, 2_000_000))
    print("table built", file=sys.stderr)

    # (2) base case of Proposition 3
    worst = 0
    for l in range(160, 320):
        lo, hi = 11034, 11033 + D(l + 1)
        seg = m[lo: hi + 1]
        assert seg.max() <= l, (l, int(seg.max()))
        worst = max(worst, int(seg.max()))
    summary.append("(2) CSW Proposition 3, base case 160 <= l <= 319: every integer of "
                   "I_l = [11034, 11033 + D(l+1)] has a chain representation with x_{k-1} <= l: VERIFIED")
    # also: 11033 itself is not representable with small last term? report smallest
    # s0 such that [s0, 11033+D(161)] is fully representable with x_{k-1} <= 160
    bad = [s for s in range(1, 11034) if m[s] > 160]
    summary.append(f"    (largest s < 11034 with no representation having x_(k-1) <= 160: {max(bad)})")

    # (3) exhaustive caterpillar counts <= 1527
    cat = {}
    for y in range(1, VMAX + 1):
        lo = max(0, T(y) - VMAX)
        for s in range(lo, T(y)):
            if s == 0 or (m[s] != INF and 2 * int(m[s]) <= y):
                v = T(y) - s
                if v not in cat:
                    cat[v] = (y, s)
    noncat = sorted(set(range(1, VMAX + 1)) - set(cat))
    csw_cat_exc = sorted([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,
                          48, 61, 64, 79, 85, 93, 96, 123, 157, 183])
    summary.append(f"(3) integers in [1,1527] that are not subtree counts of caterpillars "
                   f"C(a_0..a_k) ({len(noncat)}): {noncat}")
    summary.append(f"    agrees with CSW (34 exceptions + 48,61,64,79,85,93,96,123,157,183): {noncat == csw_cat_exc}")

    # (4) witnesses for 90..1527
    import csv
    tree_alt = {}
    with open(os.path.join(DATA, "tree_values_below_190.csv")) as f:
        for row in csv.DictReader(f):
            tree_alt[int(row["k"])] = (int(row["n"]), row["edge_list"])
    rows = []
    for k in range(90, VMAX + 1):
        if k in cat:
            y, s = cat[k]
            xs = chain_for(m, s)
            a = a_from_chain(xs, y)
            n, edges = caterpillar_edges(a)
            c1, c2 = count_tree(n, edges)
            assert c1 == k and caterpillar_formula(a) == k and c2 == k
            rows.append((k, "caterpillar", n, " ".join(map(str, a)), c1, c2))
        else:
            assert k in tree_alt, k
            n, el = tree_alt[k]
            edges = [tuple(map(int, e.split("-"))) for e in el.split()]
            c1, c2 = count_tree(n, edges)
            assert c1 == k == c2
            rows.append((k, "tree", n, el, c1, c2))
    with open(os.path.join(DATA, "tree_witnesses_90_1527.csv"), "w") as f:
        f.write("k,type,vertices,data,count_dp,count_esu\n")
        f.write("# type=caterpillar: data = a_0 ... a_k of C(a_0,...,a_k) (CSW construction);"
                " type=tree: data = edge list\n")
        for r in rows:
            f.write(",".join(map(str, r)) + "\n")
    summary.append(f"(4) explicit tree witnesses for every k in [90,1527] (k != 89 list): {len(rows)} written, "
                   f"each counted by DP and ESU; {sum(r[1] == 'tree' for r in rows)} non-caterpillars "
                   f"({[r[0] for r in rows if r[1] == 'tree']})")

    # (5a) every k in [1527, 300000] is a caterpillar count, found with y = least
    #      value such that T(y) - k has a representation with 2 x_(k-1) <= y
    KMAX = 300_000
    big_m, _ = last_term_table(3_000_000)
    for k in range(1527, KMAX + 1):
        y = max(1, isqrt(2 * k) - 1)
        while T(y) < k:
            y += 1
        found = False
        for yy in range(y, y + 400):
            s = T(yy) - k
            if s >= len(big_m):
                break
            if s == 0 or (big_m[s] != INF and 2 * int(big_m[s]) <= yy):
                found = True
                break
        assert found, k
    summary.append(f"(5a) every k in [1527, {KMAX}] is the subtree count of a caterpillar C(a_0..a_k): VERIFIED")

    # (5b) the recursive construction of the proof, on random large k
    def least_l(n):
        """least l >= 160 with n <= 11033 + D(l+1)  (n >= 11034)."""
        l = max(160, isqrt(2 * (n - 11033) // 3) - 2)
        while 11033 + D(l + 1) < n:
            l += 1
        while l > 160 and 11033 + D(l) >= n:
            l -= 1
        return l

    def rep_I(n):
        """Represent n >= 11034 as a chain sum with x_(k-1) <= least_l(n),
        following the induction of CSW Proposition 3 down to the computer
        base case 160 <= l <= 319 (iterative)."""
        tail = []
        while True:
            l = least_l(n)
            if l <= 319:
                assert m[n] <= l
                return chain_for(m, n) + list(reversed(tail))
            n2 = n - D(l)                   # last term x = l
            assert 11034 <= n2 <= 11033 + D(l // 2 + 1)
            tail.append(l)
            n = n2

    tested = 0
    for _ in range(300):
        k = rng.randint(1527, 10 ** 10)
        # find l >= 160 with k in [T(2l) - 11033 - D(l+1), T(2l) - 11034]
        l = 160
        while T(2 * l) - 11034 < k:
            l += 1
        assert T(2 * l) - 11033 - D(l + 1) <= k
        s = T(2 * l) - k
        xs = rep_I(s)
        assert all(xs[j] >= 2 * xs[j - 1] for j in range(1, len(xs))) and xs[0] >= 1 and xs[-1] <= l
        a = a_from_chain(xs, 2 * l)
        assert caterpillar_formula(a) == k
        n, edges = caterpillar_edges(a)
        # direct count on the explicit edge list (adjacency-list DP)
        assert subtree_count_lists(n, edges) == k
        tested += 1
    summary.append(f"(5b) CSW Proposition 3 construction run on {tested} random k in [1527, 10^10]; "
                   f"each caterpillar built and counted directly: all correct")

    with open(os.path.join(DATA, "caterpillars_summary.txt"), "w") as f:
        f.write("\n".join(summary) + "\n")
    print("\n".join(summary))


if __name__ == "__main__":
    main()