"""The Gaspers-Lee compound analysis for maximum degree 4 and general graphs, rerun from the repaired subcubic
measure of part (a) of Theorem C.6 (`thm:is`) (uniform separator weight; see is3_uniform_repair.py), with the potential checks of
Lemma C.11 (`lem:psi-2223`) for the new section-1 weights.

Same method as gaspers_lee_general.py (port of the authors' program deg5plus.py, five degree-4 sections with pivots
3, 16/5, 7/2, 15/4, then degrees 5 and 6), whose functions are copied here so that this script runs on its own:
explicit rational weights are certified, linear conditions exactly, branching vectors in mpmath interval arithmetic,
section running times by exact vertex enumeration. Part 2 checks the numbers used in the proof of Lemma C.11 (`lem:psi-2223`) and
the psi-adjusted section-1 rows, as is3_revision_checks.py (parts 4 and 5) does for the earlier weights.
Run (from the repository root):  .venv/bin/python experiments/algorithmic-consequences/gaspers_lee_general_uniform.py
"""
from fractions import Fraction as Fr
from itertools import combinations_with_replacement as cwr
import numpy as np
from mpmath import iv
from scipy.optimize import minimize

iv.dps = 40
fails = []

# ------------------------------------------------------------------ constraint families (generic arithmetic)


def deg4_vectors(w2, w3, w4, psi):
    """Table 1 of the journal (Fig. 4), by section. psi is the potential of the (2,2,2,2) cases."""
    return [
        [(5 * w4 - 4 * w3 + 4 * w2 + psi,) * 2,            # (2,2,2,2)(a); the potential is removed
         (4 * w4 - 2 * w3 + 3 * w2 - psi,) * 2,            # (2,2,2,2)(b); the potential is added
         (4 * w4 - 2 * w3 + 2 * w2, 4 * w4 - 2 * w3 + 3 * w2),   # (2,2,2,3); journal prints "2 w1"
         (5 * w4 - 4 * w3 + 3 * w2, 6 * w4 - 4 * w3 + 3 * w2),   # (2,2,2,4)
         (3 * w4, 5 * w4 - 2 * w3 + 2 * w2),                     # (2,2,3,3)
         (4 * w4 - 2 * w3 + w2, 5 * w4 - 2 * w3 + 2 * w2),       # (2,2,3,4)
         (5 * w4 - 4 * w3 + 2 * w2, 7 * w4 - 4 * w3 + 2 * w2)],  # (2,2,4,4)
        [(2 * w4 + 2 * w3 - 2 * w2, 4 * w4 + w2)],               # (2,3,3,3)
        [(3 * w4 - w2, 6 * w4 - 2 * w3 + w2),                    # (2,3,3,4)
         (4 * w4 - 2 * w3, 6 * w4 - 2 * w3 + w2),                # (2,3,4,4)
         (5 * w4 - 4 * w3 + w2, 8 * w4 - 4 * w3 + w2),           # (2,4,4,4)
         (w4 + 4 * w3 - 4 * w2, 5 * w4),                         # (3,3,3,3)
         (2 * w4 + 2 * w3 - 3 * w2, 5 * w4)],                    # (3,3,3,4)
        [(3 * w4 - 2 * w2, 7 * w4 - 2 * w3),                     # (3,3,4,4)
         (4 * w4 - 2 * w3 - w2, 7 * w4 - 2 * w3)],               # (3,4,4,4)
        [(5 * w4 - 4 * w3, 9 * w4 - 4 * w3)]]                    # (4,4,4,4)


SECTIONS = [(2, 3), (3, Fr(16, 5)), (Fr(16, 5), Fr(7, 2)), (Fr(7, 2), Fr(15, 4)), (Fr(15, 4), 4)]


