"""Sanity checks for 'Fast Attention after SETH: The Exact Exponent and the Hypotheses Behind It'.

Every check prints PASS/FAIL. Exact arithmetic (fractions) is used where a statement is algebraic;
the sampling check is statistical and reports the observed worst error against the proved bound.
Run:  python -I checks.py
"""
import math
import random
from fractions import Fraction

random.seed(20261030)


def ip(a, b):
    return sum(x * y for x, y in zip(a, b))


def rand_vec(m, p=0.5):
    return [1 if random.random() < p else 0 for _ in range(m)]


# ---------------------------------------------------------------------------------------------
# 1. Generalized parity identity (Theorem 4.1): with z = 1 + 1/r and S = sum_j z^{<a,b_j>},
#    r^m S is an integer and r^m S = r^m N_0 (mod r+1), N_0 = #{j : <a,b_j> = 0}.
# ---------------------------------------------------------------------------------------------
def check_mod_identity(trials=3000):
    bad = 0
    for _ in range(trials):
        m = random.randint(1, 9)
        n = random.randint(1, 14)
        r = random.choice([1, 2, 3, 4, 6, 10])
        a = rand_vec(m)
        B = [rand_vec(m, random.choice([0.2, 0.5])) for _ in range(n)]
        z = Fraction(r + 1, r)
        S = sum(z ** ip(a, b) for b in B)
        X = S * r ** m
        if X.denominator != 1:
            bad += 1
            continue
        N0 = sum(1 for b in B if ip(a, b) == 0)
        if (X.numerator - (r ** m) * N0) % (r + 1) != 0:
            bad += 1
    print(f"[1] mod-(r+1) identity on {trials} random instances: {'PASS' if bad == 0 else 'FAIL'} (bad={bad})")
    return bad == 0


# ---------------------------------------------------------------------------------------------
# 2. Reconstruction (Theorem 4.1): from any o_hat with |o_hat - o| <= eps_0,
#    where o = n/(n+S) and eps_0 = r^m / (32 n (r+1)^{2m}), the value round(r^m * n(1/o_hat - 1))
#    equals r^m S.  We test the two extreme perturbations exactly.
# ---------------------------------------------------------------------------------------------
def check_reconstruction(trials=1500):
    bad = 0
    for _ in range(trials):
        m = random.randint(1, 8)
        n = random.randint(1, 12)
        r = random.choice([1, 2, 3, 5])
        a = rand_vec(m)
        B = [rand_vec(m) for _ in range(n)]
        z = Fraction(r + 1, r)
        S = sum(z ** ip(a, b) for b in B)  # n real/filler keys
        o = Fraction(n) / (n + S)
        eps0 = Fraction(r ** m, 32 * n * (r + 1) ** (2 * m))
        err = eps0
        for sgn in (-1, 1):
            oh = o + sgn * err
            Sh = n * (1 / oh - 1)
            X = Sh * r ** m
            rec = math.floor(X + Fraction(1, 2))
            if rec != S * r ** m:
                bad += 1
    print(f"[2] exact reconstruction of r^m S from perturbed outputs: {'PASS' if bad == 0 else 'FAIL'} (bad={bad})")
    return bad == 0


# ---------------------------------------------------------------------------------------------
# 3. Isolation (Theorem 4.1): if c >= 1 elements are kept independently with prob. 1/2, the kept
#    count is nonzero mod (r+1) with probability >= 1/2.  Exact binomial computation.
# ---------------------------------------------------------------------------------------------
def check_isolation():
    worst = 1.0
    for r in range(1, 12):
        for c in range(1, 60):
            p = sum(math.comb(c, t) for t in range(c + 1) if t % (r + 1) != 0) / 2 ** c
            worst = min(worst, p)
    ok = worst >= 0.5 - 1e-12
    print(f"[3] isolation probability >= 1/2 (worst {worst:.6f}): {'PASS' if ok else 'FAIL'}")
    return ok


