import json, math, time, random
from pathlib import Path
import numpy as np

try:
    import torch
    import torch.nn.functional as F
except Exception:
    torch = None


def cyclic_design(v, base):
    blocks = [sorted({(x + t) % v for x in base}) for t in range(v)]
    return np.asarray(blocks, dtype=np.int64)


def make_designs():
    # Cyclic (7,3,1) and (13,4,1) projective-plane designs.
    return {
        "Fano_7": cyclic_design(7, [0, 1, 3]),
        "projective_13": cyclic_design(13, [0, 1, 3, 9]),
    }


def incidence(blocks, v):
    M = np.zeros((v, len(blocks)), dtype=np.int64)
    for j, b in enumerate(blocks):
        M[b, j] = 1
    return M


def math_check(name, blocks):
    v = int(blocks.max()) + 1
    M = incidence(blocks, v)
    G = M @ M.T
    r = int(M.sum(1)[0])
    off = G[~np.eye(v, dtype=bool)]
    lam = int(off[0])
    pair_vals = off.tolist()
    # Z(x,y) is the column-vector membership difference.
    norm_err = []
    for x in range(v):
        for y in range(x + 1, v):
            z = M[x] - M[y]
            norm_err.append(int(z @ z))
    return {
        "name": name, "v": v, "b": len(blocks), "k": len(blocks[0]),
        "r": r, "lambda": lam, "row_min": int(M.sum(1).min()),
        "row_max": int(M.sum(1).max()), "pair_min": int(min(pair_vals)),
        "pair_max": int(max(pair_vals)), "pair_variance": float(np.var(pair_vals)),
        "mmT_identity_maxerr": int(np.max(np.abs(G - ((r-lam)*np.eye(v, dtype=int)+lam*np.ones((v,v), dtype=int))))),
        "predicted_z_norm_sq": 2*(r-lam), "observed_z_norm_sq_min": min(norm_err),
        "observed_z_norm_sq_max": max(norm_err), "z_norm_sq_mean": float(np.mean(norm_err)),
    }


def random_blocks(v, k, b, seed=0):
    rng = np.random.default_rng(seed)
    return np.asarray([rng.choice(v, k, replace=False) for _ in range(b)], dtype=np.int64)


def contiguous_blocks(v, k, b):
    return np.asarray([[(i + t) % v for t in range(k)] for i in range(b)], dtype=np.int64)


def torch_attention(x, blocks, heads=2):
    # One-head equivalent batched block attention, with scatter-and-average.
    n, d = x.shape
    q, k, val = x, x, x
    y = torch.zeros_like(x)
    counts = torch.zeros(n, device=x.device)
    scale = math.sqrt(d)
    for idx_np in blocks:
        idx = torch.as_tensor(idx_np, device=x.device, dtype=torch.long)
        s = (q[idx] @ k[idx].T) / scale
        a = F.softmax(s, dim=-1)
        y.index_add_(0, idx, a @ val[idx])
        counts.index_add_(0, idx, torch.ones(len(idx), device=x.device))
    return y / counts[:, None]


def dense_attention(x):
    return F.softmax(x @ x.T / math.sqrt(x.shape[1]), dim=-1) @ x


def benchmark(seed=123):
    random.seed(seed); np.random.seed(seed)
    if torch is None:
        return {"error": "torch unavailable"}
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        torch.manual_seed(seed)
        if device == "cuda": torch.cuda.manual_seed_all(seed)
        # Genuine balanced designs are benchmarked on their complete point sets.
        # This avoids incorrectly claiming that disjoint copies preserve global coverage.
        cases_by_n = {
            7: {"dense": None, "contiguous": contiguous_blocks(7, 3, 7),
                "random": random_blocks(7, 3, 7, seed), "design": make_designs()["Fano_7"]},
            13: {"dense": None, "contiguous": contiguous_blocks(13, 4, 13),
                 "random": random_blocks(13, 4, 13, seed), "design": make_designs()["projective_13"]},
        }
        out = {"device": device, "results": {}}
        for n, cases in cases_by_n.items():
            d = 32; k = len(cases["design"][0]); x = torch.randn(n, d, device=device)
            out["results"][str(n)] = {}
            for name, blocks in cases.items():
                if name == "dense":
                    fn = lambda: dense_attention(x)
                    flops = n*n*d*2; pair_var = 0.0; row_var = 0.0
                else:
                    fn = lambda blocks=blocks: torch_attention(x, blocks)
                    flops = len(blocks)*k*k*d*2
                    M = incidence(blocks, n); co = M @ M.T
                    pair_var = float(np.var(co[~np.eye(n, dtype=bool)]))
                    row_var = float(np.var(M.sum(1)))
                fn()
                if device == "cuda": torch.cuda.synchronize()
                ts=[]
                for _ in range(5):
                    t=time.perf_counter(); fn()
                    if device == "cuda": torch.cuda.synchronize()
                    ts.append(time.perf_counter()-t)
                out["results"][str(n)][name] = {
                    "median_ms": 1000*float(np.median(ts)),
                    "attention_score_flops": int(flops),
                    "pair_coverage_variance": pair_var, "row_degree_variance": row_var}
        return out
    except Exception as e:
        return {"error": repr(e)}


def main():
    designs = make_designs()
    checks = [math_check(name, blocks) for name, blocks in designs.items()]
    # Prediction sweep: for every exact design, all pair norms must equal 2(r-lambda).
    # A second prediction is invariance under scaling the incidence matrix by s copies:
    # r, lambda, and ||Z||^2 scale by s, while pair variance remains zero.
    base = designs["Fano_7"]
    sweep=[]
    for copies in [1,2,4,8]:
        b=np.tile(base, (copies,1))
        c=math_check(f"Fano copies={copies}", b)
        c["predicted_r"] = 3*copies
        c["predicted_lambda"] = 1*copies
        c["predicted_z_norm_sq"] = 2*(3-1)*copies
        sweep.append(c)
    report={"math_checks": checks, "scaling_sweep": sweep, "benchmark": benchmark()}
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()