def out_edges(i):
    """Lemma 7: number of edges leaving N[v] for neighbor degrees i of v (degree 5 or 6)."""
    if len(i) == 5:
        o = 3 if sum(i) % 2 == 0 else 4
        return 5 if i == (2, 2, 2, 2, 2) else o
    o = 3 if sum(i) % 2 == 1 else 4
    if i == (2, 2, 2, 2, 2, 2):
        return 6
    return 5 if i == (2, 2, 2, 2, 2, 3) else o


def deg56_vectors(w):
    """Lemma 8 (as in deg5plus.py): w = [0, 0, w2, ..., w6]."""
    res = []
    for d, rng in ((5, [2, 3, 4, 5]), (6, [2, 3, 4, 5, 6])):
        for i in cwr(rng, d):
            if i == (d,) * d:
                continue                       # regular graph: happens once per component
            deg2 = sum(1 for k in i if k == 2)
            f1 = w[d] + sum(w[k] for k in i) + out_edges(i) * (w[d] - w[d - 1])
            f2 = w[d] + sum(w[k] - w[k - 1] for k in i) + deg2 * (w[d] - w[d - 1])
            res.append((f1, f2))
    return res


def deg56_linear(w, w4final):
    """Handover from degree 4 (weights at least the last degree-4 section), monotonicity and concavity."""
    return [w[2] - w4final[0], w[3] - w4final[1], w[4] - w4final[2],
            w[6] - w[5], w[5] - w[4], w[4] - w[3], w[3] - w[2],
            (w[5] - w[4]) - (w[6] - w[5]), (w[4] - w[3]) - (w[5] - w[4]), (w[3] - w[2]) - (w[4] - w[3]),
            w[2] - (w[3] - w[2])]


# ------------------------------------------------------------------ floating-point optimization


def solve(obj, ineqs, eqs, x0, starts=25):
    cons = [{"type": "ineq", "fun": lambda x: np.array(ineqs(x))}]
    if eqs:
        cons.append({"type": "eq", "fun": lambda x: np.array(eqs(x))})
    best = None
    rng = np.random.default_rng(0)
    for k in range(starts):
        xs = np.array(x0) * (1 + (0.1 * rng.standard_normal(len(x0)) if k else 0))
        r = minimize(obj, xs, constraints=cons, method="SLSQP", options={"ftol": 1e-14, "maxiter": 5000})
        if r.success and min(ineqs(r.x)) > -1e-9 and (not eqs or max(map(abs, eqs(r.x))) < 1e-8):
            if best is None or r.fun < best.fun:
                best = r
    return best.x


def br(v):
    return [1 - 2.0 ** -a - 2.0 ** -b for a, b in v]


def float_chain(W2, W3, margin=0.0):
    rows = []
    # section 1: x = (w4, psi)
    x = solve(lambda x: x[0], lambda x: br(deg4_vectors(W2, W3, x[0], x[1])[0]), None, [0.33, 0.0])
    w = [W2, W3, x[0]]
    rows.append((x[0], x[1], 0.0))
    for k in range(1, 5):
        p0 = float(SECTIONS[k][0])
        f = lambda a, w=w, p0=p0: [w[0] - (2 - p0) * a[0], w[1] - (3 - p0) * a[0], w[2] - (4 - p0) * a[0]]
        a = solve(lambda a: -a[0], lambda a, f=f, k=k: br(deg4_vectors(*f(a), 0)[k]), None, [0.01])
        w = f(a)
        rows.append((w[2], 0.0, a[0]))
    wf = w
    x = solve(lambda x: x[4], lambda x: [g - margin for g in br(deg56_vectors([0, 0, *x]))]
              + [g - margin for g in deg56_linear([0, 0, *x], wf)], None, [wf[0], wf[1], wf[2], 0.3, 0.31])
    return rows, wf, x