# ---------------------------------------------------------------------------------------------
# helpers: exact-ish attention output for one query via log-sum-exp in float (logits are moderate)
# ---------------------------------------------------------------------------------------------
def softmax_out(logits, values):
    mx = max(logits)
    w = [math.exp(l - mx) for l in logits]
    s = sum(w)
    return sum(wi * vi for wi, vi in zip(w, values)) / s


# ---------------------------------------------------------------------------------------------
# 4. Gap gadget (Theorem 3.3): Max-IP gap instances with all inner products in [0, T], YES pair at T,
#    NO pairs at most T/tau.  eta = (1-2 eps)/4, xi = ceil(10 ln((1-eta)/eta))/10, and
#    lambda T (1-1/tau) = ceil(10 (ln n + 2 ln(1/eta) + 1/10))/10, dummy logit lambda*T - xi.
#    Expect: YES row output >= 1-eta, every NO output <= eta, spread <= (ln n + 2 ln(1/eta) + 1/5)/(1-1/tau).
# ---------------------------------------------------------------------------------------------
def check_gap_gadget(trials=400):
    bad = 0
    for _ in range(trials):
        n = random.choice([10, 50, 200, 1000])
        T = random.choice([8, 16, 32, 64])
        tau = random.choice([2, 4, 8, 16])
        eps = random.choice([1e-6, 1 / 6, 1 / 4, 0.45])
        eta = (1 - 2 * eps) / 4
        xi = math.ceil(10 * math.log((1 - eta) / eta)) / 10
        lam = math.ceil(10 * (math.log(n) + 2 * math.log(1 / eta) + 0.1)) / (10 * T * (1 - 1 / tau))
        dummy = lam * T - xi
        yes = random.random() < 0.5
        # adversarial NO-type inner products: as large as allowed
        ips = [math.floor(T / tau) for _ in range(n)]
        if yes:
            ips[random.randrange(n)] = T
        logits = [lam * s for s in ips] + [dummy]
        vals = [1.0] * n + [0.0]
        o = softmax_out(logits, vals)
        spread = max(logits) - min(min(logits), 0.0)
        bound = (math.log(n) + 2 * math.log(1 / eta) + 0.2) / (1 - 1 / tau)
        if yes and o < 1 - eta:
            bad += 1
        if (not yes) and o > eta:
            bad += 1
        if lam * T > bound + 1e-12 or dummy < 0:
            bad += 1
    print(f"[4] gap gadget thresholds 1-eta vs eta and spread bound: {'PASS' if bad == 0 else 'FAIL'} (bad={bad})")
    return bad == 0


# ---------------------------------------------------------------------------------------------
# 5. Exact-OV gadget (Proposition 5.1): logits -lambda<a,b>, dummy -11/10, lambda = ceil(10 ln(14n))/10.
#    Expect YES row >= 3/4, NO rows <= 1/4, spread <= D (ln(14n)+1/10).
# ---------------------------------------------------------------------------------------------
def check_exact_ov_gadget(trials=300):
    bad = 0
    for _ in range(trials):
        n = random.choice([8, 32, 128])
        D = random.choice([6, 10, 16])
        A = [rand_vec(D, 0.5) for _ in range(n)]
        B = [rand_vec(D, 0.5) for _ in range(n)]
        lam = math.ceil(10 * math.log(14 * n)) / 10
        exists = any(ip(a, b) == 0 for a in A for b in B)
        outs = []
        for a in A:
            logits = [-lam * ip(a, b) for b in B] + [-1.1]
            outs.append(softmax_out(logits, [1.0] * n + [0.0]))
            if lam * D > D * (math.log(14 * n) + 0.1) + 1e-12:
                bad += 1
        if exists and max(outs) < 0.75:
            bad += 1
        if not exists and max(outs) > 0.25:
            bad += 1
        for a, o in zip(A, outs):
            has = any(ip(a, b) == 0 for b in B)
            if has and o < 0.75:
                bad += 1
            if (not has) and o > 0.25:
                bad += 1
    print(f"[5] exact-OV gadget (per row): {'PASS' if bad == 0 else 'FAIL'} (bad={bad})")
    return bad == 0


