import json
import math
from collections import deque
import numpy as np

SEED = 376


def softmax(x):
    z = x - np.max(x, axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)


def kappa(p, eps=0.0):
    return (p > eps).astype(np.int8)


def verify_support_order(rng, trials=5000):
    # Finite Boolean-support version of q <= p iff kappa(q) <= kappa(p).
    for _ in range(trials):
        p = rng.random(8)
        q = rng.random(8)
        # Make arbitrary nonnegative distributions, including zeros.
        p[rng.random(8) < .35] = 0
        q[rng.random(8) < .35] = 0
        p /= p.sum() if p.sum() else 1
        q /= q.sum() if q.sum() else 1
        lhs = np.all((q > 0) <= (p > 0))
        rhs = np.all(kappa(q) <= kappa(p))
        if lhs != rhs:
            return False, _
    return True, trials


def windows(n, width=5, stride=3):
    out = []
    start = 0
    while start < n:
        out.append(np.arange(start, min(n, start + width)))
        if start + width >= n:
            break
        start += stride
    return out


def independent_mask(prob, eps):
    return prob > eps


def continuation_violations(mask, wins):
    # A local-surjectivity analogue: in every shared query of neighboring
    # windows, both retained fibers must share at least one retained key.
    bad = 0
    pairs = 0
    for a, b in zip(wins[:-1], wins[1:]):
        shared_q = np.intersect1d(a, b)
        shared_k = np.intersect1d(a, b)
        for q in shared_q:
            pairs += 1
            if not np.any(mask[q, shared_k]):
                bad += 1
    return bad, pairs


def row_empty(mask):
    return int(np.sum(mask.sum(axis=1) == 0))


def repair_continuations(prob, mask, wins):
    """Add one maximum-weight common-key witness per violated overlap fiber."""
    mask = mask.copy()
    for a, b in zip(wins[:-1], wins[1:]):
        shared_q = np.intersect1d(a, b)
        shared_k = np.intersect1d(a, b)
        for q in shared_q:
            if not np.any(mask[q, shared_k]):
                k = shared_k[np.argmax(prob[q, shared_k])]
                mask[q, k] = True
    return mask


def structured_clean(prob, initial, wins):
    mask = initial.copy()
    n = mask.shape[0]
    # Remove low-weight edges only when all fibers and overlap continuations
    # remain nonempty. This is the finite mask certificate used by the MVP.
    candidates = [(float(prob[q, k]), q, k) for q in range(n) for k in range(n)
                  if mask[q, k]]
    candidates.sort()
    for _, q, k in candidates:
        if mask[q].sum() <= 1:
            continue
        mask[q, k] = False
        bad, _ = continuation_violations(mask, wins)
        if bad or row_empty(mask):
            mask[q, k] = True
    return mask


def window_graph(mask, wins):
    g = {i: set() for i in range(len(wins))}
    for i in range(len(wins)):
        for j in range(i + 1, len(wins)):
            shared_q = np.intersect1d(wins[i], wins[j])
            shared_k = np.intersect1d(wins[i], wins[j])
            # A bijective restriction in this row-wise toy means each shared
            # query has exactly one common surviving key.
            bij = bool(len(shared_q) and all(np.sum(mask[q, shared_k]) == 1
                                             for q in shared_q))
            if bij:
                g[i].add(j); g[j].add(i)
    return g


def components(g):
    seen, cs = set(), []
    for s in g:
        if s in seen: continue
        c, todo = [], [s]; seen.add(s)
        while todo:
            u = todo.pop(); c.append(u)
            for v in g[u]:
                if v not in seen: seen.add(v); todo.append(v)
        cs.append(c)
    return cs


def attention_error(logits, values, mask, dense):
    masked = np.where(mask, logits, -1e9)
    out = softmax(masked) @ values
    return float(np.mean((out - dense) ** 2))


def one_trial(seed, n=12, d=8):
    rng = np.random.default_rng(seed)
    wins = windows(n, width=6, stride=3)
    q = rng.normal(size=(n, d)); k = rng.normal(size=(n, d)); v = rng.normal(size=(n, d))
    raw = q @ k.T / math.sqrt(d)
    local = np.zeros((n, n), dtype=bool)
    for w in wins: local[np.ix_(w, w)] = True
    logits = np.where(local, raw, -1e9)
    prob = softmax(logits)
    eps = 0.08
    base = prob > eps
    idea = structured_clean(prob, repair_continuations(prob, base, wins), wins)
    dense = prob @ v
    bad0, pairs = continuation_violations(base, wins)
    bad1, _ = continuation_violations(idea, wins)
    return {
        "dense_edges": int(local.sum()), "dense_error": 0.0,
        "independent_edges": int(base.sum()), "independent_error": attention_error(logits, v, base, dense),
        "independent_violations": bad0,
        "structured_edges": int(idea.sum()), "structured_error": attention_error(logits, v, idea, dense),
        "structured_violations": bad1, "checks": pairs,
        "independent_components": len(components(window_graph(base, wins))),
        "structured_components": len(components(window_graph(idea, wins)))
    }


def run():
    rng = np.random.default_rng(SEED)
    support_ok, support_trials = verify_support_order(rng)
    trials = [one_trial(SEED + i) for i in range(10)]
    keys = trials[0].keys()
    summary = {}
    for key in keys:
        vals = [x[key] for x in trials]
        summary[key] = {"mean": float(np.mean(vals)), "std": float(np.std(vals)), "values": vals}
    result = {
        "seed": SEED, "epsilon": 0.08, "math_support_order": {"passed": support_ok, "trials": support_trials},
        "trials": 10, "summary": summary
    }
    with open("results.json", "w") as f: json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    run()