def section_time(w, p0, p1):
    """max sum a_i w_i over distributions on degrees {2,3,4} with average degree in [p0, p1] (exact)."""
    best = None
    degs = [2, 3, 4]
    cands = [{d: Fr(1)} for d in degs if p0 <= d <= p1]
    for i in degs:
        for j in degs:
            if i < j:
                for p in (Fr(p0), Fr(p1)):
                    if i <= p <= j:
                        t = (p - i) / (j - i)
                        cands.append({i: 1 - t, j: t})
    for c in cands:
        v = sum(c[d] * w[d - 2] for d in c)
        best = v if best is None or v > best else best
    return best


def ivq(q):
    return iv.mpf(q.numerator) / iv.mpf(q.denominator)


def check_vectors(vs, label):
    worst = 0
    for a, b in vs:
        s = iv.mpf(2) ** (-ivq(a)) + iv.mpf(2) ** (-ivq(b))
        worst = max(worst, float(s.b))
        if not s.b <= 1:
            fails.append((label, float(a), float(b)))
    return worst


def q(x, digits=7, up=None):
    s = 10 ** digits
    v = Fr(round(x * s), s)
    if up is True:
        v += Fr(1, s)
    if up is False:
        v -= Fr(1, s)
    return v


def certify(W2, W3, label, chain_rows=None):
    rows, wf, x = float_chain(float(W2), float(W3))
    print("%s, floating point: degree-4 final %s, w6 = %.7f, 2^w6 = %.5f" % (label, np.round(wf, 7), x[4], 2 ** x[4]))
    _, _, x = float_chain(float(W2), float(W3), margin=1e-6)       # degrees 5-6 with slack, for rounding
    w4, psi = q(rows[0][0], up=True), q(rows[0][1])
    if psi < 0:
        fails.append((label, "psi < 0"))
    w = [W2, W3, w4]
    secw = [w]
    worst = [check_vectors(deg4_vectors(W2, W3, w4, psi)[0], label + " section 1")]
    times = [section_time(w, *SECTIONS[0])]
    for k in range(1, 5):
        p0 = Fr(SECTIONS[k][0])
        alpha = q(rows[k][2], up=False)
        if alpha < 0:
            fails.append((label, "alpha < 0", k))
        w = [w[0] - (2 - p0) * alpha, w[1] - (3 - p0) * alpha, w[2] - (4 - p0) * alpha]   # pivot shift, exact
        secw.append(w)
        worst.append(check_vectors(deg4_vectors(*w, 0)[k], "%s section %d" % (label, k + 1)))
        times.append(section_time(w, *SECTIONS[k]))
    # A case with highest average degree h can occur whenever d(G) <= h, so the weights of every section must also
    # satisfy the vectors of all later sections (the measure is the minimum over sections).
    cross = 0
    for k in range(5):
        for j in range(k + 1, 5):
            cross = max(cross, check_vectors(deg4_vectors(*secw[k], 0)[j], "%s cross %d-%d" % (label, k + 1, j + 1)))
    t4 = max(times)
    W = [0, 0] + [q(v, up=True) for v in x[:4]] + [q(x[4], up=True)]
    W[2], W[3], W[4] = max(W[2], w[0]), max(W[3], w[1]), max(W[4], w[2])
    for i, g in enumerate(deg56_linear(W, w)):
        if g < 0 or (i >= 3 and g == 0):
            fails.append((label, "degree 5-6 linear", i, float(g)))
    worst.append(check_vectors(deg56_vectors(W), label + " degrees 5-6"))
    base4 = float((iv.mpf(2) ** ivq(t4)).b)
    base = (iv.mpf(2) ** ivq(W[6]))
    if not TAU18 < float(base.a):
        fails.append((label, "tau(1,8) larger"))
    print("   certified: section weights (w2, w3, w4) ending at %s; largest branching sum %.7f; later-section cases"
          " under earlier-section weights: largest sum %.4f" % ([round(float(v), 7) for v in w], max(worst), cross))
    print("   maximum degree 4: 2^(%.7f n) < %.5f^n; general graphs: w2..w6 = %s, 2^(%s n) < %.5f^n"
          % (float(t4), base4, [round(float(v), 7) for v in W[2:]], W[6], float(base.b)))
    return base4, float(base.b), secw, psi, W, rows