# ---------------------------------------------------------------------------------------------
# 6. Sampling (Proposition 3.1): m = ceil(40 e^W ln(2 nq (1+dv)/delta)/eps^2) samples per row give
#    error <= eps for all rows w.p. >= 1-delta.  Statistical check on random instances.
# ---------------------------------------------------------------------------------------------
def check_sampling(trials=30):
    worst_ratio = 0.0
    fails = 0
    for _ in range(trials):
        nk = 2000
        W = random.choice([1.0, 2.0, 3.0])
        eps = random.choice([0.2, 0.3])
        nq, dv, delta = 5, 1, 0.1
        m = math.ceil(40 * math.exp(W) * math.log(2 * nq * (1 + dv) / delta) / eps ** 2)
        for _ in range(nq):
            # logits spread exactly W; heavy keys rare (worst case for uniform sampling)
            logits = [W if random.random() < 0.01 else 0.0 for _ in range(nk)]
            vals = [1.0 if l > 0 else -1.0 for l in logits]
            o = softmax_out(logits, vals)
            idx = [random.randrange(nk) for _ in range(m)]
            w = [math.exp(logits[j]) for j in idx]
            oh = sum(wi * vals[j] for wi, j in zip(w, idx)) / sum(w)
            worst_ratio = max(worst_ratio, abs(oh - o) / eps)
            if abs(oh - o) > eps:
                fails += 1
    ok = fails <= 1  # allowed failure probability 0.1 per instance batch
    print(f"[6] sampling estimator: worst |err|/eps = {worst_ratio:.3f}, failures {fails}: {'PASS' if ok else 'FAIL'}")
    return ok


# ---------------------------------------------------------------------------------------------
# 7. Taylor bound (for Proposition 4.4): |P_g(x)/e^x - 1| <= e^{2R} (eR/g)^g on [-R, R].
# ---------------------------------------------------------------------------------------------
def check_taylor():
    ok = True
    for R in [0.5, 1.0, 2.0, 3.0]:
        for g in [4, 8, 12, 20, 30]:
            bound = math.exp(2 * R) * (math.e * R / g) ** g
            worst = 0.0
            for i in range(2001):
                x = -R + 2 * R * i / 2000
                xf = Fraction(x)
                P = sum(xf ** t / math.factorial(t) for t in range(g))
                worst = max(worst, abs(float(P) / math.exp(x) - 1))
            if worst > bound * (1 + 1e-9) + 1e-15:
                ok = False
                print(f"    Taylor bound violated: R={R} g={g} worst={worst} bound={bound}")
    print(f"[7] Taylor relative-error bound: {'PASS' if ok else 'FAIL'}")
    return ok


# ---------------------------------------------------------------------------------------------
# 8. End-to-end run of the OV algorithm of Theorem 4.1 on random instances.
#    The attention oracle is simulated exactly (weights (1+1/r)^s are rational) and its answer is
#    perturbed by +-eps_0 in every coordinate. Four isolation rounds.
# ---------------------------------------------------------------------------------------------
def ov_via_parity(A, B, r, rounds=4, rng=random):
    m = len(A[0])
    n = len(A)
    if any(sum(a) == 0 for a in A):
        return True
    ones = [1] * m
    z = Fraction(r + 1, r)
    N = 2 * n
    c_eff = max(1.0, m / max(1.0, math.log2(n)))  # m = ceil(c log n) with c_eff >= m/log n
    eps0 = Fraction(r ** m, 32 * n * (r + 1) ** (2 * m))
    inv = pow(r ** m, -1, r + 1) if r + 1 > 1 else 0
    for _ in range(rounds):
        keep = [rng.random() < 0.5 for _ in B]
        keys = [b if k else ones for b, k in zip(B, keep)]
        for a in A:
            S = sum(z ** ip(a, kk) for kk in keys)
            o = Fraction(n) / (n + S)
            oh = o + rng.choice([-1, 1]) * eps0
            X = math.floor(r ** m * n * (1 / oh - 1) + Fraction(1, 2))
            if (X * inv) % (r + 1) != 0:
                return True
    return False


