"""Statistical regularity tests on r-step SHA-256 (single block, standard IV, with feed-forward), C1.

For N random 512-bit messages (seeded) and r in R = {1,2,3,4,5,6,8,10,12,16,20,24,32,64}:
 (i)   bias of each of the 256 output bits, |mean - 0.5|; summary = max and number above the 4-sigma level
       4 * 0.5/sqrt(N);
 (ii)  linear correlation c(i,j) = E[(-1)^(x_i xor y_j)] between each of 64 randomly sampled input bits
       (from W_0..W_15) and every output bit; count of pairs with |c| > 4/sqrt(N) (under H0: c ~ N(0, 1/N), so the
       expected count under uniformity is 64*256*P(|Z|>4) ~ 1.04);
 (iii) chi-square uniformity test on each of the 32 output bytes (255 d.f.), min p-value and number of bytes
       rejected at alpha = 0.01 / 32 (Bonferroni).
Controls: r = 64 (full compression) and a random oracle (hashlib SHA-256 over a secret salt and the message).
One 64-step trace per message gives every r at once (state after r steps + IV feed-forward).

Usage: PYTHONPATH=src .venv/bin/python experiments/reduced_round_sha256/bias_test.py --seed 1 --n 20000 \
           --out results/reduced_round_sha256/bias.json
"""
from __future__ import annotations
import argparse, os, random, sys, time
import numpy as np
from scipy import stats as sps
from common import S, provenance, random_block, state_to_int, random_oracle, peak_rss_mb, MASK

R_LIST = [1, 2, 3, 4, 5, 6, 8, 10, 12, 16, 20, 24, 32, 64]


def bits_matrix(ints, nbits=256) -> np.ndarray:
    """N x nbits uint8 matrix, column 0 = most significant bit."""
    arr = np.array([[(v >> (nbits - 1 - i)) & 1 for i in range(nbits)] for v in ints], dtype=np.uint8)
    return arr


def analyse(Y: np.ndarray, X: np.ndarray, n: int, label: str):
    thr = 4 * 0.5 / n ** 0.5
    bias = np.abs(Y.mean(axis=0) - 0.5)
    Xs = 1.0 - 2.0 * X.astype(np.float64)
    Ys = 1.0 - 2.0 * Y.astype(np.float64)
    C = Xs.T @ Ys / n  # 64 x 256
    cthr = 4 / n ** 0.5
    big = int((np.abs(C) > cthr).sum())
    # chi-square per output byte
    pvals, chis = [], []
    for byte in range(32):
        cols = Y[:, byte * 8:(byte + 1) * 8].astype(np.int64)
        vals = (cols * (1 << np.arange(7, -1, -1))).sum(axis=1)
        counts = np.bincount(vals, minlength=256)
        chi, p = sps.chisquare(counts)
        pvals.append(float(p)); chis.append(float(chi))
    alpha = 0.01 / 32
    res = {
        "max_bit_bias": float(bias.max()), "mean_bit_bias": float(bias.mean()),
        "bias_threshold_4sigma": thr, "n_bits_bias_over_4sigma": int((bias > thr).sum()),
        "bit_bias": bias.round(6).tolist(),
        "max_abs_corr": float(np.abs(C).max()), "corr_threshold_4sigma": cthr,
        "n_pairs_corr_over_4sigma": big, "expected_pairs_under_null": float(64 * 256 * 2 * sps.norm.sf(4)),
        "argmax_corr_pair": [int(x) for x in np.unravel_index(np.abs(C).argmax(), C.shape)],
        "chi2_min_p": float(min(pvals)), "chi2_n_bytes_rejected_bonferroni": int(sum(p < alpha for p in pvals)),
        "chi2_p_by_byte": pvals, "chi2_stat_by_byte": chis,
    }
    print(f"[{label}] max bias {res['max_bit_bias']:.4f} (#>4s {res['n_bits_bias_over_4sigma']}), "
          f"max|corr| {res['max_abs_corr']:.4f} (#>4s {big}), chi2 min p {res['chi2_min_p']:.3g} "
          f"(rejected {res['chi2_n_bytes_rejected_bonferroni']}/32)", file=sys.stderr)
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--n", type=int, default=20000)
    ap.add_argument("--n-input-bits", type=int, default=64)
    ap.add_argument("--out", default="results/reduced_round_sha256/bias.json")
    a = ap.parse_args()
    t0 = time.time()
    rng = random.Random(a.seed)
    salt = os.urandom(32)
    in_bits = sorted(rng.sample(range(512), a.n_input_bits))
    blocks = [random_block(rng) for _ in range(a.n)]
    traces = []
    for blk in blocks:
        tr = []
        S.compress(S.H0, blk, 64, feedforward=False, trace=tr)
        traces.append(tr)
    msg_ints = [int.from_bytes(b, "big") for b in blocks]
    X = bits_matrix(msg_ints, 512)[:, in_bits]
    meas = {"input_bits_sampled": in_bits, "by_r": {}}
    for r in R_LIST:
        outs = [state_to_int([(x + y) & MASK for x, y in zip(S.H0, tr[r - 1])]) for tr in traces]
        meas["by_r"][str(r)] = analyse(bits_matrix(outs), X, a.n, f"r={r}")
    ro = [random_oracle(salt, b) for b in blocks]
    meas["random_oracle_control"] = analyse(bits_matrix(ro), X, a.n, "RO")
    stats = {"compress_calls": a.n, "wall_s": round(time.time() - t0, 3), "peak_rss_mb": peak_rss_mb()}
    provenance.record("sha256_bias_test", {"n": a.n, "r_list": R_LIST, "n_input_bits": a.n_input_bits,
                                           "feedforward": True, "iv": "standard H0"},
                      meas, seed=a.seed, extra={"op_stats": stats}, out=a.out, started=t0)
    print("wrote", a.out, stats)


if __name__ == "__main__":
    main()