from __future__ import annotations

import argparse

import numpy as np
from model.ref_model import match_top1, random_hashes, random_noise_masks


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(
        "--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]

    print("rate,top1_stability,avg_clean_margin")
    for rate in args.rates:
        stable = 0
        margins = []

        for q, clean in zip(queries, clean_results):
            noise_masks = random_noise_masks(
                rng,
                args.rows,
                width=args.width,
                bit_flip_rate=rate,
            )
            noisy = match_top1(q, rows, width=args.width, noise_masks=noise_masks)

            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},{stable / args.queries:.6f},{np.mean(margins):.3f}")


if __name__ == "__main__":
    main()