def check_parity_end_to_end(trials=400):
    errs_yes = errs_no = n_yes = n_no = 0
    for _ in range(trials):
        m = random.randint(3, 7)
        n = random.randint(2, 9)
        r = random.choice([1, 2, 3, 5])
        A = [rand_vec(m, 0.6) for _ in range(n)]
        B = [rand_vec(m, 0.6) for _ in range(n)]
        truth = any(ip(a, b) == 0 for a in A for b in B)
        got = ov_via_parity(A, B, r)
        if truth:
            n_yes += 1
            errs_yes += (not got)
        else:
            n_no += 1
            errs_no += got
    # NO instances must never be accepted; YES instances may fail with prob <= 1/16 per instance
    ok = errs_no == 0 and errs_yes <= max(3, n_yes // 8)
    print(f"[8] end-to-end parity algorithm: NO-errors {errs_no}/{n_no}, YES-misses {errs_yes}/{n_yes}: "
          f"{'PASS' if ok else 'FAIL'}")
    return ok


# ---------------------------------------------------------------------------------------------
# 9. Block embedding (Theorem 4.1 with blocks): F(a,b) = #blocks in which a and b intersect equals
#    <phi(a), psi(b)> with phi(a)_{t,S} = (-1)^{|S|+1} prod_{k in S} a_k, psi(b)_{t,S} = prod_{k in S} b_k,
#    over nonempty S inside block t; dimension sum_t (2^{|B_t|} - 1). Then an end-to-end run of the
#    counting algorithm with K = #blocks digits in place of m.
# ---------------------------------------------------------------------------------------------
from itertools import combinations


def blocks_of(m, bsz):
    return [list(range(s, min(s + bsz, m))) for s in range(0, m, bsz)]


def phi_psi(a, b, blocks):
    ph, ps = [], []
    for blk in blocks:
        for size in range(1, len(blk) + 1):
            for S in combinations(blk, size):
                ph.append((-1) ** (size + 1) * math.prod(a[k] for k in S))
                ps.append(math.prod(b[k] for k in S))
    return ph, ps


def F_blocks(a, b, blocks):
    return sum(1 for blk in blocks if any(a[k] and b[k] for k in blk))


def check_block_embedding(trials=2000):
    bad = 0
    for _ in range(trials):
        m = random.randint(1, 10)
        bsz = random.randint(1, 4)
        blocks = blocks_of(m, bsz)
        a, b = rand_vec(m), rand_vec(m)
        ph, ps = phi_psi(a, b, blocks)
        if sum(x * y for x, y in zip(ph, ps)) != F_blocks(a, b, blocks):
            bad += 1
        if len(ph) != sum(2 ** len(blk) - 1 for blk in blocks):
            bad += 1
    print(f"[9a] block embedding identity and dimension: {'PASS' if bad == 0 else 'FAIL'} (bad={bad})")
    return bad == 0


def ov_via_blocks(A, B, r, bsz, rounds=4, rng=random):
    m = len(A[0])
    n = len(A)
    if any(sum(a) == 0 for a in A):
        return True
    blocks = blocks_of(m, bsz)
    K = len(blocks)
    ones = [1] * m
    z = Fraction(r + 1, r)
    eps0 = Fraction(r ** K, 32 * n * (r + 1) ** (2 * K))
    inv = pow(r ** K, -1, r + 1)
    for _ in range(rounds):
        keep = [rng.random() < 0.5 for _ in B]
        keys = [b if k else ones for b, k in zip(B, keep)]
        for a in A:
            S = sum(z ** F_blocks(a, kk, blocks) for kk in keys)
            o = Fraction(n) / (n + S)
            oh = o + rng.choice([-1, 1]) * eps0
            X = math.floor(r ** K * n * (1 / oh - 1) + Fraction(1, 2))
            if (X * inv) % (r + 1) != 0:
                return True
    return False


def check_blocks_end_to_end(trials=400):
    errs_yes = errs_no = n_yes = n_no = 0
    for _ in range(trials):
        m = random.randint(3, 9)
        n = random.randint(2, 9)
        r = random.choice([1, 2, 3])
        bsz = random.randint(1, 4)
        A = [rand_vec(m, 0.6) for _ in range(n)]
        B = [rand_vec(m, 0.6) for _ in range(n)]
        truth = any(ip(a, b) == 0 for a in A for b in B)
        got = ov_via_blocks(A, B, r, bsz)
        if truth:
            n_yes += 1
            errs_yes += (not got)
        else:
            n_no += 1
            errs_no += got
    ok = errs_no == 0 and errs_yes <= max(3, n_yes // 8)
    print(f"[9b] end-to-end block counting: NO-errors {errs_no}/{n_no}, YES-misses {errs_yes}/{n_yes}: "
          f"{'PASS' if ok else 'FAIL'}")
    return ok


def log_softmax_out(groups):
    """groups: list of (log count, logit, value); returns the attention output, computed stably."""
    top = max(lc + l for lc, l, v in groups)
    num = sum(math.exp(lc + l - top) * v for lc, l, v in groups)
    den = sum(math.exp(lc + l - top) for lc, l, v in groups)
    return num / den


def check_curve_gadget():
    """Gadget of Theorem 3.6 with worst-case counts: n_B = ceil(n^alpha) real keys, all at the NO logit
    lambda*floor(T/tau) except possibly one at lambda*T, each key and the dummy copied M times."""
    ok = True
    worst_no, worst_yes, worst_spread = 0.0, 1.0, 0.0
    for alpha in (0.1, 0.3, 0.5, 0.9, 1.0):
        for theta in (1.0, 0.3, 0.05):
            for eps in (1e-6, 1 / 6, 0.45):
                tau = math.ceil(4 / theta)
                eta = (1 - 2 * eps) / 4
                # n1 of step <1>5: (1+theta/2)(2 ln(1/eta) + 11/10) <= (theta/2) alpha ln n
                lnn1 = (1 + theta / 2) * (2 * math.log(1 / eta) + 1.1) * 2 / (theta * alpha)
                for logn in (200, 2000, 20000, lnn1 / math.log(2), 3 * lnn1 / math.log(2)):
                    lnn = logn * math.log(2)
                    T = 7 * logn
                    lognB = alpha * lnn + math.log1p(math.exp(-alpha * lnn))   # log(n^alpha + 1) >= log ceil(n^alpha)
                    lognB1 = lognB + math.log1p(-math.exp(-lognB))             # log(nB - 1)
                    xi = math.ceil(10 * math.log((1 - eta) / eta)) / 10
                    lam = math.ceil(10 * (alpha * lnn + 2 * math.log(1 / eta) + 1)) / (10 * T * (1 - 1 / tau))
                    tstar = lam * T - xi
                    lM = math.log(1000.0)
                    no_logit = lam * math.floor(T / tau)
                    o_no = log_softmax_out([(lM + lognB, no_logit, 1.0), (lM, tstar, 0.0)])
                    o_yes = log_softmax_out([(lM, lam * T, 1.0), (lM + lognB1, no_logit, 1.0), (lM, tstar, 0.0)])
                    o_yes_nodup = log_softmax_out([(0.0, lam * T, 1.0), (lognB1, no_logit, 1.0), (0.0, tstar, 0.0)])
                    spread = lam * T / lnn
                    big = lnn >= lnn1 * (1 + 1e-9)    # n >= n1(alpha, theta, eps) of step <1>5
                    if o_no > eta or o_yes < 1 - eta or abs(o_yes - o_yes_nodup) > 1e-12 or tstar < 0:
                        ok = False
                    # step <2>3 holds for every n; the spread bound for n >= n1
                    if lam * T > (1 + theta / 2) * (alpha * lnn + 2 * math.log(1 / eta) + 1.1) * (1 + 1e-12):
                        ok = False
                    if big and spread > (1 + theta) * alpha:
                        ok = False
                    worst_no = max(worst_no, o_no / eta)
                    worst_yes = min(worst_yes, (1 - o_yes) / eta)
                    if big:
                        worst_spread = max(worst_spread, spread / ((1 + theta) * alpha))
    print(f"[10] curve gadget with duplication: max NO/eta = {worst_no:.3f}, max (1-YES)/eta = "
          f"{1 - worst_yes if worst_yes < 1 else 0:.3f}, max spread/((1+theta)alpha ln n) = {worst_spread:.3f}: "
          f"{'PASS' if ok else 'FAIL'}")
    return ok


def check_sign_estimator(trials=400):
    """Random-sign protocol of Theorem 5.6: empirical failure rate of |Z - <q,k>| < Delta/2 (evidence only)."""
    rng = random.Random(7)
    worst = 0.0
    for d, Gamma in ((8, 0.5), (32, 1.0), (64, 2.0), (16, 3.0)):
        q = [rng.gauss(0, 1) for _ in range(d)]
        k = [rng.gauss(0, 1) for _ in range(d)]
        nq = math.sqrt(sum(x * x for x in q))
        nk = math.sqrt(sum(x * x for x in k))
        sigma = nq * nk
        Delta = sigma / Gamma
        true = sum(a * b for a, b in zip(q, k))
        u = math.ceil(384 * Gamma ** 2)
        t0 = math.sqrt(2 * math.log(32 * u))
        w = math.ceil(math.log2(8 * t0 ** 2 * Gamma))
        fails = 0
        for _ in range(trials):
            Z = 0.0
            for _ in range(u):
                s = [rng.choice((-1, 1)) for _ in range(d)]
                sk = sum(a * b for a, b in zip(s, k))
                sq = sum(a * b for a, b in zip(s, q))
                Y = max(-t0 * nk, min(t0 * nk, sk))
                kh = round((2 ** w) * Y / (t0 * nk))
                Z += sq * t0 * nk * (2 ** -w) * kh
            Z /= u
            fails += abs(Z - true) >= Delta / 2
        worst = max(worst, fails / trials)
    ok = worst <= 0.25
    print(f"[11] random-sign estimator: worst empirical failure rate {worst:.3f} (claimed <= 1/4): "
          f"{'PASS' if ok else 'FAIL'}")
    return ok


def check_pm_values():
    """Remark 3.5: the gadget at precision eps/2 on values {0,1} separates at precision eps on values {-1,1}."""
    ok = True
    for eps in (1e-6, 0.1, 1 / 3, 0.5, 0.9, 0.999):
        eta = (1 - eps) / 4                   # eta = (1 - 2(eps/2))/4
        xi = math.ceil(10 * math.log((1 - eta) / eta)) / 10
        for alpha, lnn in ((1.0, 400.0), (0.3, 2000.0)):
            lognB = alpha * lnn + math.log1p(math.exp(-alpha * lnn))
            lamT = math.ceil(10 * (alpha * lnn + 2 * math.log(1 / eta) + 1)) / (10 * (1 - 1 / 8))
            tstar = lamT - xi
            no_logit = lamT / 8
            yes = log_softmax_out([(0.0, lamT, 1.0), (lognB, no_logit, 1.0), (0.0, tstar, -1.0)])
            no = log_softmax_out([(lognB, no_logit, 1.0), (0.0, tstar, -1.0)])
            if not (yes >= 1 - 2 * eta and no <= -1 + 2 * eta and yes - eps > 0 > no + eps):
                ok = False
    print(f"[12] values in {{-1,1}}, every eps < 1: {'PASS' if ok else 'FAIL'}")
    return ok


if __name__ == "__main__":
    results = [check_mod_identity(), check_reconstruction(), check_isolation(), check_gap_gadget(),
               check_exact_ov_gadget(), check_sampling(), check_taylor(), check_parity_end_to_end(),
               check_block_embedding(), check_blocks_end_to_end(), check_curve_gadget(),
               check_sign_estimator(), check_pm_values()]
    print("ALL PASS" if all(results) else "SOME CHECK FAILED")