"""Sweep write-noise rates and measure top-1 stability.

Applies write-noise flip masks to stored rows (simulating noisy writes),
then queries the noisy rows and compares top-1 results against clean rows.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

SIM_ROOT = Path(__file__).resolve().parents[1]
if str(SIM_ROOT) not in sys.path:
    sys.path.insert(0, str(SIM_ROOT))

import numpy as np
from model.ref_model import (
    generate_write_flip_mask,
    match_top1,
    random_hashes,
)


def apply_write_noise(
    rows: list[int],
    *,
    width: int,
    rate_num: int,
    rate_den: int,
    noise_bits: int = 8,
    seed: int = 0,
) -> list[int]:
    """Apply write-noise flip masks to every row, returning noisy copies.

    *seed* is a 64-bit value (RTL NOISE_SEED). It is duplicated to form
    the 128-bit xorshift initial state: {seed, seed}.
    """
    noisy: list[int] = []
    state = (seed << 64) | seed
    mask_w = (1 << width) - 1
    for row in rows:
        flip, state = generate_write_flip_mask(
            state, width, noise_bits, rate_num, rate_den
        )
        noisy.append((row ^ flip) & mask_w)
    return noisy


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--rows", type=int, default=512)
    parser.add_argument("--queries", type=int, default=128)
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument("--noise-bits", type=int, default=8)
    parser.add_argument(
        "--rates",
        type=float,
        nargs="+",
        default=[0.0, 0.001, 0.005, 0.01, 0.02, 0.05],
    )
    args = parser.parse_args()

    rng = np.random.default_rng(args.seed)
    rows = random_hashes(rng, args.rows, width=args.width)

    # Construct simple positive queries by selecting existing rows.
    query_indices = rng.integers(0, args.rows, size=args.queries)
    queries = [rows[int(i)] for i in query_indices]

    clean_results = [match_top1(q, rows, width=args.width) for q in queries]

    # Use a fixed denominator matching the 8-bit sample space (2^8 = 256).
    # Note: floor() is used, matching RTL threshold = (rate_num * 256) // rate_den.
    # Rates below 1/256 (≈0.39%) collapse to zero under this scheme.
    rate_den = 256

    print("rate,rate_num,effective_prob,top1_stability,avg_clean_margin")
    for rate in args.rates:
        rate_num = int(rate * rate_den)
        effective = rate_num / rate_den if rate_den > 0 else 0.0
        stable = 0
        margins = []

        noisy_rows = apply_write_noise(
            rows,
            width=args.width,
            rate_num=rate_num,
            rate_den=rate_den,
            noise_bits=args.noise_bits,
            seed=args.seed,
        )

        for q, clean in zip(queries, clean_results):
            noisy = match_top1(q, noisy_rows, width=args.width)

            if noisy.top1_index == clean.top1_index:
                stable += 1

            sorted_scores = np.sort(clean.scores)
            margin = int(sorted_scores[-1] - sorted_scores[-2])
            margins.append(margin)

        print(f"{rate},{rate_num},{effective:.6f},{stable / args.queries:.6f},{np.mean(margins):.3f}")


if __name__ == "__main__":
    main()