import json
import numpy as np
from toposcan import gf2_rank, betti_graph, interlevel_betti, grid_tokens


def wedge_triangles(q, level=0.35):
    """q triangles sharing vertex 0; each contributes one independent cycle."""
    edges = []
    h = [0.0]
    for j in range(q):
        u, v = 1 + 2*j, 2 + 2*j
        h.extend([level, level])
        edges.extend([(0, u), (u, v), (v, 0)])
    return len(h), edges, np.asarray(h)


def direct_formula_check(rng):
    # Verify beta0=n-rank(B0)-rank(B1), beta1=n1-rank(B1) on random graphs.
    rows = []
    for n in [4, 7, 10]:
        for _ in range(20):
            edges = [(u, v) for u in range(n) for v in range(u+1, n)
                     if rng.random() < 0.28]
            b0, b1 = betti_graph(n, edges)
            B1 = np.zeros((n, len(edges)), dtype=np.uint8)
            for j, (u, v) in enumerate(edges):
                B1[u, j] = B1[v, j] = 1
            rank = gf2_rank(B1)
            predicted = (n-rank, len(edges)-rank)
            rows.append((b0, b1) == predicted)
    return {"cases": len(rows), "passed": int(sum(rows))}


def main():
    rng = np.random.default_rng(7)
    results = {"formula_check": direct_formula_check(rng)}

    # Prediction 1: each added independent cycle adds exactly one to beta1.
    scaling = []
    for q in range(1, 7):
        n, edges, h = wedge_triangles(q)
        (b0, b1), _ = interlevel_betti(n, edges, h, 0.0, 0.36)
        scaling.append({"q": q, "observed_beta1": b1, "predicted_beta1": q})
    results["cycle_scaling"] = scaling

    # Prediction 2: beta1 turns on exactly when the upper interval reaches
    # the common upper-star edge value .35.
    n, edges, h = wedge_triangles(3)
    transition = []
    for hi in [0.20, 0.34, 0.35, 0.36, 0.50]:
        (b0, b1), (nv, ne) = interlevel_betti(n, edges, h, 0.0, hi)
        transition.append({"hi": hi, "observed_beta1": b1,
                           "predicted_beta1": 0 if hi < 0.35 else 3,
                           "vertices": nv, "edges": ne})
    results["window_transition"] = transition

    # Prediction 3: with fixed hi=.50, perturbations below the .15 margin
    # cannot change token beta1; once the margin is crossed, positive shifts
    # remove all cycle edges from this interval.
    robustness = []
    for eps in [0.00, 0.05, 0.14, 0.15, 0.16, 0.25]:
        hp = h + eps
        (b0, b1), _ = interlevel_betti(n, edges, hp, 0.0, 0.50)
        robustness.append({"epsilon": eps, "observed_beta1": b1,
                           "predicted_beta1": 3 if eps < 0.15 else 0})
    results["perturbation_margin"] = robustness

    # Exercise the shared-grid token representation and report its shape.
    tokens = grid_tokens(n, edges, h, np.linspace(0, 1, 11), m=4, stride=2)
    results["token_example"] = {"shape": list(tokens.shape),
                                 "tokens": tokens.tolist()}

    # Aggregate exact prediction success, excluding floating boundary cases
    # where representation of .35/.15 can decide inclusion.
    scaling_ok = all(x["observed_beta1"] == x["predicted_beta1"] for x in scaling)
    transition_ok = all(x["observed_beta1"] == x["predicted_beta1"]
                        for x in transition)
    robust_ok = all(x["observed_beta1"] == x["predicted_beta1"]
                    for x in robustness if x["epsilon"] != 0.15)
    results["prediction_summary"] = {
        "scaling_exact": scaling_ok,
        "transition_exact": transition_ok,
        "submargin_robust_exact": robust_ok,
        "boundary_case_epsilon_0.15_observed": robustness[3]["observed_beta1"],
        "boundary_tolerance_note": "At equality, closed interval inclusion is expected; epsilon=.15 is reported separately."
    }
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()