"""Avalanche / diffusion analysis of the SHA-256 compression function, steps r = 1..64 (no feed-forward).

Research question C1 (diffusion). For a single-bit input difference we measure, after every step t = 1..64,
the fraction of the 256 state bits that flip, averaged over N random messages (seeded). Four configurations:
  W0, W5, W15 : one random bit of message word W_j is flipped (the difference is tracked through the message
                schedule, so for t >= 16 several expanded words differ);
  IV          : one random bit of the 256-bit initial state is flipped (message unchanged).
Negative control: a 256-bit random oracle (hashlib SHA-256 over a secret os.urandom salt and the input) evaluated on
the same input pairs — its flip ratio must be ~0.5 with sd ~ 0.5/sqrt(256) = 0.031 from the start.

Outputs (JSON via wp34.provenance.record): per configuration the [64 x 256] mean-flip matrix, the per-step mean and sd
of the Hamming weight of the state difference, and `full_diffusion_step` = first step with mean flip ratio >= 0.49.
Only one trace per message pair is needed: the state after r steps without feed-forward is trace[r-1].

Usage: PYTHONPATH=src .venv/bin/python experiments/reduced_round_sha256/diffusion.py --seed 1 --n 2000 \
           --out results/reduced_round_sha256/diffusion.json
"""
from __future__ import annotations
import argparse, os, random, sys, time
import numpy as np
from common import S, provenance, random_block, state_to_int, int_to_bits, random_oracle, peak_rss_mb, bootstrap_ci

CONFIGS = ["W0", "W5", "W15", "IV"]


def flip_block_bit(block: bytes, word: int, bit: int) -> bytes:
    words = S.bytes_to_words(block)
    words[word] ^= 1 << bit
    return S.words_to_bytes(words)


def run(seed: int, n: int, threshold: float):
    rng = random.Random(seed)
    salt = os.urandom(32)  # secret salt for the random-oracle control (never saved)
    out = {}
    stats = {"compress_calls": 0}
    for cfg in CONFIGS:
        flips = np.zeros((64, 256), dtype=np.int64)
        hw = np.zeros((n, 64), dtype=np.int64)
        ro_flips = np.zeros(n, dtype=np.int64)
        in_pos = []
        for i in range(n):
            block = random_block(rng)
            state = list(S.H0)
            if cfg == "IV":
                w, b = rng.randrange(8), rng.randrange(32)
                state2 = list(state)
                state2[w] ^= 1 << b
                block2 = block
                in_pos.append(("IV", w, b))
                x1, x2 = S.words_to_bytes(state) + block, S.words_to_bytes(state2) + block
            else:
                j = int(cfg[1:])
                b = rng.randrange(32)
                state2 = state
                block2 = flip_block_bit(block, j, b)
                in_pos.append(("W", j, b))
                x1, x2 = block, block2
            tr1, tr2 = [], []
            S.compress(state, block, 64, feedforward=False, trace=tr1)
            S.compress(state2, block2, 64, feedforward=False, trace=tr2)
            stats["compress_calls"] += 2
            for t in range(64):
                d = state_to_int(tr1[t]) ^ state_to_int(tr2[t])
                bits = int_to_bits(d, 256)
                flips[t] += bits
                hw[i, t] = sum(bits)
            ro_flips[i] = bin(random_oracle(salt, x1) ^ random_oracle(salt, x2)).count("1")
        mean_matrix = flips / n
        ratio = mean_matrix.mean(axis=1)  # mean flip ratio per step
        full = next((t + 1 for t in range(64) if ratio[t] >= threshold), None)
        brng = random.Random(seed + 1000)
        out[cfg] = {
            "n": n,
            "mean_flip_matrix": mean_matrix.round(5).tolist(),
            "mean_flip_ratio_by_step": ratio.round(6).tolist(),
            "hamming_mean_by_step": hw.mean(axis=0).round(4).tolist(),
            "hamming_sd_by_step": hw.std(axis=0, ddof=1).round(4).tolist(),
            "hamming_ci95_step64": bootstrap_ci(hw[:, 63].tolist(), brng),
            "full_diffusion_step": full,
            "threshold": threshold,
            "random_oracle_control": {
                "flip_ratio_mean": float(ro_flips.mean() / 256),
                "flip_ratio_sd": float(ro_flips.std(ddof=1) / 256),
                "hamming_mean": float(ro_flips.mean()),
                "hamming_sd": float(ro_flips.std(ddof=1)),
            },
            "input_bit_positions_sample": in_pos[:20],
        }
        print(f"[{cfg}] full diffusion (>= {threshold}) at step {full}; ratio by step:",
              " ".join("%.3f" % x for x in ratio[:24]), "...", file=sys.stderr)
    return out, stats


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--n", type=int, default=2000)
    ap.add_argument("--threshold", type=float, default=0.49)
    ap.add_argument("--out", default="results/reduced_round_sha256/diffusion.json")
    a = ap.parse_args()
    t0 = time.time()
    meas, stats = run(a.seed, a.n, a.threshold)
    stats["wall_s"] = round(time.time() - t0, 3)
    stats["peak_rss_mb"] = peak_rss_mb()
    provenance.record("sha256_diffusion", {"n": a.n, "steps": 64, "threshold": a.threshold, "configs": CONFIGS,
                                           "feedforward": False, "iv": "standard H0"},
                      meas, seed=a.seed, extra={"op_stats": stats}, out=a.out, started=t0)
    print("wrote", a.out, "in", stats["wall_s"], "s")


if __name__ == "__main__":
    main()