TAU18 = 1.2320547            # upper bound for the root 1.23205463 of x^8 = x^7 + 1 (degree >= 7)

# ------------------------------------------------------------------ 1. the chain
W2P, W3P = Fr("0.0052182"), Fr("0.176755")      # (w2', w3') of part (a) of Theorem C.6 (`thm:is`), certified in is3_uniform_repair.py
print("1. Polynomial space, from the subcubic weights (w2', w3') = (%s, %s) of the repaired thm:is:3" % (W2P, W3P))
b4, b, secw, psi, W, rows = certify(W2P, W3P, "   repaired")
alphas = [Fr(0)] + [q(rows[k][2], up=False) for k in range(1, 5)]
print("   section weights (exact decimals):")
for k, (w, al) in enumerate(zip(secw, alphas)):
    print("     section %d: alpha = %s, (w2, w3, w4) = (%s, %s, %s)" % (k + 1, float(al), *[float(x) for x in w]))
for k, w in enumerate(secw):
    print("     section %d time exponent tau_%d = %s" % (k + 1, k + 1, float(section_time(w, *SECTIONS[k]))))
t4 = max(section_time(w, *SECTIONS[k]) for k, w in enumerate(secw))
print("   2^max tau_k <= %.7f; 2^w6 in [%.7f, %.7f]"
      % (float((iv.mpf(2) ** ivq(t4)).b), float((iv.mpf(2) ** ivq(W[6])).a), float((iv.mpf(2) ** ivq(W[6])).b)))
reg = []
for d in (5, 6):
    i = (d,) * d
    f1 = W[d] + sum(W[k] for k in i) + out_edges(i) * (W[d] - W[d - 1])
    f2 = W[d] + sum(W[k] - W[k - 1] for k in i)
    reg.append(float((iv.mpf(2) ** (-ivq(f1)) + iv.mpf(2) ** (-ivq(f2))).b))
print("   regular neighborhoods (5,...,5) and (6,...,6), not required: sums %s (why the selection rule is changed)"
      % [round(x, 4) for x in reg])
print("   psi = %s; degrees 5-6: (w2, ..., w6) = %s" % (float(psi), [float(x) for x in W[2:]]))
print("   maximum degree 4: O(%.5f^n); general graphs O(%.5f^n); coloring O((1 + %.5f)^n)" % (b4, b, b))
for lab, val, lim in (("maximum degree 4 base < 1.2037", b4, 1.2037), ("general base < 1.2341", b, 1.2341)):
    print("   %-40s %s" % (lab, "ok" if val < lim else "FAIL"))
    if not val < lim:
        fails.append(lab)

# ------------------------------------------------------------------ 2. the potential psi (Lemma C.11, `lem:psi-2223`)
print("2. Potential psi in section 1 (lem:psi-2223) and the psi-adjusted rows")
w2, w3, w4 = secw[0]
dw3, dw4 = w3 - w2, w4 - w3
e = dw3 - dw4
print("   psi = %s, e = (w3-w2)-(w4-w3) = %s = %.7f" % (psi, e, float(e)))
excess = {
    "k=1, z remains: dw3 - dw4": dw3 - dw4,
    "k=1, z deleted: w3 - dw4": w3 - dw4,
    "k=0: dw3": dw3,
    "k=2, z'' uncharged: w3 + dw3 - 2 dw4": w3 + dw3 - 2 * dw4,
    "k=2, z'' third far end of degree 4: w3 + w4 - w2 - 3 dw4": w3 + w4 - w2 - 3 * dw4,
    "z'' = u, first branch: 2 w3 - 2 dw4": 2 * w3 - 2 * dw4,
    "z'' = u, second branch: w3 + dw3 - 2 dw4": w3 + dw3 - 2 * dw4,
    "u a far end, z = q, k=0, first branch: w3 - dw4": w3 - dw4,
}
for name, v in excess.items():
    ok = v > psi
    print("   excess %-58s = %.7f > psi: %s" % (name, float(v), ok))
    if not ok:
        fails.append(name)
