"""Toy verification of the cross-degree anti-oscillation certificate.

Dynamics: x[t+1] = 1{W x[t] >= R}.  For a bipartite graph, initializing one
side to one and the other to zero gives an exact period-two orbit whenever
minimum cross-gain >= R (with no within-side edges).  This script checks the
boundary, scaling, and threshold suppression predictions, then compares a
baseline threshold choice with a certificate-aware choice.
"""
import itertools
import json
import numpy as np

SEED = 1629


def balanced_partitions(n):
    # Fix vertex 0 in side 0 to remove complement duplication.
    for bits in itertools.product((0, 1), repeat=n - 1):
        p = np.array((0,) + bits, dtype=np.int8)
        if 0 < p.sum() < n:
            yield p


def min_cross_gain(W, part):
    vals = []
    for i in range(len(part)):
        opposite = part != part[i]
        vals.append(float(np.abs(W[i, opposite]).sum()))
    return min(vals)


def exact_certificate(W):
    vals = [(min_cross_gain(W, p), p.copy()) for p in balanced_partitions(len(W))]
    return max(vals, key=lambda z: z[0])


def threshold_step(W, x, R):
    return (W @ x >= R).astype(np.int8)


def orbit(W, x0, R, steps=12):
    x = np.asarray(x0, dtype=np.int8).copy()
    out = [x.copy()]
    for _ in range(steps):
        x = threshold_step(W, x, R)
        out.append(x.copy())
    return np.asarray(out)


def is_period_two_from(out, start=0):
    return (len(out) >= start + 3 and
            np.array_equal(out[start], out[start + 2]) and
            not np.array_equal(out[start], out[start + 1]))


def make_bipartite(left, right, p, weight=1.0, seed=0):
    rng = np.random.default_rng(seed)
    n = left + right
    W = np.zeros((n, n), dtype=float)
    for i in range(left):
        for j in range(left, n):
            if rng.random() < p:
                W[i, j] = W[j, i] = weight
    # Ensure every vertex has at least one cross edge for useful sweeps.
    for i in range(n):
        if not np.any(W[i]):
            j = rng.integers(left, n) if i < left else rng.integers(0, left)
            W[i, j] = W[j, i] = weight
    return W


def boundary_sweep(W, x0, cert, thresholds):
    return [int(is_period_two_from(orbit(W, x0, R), 0)) for R in thresholds]


def main():
    # Prediction 1: exact transition is R <= c_*; test every integer R.
    W = make_bipartite(4, 4, p=0.72, weight=1.0, seed=1)
    c_star, part = exact_certificate(W)
    x0 = (part == 1).astype(np.int8)
    thresholds = list(range(0, int(c_star) + 3))
    observed = boundary_sweep(W, x0, c_star, thresholds)
    predicted = [int(1 <= R <= c_star) for R in thresholds]
    boundary_ok = observed == predicted

    # Prediction 2: scaling all weights scales c_* and the boundary linearly.
    scales = [0.5, 1.0, 1.5, 2.0]
    scaling_rows = []
    for s in scales:
        Ws = s * W
        cs, ps = exact_certificate(Ws)
        xs = (ps == 1).astype(np.int8)
        # Probe integer and half-integer thresholds around the predicted c*.
        Rs = np.arange(max(0.0, cs - 1.0), cs + 1.01, 0.25)
        flags = [is_period_two_from(orbit(Ws, xs, float(R)), 0) for R in Rs]
        max_observed = max((float(R) for R, ok in zip(Rs, flags) if ok), default=-1.0)
        scaling_rows.append({"scale": s, "predicted_c_star": cs,
                             "observed_largest_period2_R": max_observed,
                             "boundary_error": abs(max_observed - cs)})
    scaling_ok = all(row["boundary_error"] <= 0.26 for row in scaling_rows)

    # Prediction 3: choosing R 10% above the certificate blocks this certified
    # orbit, while choosing R at or below it preserves it.
    suppress_rows = []
    for s in [0.5, 1.0, 1.5, 2.0]:
        Ws = s * W
        cs, ps = exact_certificate(Ws)
        xs = (ps == 1).astype(np.int8)
        at = is_period_two_from(orbit(Ws, xs, cs), 0)
        above = is_period_two_from(orbit(Ws, xs, 1.1 * cs), 0)
        suppress_rows.append({"scale": s, "c_star": cs,
                              "at_certificate": bool(at),
                              "at_110_percent": bool(above)})
    suppression_ok = all(r["at_certificate"] and not r["at_110_percent"] for r in suppress_rows)

    # Baseline vs idea on random initial states. Baseline uses R=1, while
    # certificate-aware selection uses R=1.1*c_* (minimum integer-like safe
    # threshold is represented directly here because the toy permits reals).
    rng = np.random.default_rng(SEED)
    n_trials = 300
    base_flags, idea_flags = [], []
    safe_R = 1.1 * c_star
    for _ in range(n_trials):
        z = rng.integers(0, 2, size=len(W), dtype=np.int8)
        base_flags.append(is_period_two_from(orbit(W, z, 1.0), 2))
        idea_flags.append(is_period_two_from(orbit(W, z, safe_R), 2))
    baseline_rate = float(np.mean(base_flags))
    idea_rate = float(np.mean(idea_flags))

    result = {
        "seed": SEED,
        "graph_nodes": len(W),
        "edges": int(np.count_nonzero(W) // 2),
        "exact_c_star": c_star,
        "certificate_partition": part.tolist(),
        "prediction_1_boundary": {"thresholds": thresholds, "observed": observed,
                                  "predicted": predicted, "confirmed": boundary_ok},
        "prediction_2_scaling": scaling_rows,
        "prediction_2_confirmed": scaling_ok,
        "prediction_3_suppression": suppress_rows,
        "prediction_3_confirmed": suppression_ok,
        "baseline_vs_idea": {"baseline_R": 1.0, "idea_R": safe_R,
                              "trials": n_trials,
                              "period2_rate_baseline": baseline_rate,
                              "period2_rate_idea": idea_rate},
        "all_mechanism_predictions_confirmed": bool(boundary_ok and scaling_ok and suppression_ok)
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()