import json
import math
import random
from itertools import combinations
import numpy as np
from scipy.optimize import minimize_scalar


def mu(G, q):
    q = np.asarray(q, dtype=float)
    d = np.exp(q - np.max(q))
    return float(np.max((G @ d) / d))


def rho(G):
    return float(np.max(np.abs(np.linalg.eigvals(G))))


def optimize_q(G):
    # For this sanity check, optimize the exact nonsmooth certificate in the
    # only meaningful relative coordinate for two nodes.
    if G.shape == (2, 2):
        f = lambda x: mu(G, [x, -x])
        r = minimize_scalar(f, bounds=(-10, 10), method="bounded", options={"xatol": 1e-12})
        return np.array([r.x, -r.x]), float(r.fun)
    q = np.zeros(G.shape[0])
    return q, mu(G, q)


def dominant_cycle_matrix(c, n=6, reverse_fraction=0.2):
    G = np.zeros((n, n), dtype=float)
    # Dominant directed 3-cycle has product c^3. Reciprocal weak links make
    # the proposed symmetric pair score able to discover its node pairs.
    G[1, 0] = c * 4.0
    G[2, 1] = c
    G[0, 2] = c / 4.0
    G[0, 1] = reverse_fraction * c * 4.0
    G[1, 2] = reverse_fraction * c
    G[2, 0] = reverse_fraction * c / 4.0
    for i in range(3, n):
        G[i, i - 1] = 0.08
        G[i - 1, i] = 0.05
    return G


def edge_score(G, i, j, delta=1e-8):
    return math.log(G[i, j] + delta) + math.log(G[j, i] + delta)


def greedy_partition(G, max_size=3):
    clusters = [{i} for i in range(G.shape[0])]
    while True:
        candidates = []
        for a, b in combinations(range(len(clusters)), 2):
            if len(clusters[a] | clusters[b]) <= max_size:
                score = max(edge_score(G, i, j) for i in clusters[a] for j in clusters[b])
                candidates.append((score, a, b))
        if not candidates:
            break
        score, a, b = max(candidates)
        if score < math.log(0.08 * 0.05):
            break
        merged = clusters[a] | clusters[b]
        clusters = [x for k, x in enumerate(clusters) if k not in (a, b)] + [merged]
    return sorted([sorted(x) for x in clusters], key=lambda x: x[0])


def cycle_internal(P, cycle=(0, 1, 2)):
    owner = {i: k for k, S in enumerate(P) for i in S}
    return len({owner[i] for i in cycle}) == 1


def visible_cycle_gain(G, P, cycle=(0, 1, 2)):
    return 0.0 if cycle_internal(P, cycle) else float(np.prod([G[cycle[(k + 1) % 3], cycle[k]] for k in range(3)]))


def main():
    np.random.seed(7)
    random.seed(7)
    out = {"normalization_sweep": [], "imbalance_sweep": [], "cluster_sweep": []}

    # Prediction 1: diagonal scaling cannot beat the cycle boundary c=1.
    for c in [0.25, 0.5, 0.9, 1.0, 1.1, 1.5, 2.0]:
        G = np.array([[0., 4*c], [c/4, 0.]])
        q, best = optimize_q(G)
        out["normalization_sweep"].append({"c": c, "rho": rho(G), "mu_q0": mu(G, [0, 0]),
            "mu_optimized": best, "q": q.tolist(), "predicted_optimum": c,
            "stable_by_certificate": best < 1.0})

    # Prediction 2: unscaled certificate depends on imbalance, optimized one does not.
    for a in [1., 2., 4., 8.]:
        c = 0.7
        G = np.array([[0., a*c], [c/a, 0.]])
        _, best = optimize_q(G)
        out["imbalance_sweep"].append({"imbalance": a, "mu_unscaled": mu(G, [0, 0]),
            "mu_optimized": best, "predicted_optimized": c})

    # Prediction 3: greedy gain-weighted merging internalizes the cycle when
    # cluster capacity is 3; the fixed random partition leaves it exposed.
    singleton_P = [[0], [1], [2], [3], [4], [5]]
    random_P = [[0, 2], [1], [3, 4, 5]]
    for c in [0.5, 0.9, 1.0, 1.1, 1.5]:
        G = dominant_cycle_matrix(c)
        greedy = greedy_partition(G, max_size=3)
        out["cluster_sweep"].append({"c": c, "rho": rho(G), "greedy_partition": greedy,
            "singleton_internal": cycle_internal(singleton_P),
            "singleton_boundary_cycle_gain": visible_cycle_gain(G, singleton_P),
            "greedy_internal": cycle_internal(greedy),
            "greedy_boundary_cycle_gain": visible_cycle_gain(G, greedy),
            "random_internal": cycle_internal(random_P),
            "random_boundary_cycle_gain": visible_cycle_gain(G, random_P),
            "cycle_product_predicted": c**3})

    norm_err = max(abs(x["mu_optimized"] - x["predicted_optimum"]) for x in out["normalization_sweep"])
    imb_err = max(abs(x["mu_optimized"] - x["predicted_optimized"]) for x in out["imbalance_sweep"])
    boundary_ok = all((x["c"] < 1) == x["stable_by_certificate"] for x in out["normalization_sweep"] if x["c"] != 1)
    boundary_at_one = min(abs(x["mu_optimized"] - 1) for x in out["normalization_sweep"] if x["c"] == 1) < 1e-8
    transition_ok = all(x["greedy_internal"] and not x["singleton_internal"] and not x["random_internal"] and
        x["greedy_boundary_cycle_gain"] == 0.0 and
        abs(x["singleton_boundary_cycle_gain"] - x["cycle_product_predicted"]) < 1e-10 and
        abs(x["random_boundary_cycle_gain"] - x["cycle_product_predicted"]) < 1e-10
        for x in out["cluster_sweep"])
    out["predictions"] = {
        "boundary_predicted_c": 1.0, "boundary_confirmed": boundary_ok and boundary_at_one,
        "normalization_max_abs_error": norm_err,
        "imbalance_max_abs_error": imb_err,
        "cluster_transition_confirmed": transition_ok,
        "mechanism_scope": "clustering removes the dominant channel from the boundary graph; it does not reduce full spectral radius"
    }
    print(json.dumps(out, indent=2, sort_keys=True))

if __name__ == "__main__":
    main()