print("   minimum excess = %.7f (= e)" % float(min(excess.values())))
closed = {
    "closed part, b >= K_m >= 391: 387 w2 - 3 dw4": 387 * w2 - 3 * dw4,
    "closed part, b = K >= 65, two vertices of degree 3: 2 w3 + 60 w2 - 3 dw4": 2 * w3 + 60 * w2 - 3 * dw4,
}
for name, v in closed.items():
    ok = v > e
    print("   excess %-72s = %.7f > e: %s" % (name, float(v), ok))
    if not ok:
        fails.append(name)
# hypothesis of the hybrid rule (lem:isg-one-section): 2 dw_D / w2 <= K + 1 - Delta for D = d(v) >= 4
ratios = [("section %d, Delta = D = 4" % (k + 1), 2 * (w[2] - w[1]) / w[0]) for k, w in enumerate(secw)]
ratios += [("degrees 5-6 weights, Delta = %d, D = %d" % (Dl, D), 2 * (W[D] - W[D - 1]) / W[2])
           for Dl in (5, 6) for D in range(4, Dl + 1)]
rname, rmax = max(ratios, key=lambda t: t[1])
print("   largest 2 dw_D / w2 over the section weights and the degree 5-6 weights: %.4f (%s) < 62: %s"
      % (float(rmax), rname, rmax < 62))
if not rmax < 62:
    fails.append("2 dw_D / w2 < 62")
ok = alphas[1] / 3 - psi > 0
print("   alpha_2/3 - psi = %.7f > 0: %s" % (float(alphas[1] / 3 - psi), ok))
if not ok:
    fails.append("alpha2/3 > psi")


def bs(a, b_):
    return iv.mpf(2) ** (-ivq(a)) + iv.mpf(2) ** (-ivq(b_))


names1 = ["(2,2,2,2)(a)", "(2,2,2,2)(b)", "(2,2,2,3)", "(2,2,2,4)", "(2,2,3,3)", "(2,2,3,4)", "(2,2,4,4)"]
V = deg4_vectors(w2, w3, w4, psi)
worst_other = 0
for nm, (a, b_) in zip(names1, V[0]):
    if nm.startswith("(2,2,2,2)"):
        s = bs(a, b_)
        print("   %-14s as printed (psi in the vector): sum <= %.7f" % (nm, float(s.b)))
        ok = s.b <= 1
    elif nm == "(2,2,2,3)":
        sums = [bs(a + e - psi, b_), bs(a, b_ + e - psi), bs(a + e - psi, b_ + e - psi)]
        print("   (2,2,2,3) with e - psi added in the first, second, both branches: sums <= %s"
              % [round(float(x.b), 5) for x in sums])
        s = bs(a - psi, b_ - psi)
        print("   (2,2,2,3) with -psi in both branches and no excess: sum >= %.5f (> 1, so the lemma is needed)"
              % float(s.a))
        ok = all(x.b <= 1 for x in sums)
    else:
        s = bs(a - psi, b_ - psi)
        worst_other = max(worst_other, float(s.b))
        ok = s.b <= 1
    if not ok:
        fails.append(nm)
print("   other section-1 rows with -psi in both branches: largest sum <= %.5f" % worst_other)
worst = 0
for k, sec in enumerate(deg4_vectors(w2, w3, w4, 0)[1:], start=2):
    for (a, b_) in sec:
        s = bs(a - psi, b_ - psi)
        worst = max(worst, float(s.b))
        if not s.b <= 1:
            fails.append("later section %d under section-1 weights" % k)
print("   later-section rows under section-1 weights with -psi in both branches: largest sum <= %.5f" % worst)
print("ALL CHECKS PASSED" if not fails else "FAILURES: %s" % fails)