#!/usr/bin/env python3
"""Run CAM retrieval benchmark noise sweep (0%–100%, step 10%) and generate summary.

Usage:
    python scripts/run_retrieval_noise_sweep.py \
        --dataset outputs/cam_retrieval_benchmark/datasets/cifar10_hash512_rows512_queries128.npz \
        --num-rows 512 \
        --output docs/cam_retrieval_noise_sweep.md
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Iterator

PROJECT_ROOT = Path(__file__).resolve().parents[1]

DEFAULT_NOISE_RATES = list(range(0, 101, 10))  # 0, 10, 20, ..., 100


@dataclass
class RunResult:
    noise_pct: int
    run_id: str
    metrics: dict  # {"1": {...}, "5": {...}}
    success: bool
    error_msg: str = ""
    params: dict = field(default_factory=dict)
    dataset_info: dict = field(default_factory=dict)


def run_single(
    dataset: str,
    num_rows: int,
    noise_pct: int,
    topk_k: int = 5,
    hash_bits: int = 512,
) -> RunResult:
    """Run a single benchmark with the given noise rate."""
    write_noise_en = 0 if noise_pct == 0 else 1
    # Include dataset stem in run_id to avoid cross-dataset overwrites
    dataset_stem = Path(dataset).stem
    run_id = f"noise_sweep_{dataset_stem}_{noise_pct:03d}pct"

    env = os.environ.copy()
    env["CAM_RETRIEVAL_RUN_ID"] = run_id
    env["CAM_RETRIEVAL_DATASET"] = dataset

    make_args = [
        "make",
        "-C", "hw/sim",
        f"TOPK_K={topk_k}",
        f"NUM_ROWS={num_rows}",
        f"HASH_BITS={hash_bits}",
        f"WRITE_NOISE_EN={write_noise_en}",
    ]
    if write_noise_en:
        make_args.extend([
            f"WRITE_NOISE_RATE_NUM={noise_pct}",
            "WRITE_NOISE_RATE_DEN=100",
        ])

    clean_cmd = ["make", "-C", "hw/sim", "clean"]
    test_cmd = make_args + ["test-benchmark-retrieval"]

    cwd = str(PROJECT_ROOT)

    # Clean — stream output
    print("  [clean]", flush=True)
    result_clean = subprocess.run(
        clean_cmd, cwd=cwd, env=env,
        timeout=120,
    )
    if result_clean.returncode != 0:
        return RunResult(
            noise_pct=noise_pct,
            run_id=run_id,
            metrics={},
            success=False,
            error_msg=f"clean failed (rc={result_clean.returncode})",
        )

    # Run benchmark — stream output but capture stderr for error reporting
    print("  [make test-benchmark-retrieval]", flush=True)
    try:
        result = subprocess.run(
            test_cmd, cwd=cwd, env=env,
            stderr=subprocess.PIPE,
            text=True,
            timeout=1800,  # 30 min per run
        )
    except subprocess.TimeoutExpired:
        return RunResult(
            noise_pct=noise_pct,
            run_id=run_id,
            metrics={},
            success=False,
            error_msg="test timed out (30 min)",
        )

    if result.returncode != 0:
        return RunResult(
            noise_pct=noise_pct,
            run_id=run_id,
            metrics={},
            success=False,
            error_msg=f"test failed (rc={result.returncode}):\n{result.stderr[-2000:]}",
        )

    # Read result from output
    return _read_result(noise_pct, run_id)


def _read_result(noise_pct: int, run_id: str) -> RunResult:
    """Read benchmark results from output directory."""
    out_dir = PROJECT_ROOT / "outputs" / "cam_retrieval_benchmark" / run_id
    metrics_file = out_dir / "metrics.json"

    if not metrics_file.exists():
        return RunResult(
            noise_pct=noise_pct,
            run_id=run_id,
            metrics={},
            success=False,
            error_msg=f"metrics.json not found at {metrics_file}",
        )

    try:
        data = json.loads(metrics_file.read_text(encoding="utf-8"))
    except Exception as exc:
        return RunResult(
            noise_pct=noise_pct,
            run_id=run_id,
            metrics={},
            success=False,
            error_msg=f"Failed to parse metrics.json: {exc}",
        )

    return RunResult(
        noise_pct=noise_pct,
        run_id=run_id,
        metrics=data.get("metrics", {}),
        success=data.get("status") == "pass",
        params=data.get("params", {}),
        dataset_info=data.get("dataset", {}),
    )


def iter_noise_rates() -> Iterator[int]:
    return iter(DEFAULT_NOISE_RATES)


def run_sweep(
    dataset: str,
    num_rows: int,
    topk_k: int = 5,
    hash_bits: int = 512,
) -> list[RunResult]:
    """Run the full noise sweep, return all results."""
    results: list[RunResult] = []
    total = len(DEFAULT_NOISE_RATES)

    for idx, noise_pct in enumerate(DEFAULT_NOISE_RATES):
        pct_str = f"{noise_pct:3d}%"
        print(f"\n{'='*60}")
        print(f"  [{idx+1:2d}/{total}]  Noise rate: {pct_str}")
        print(f"{'='*60}\n", flush=True)

        run_result = run_single(
            dataset=dataset,
            num_rows=num_rows,
            noise_pct=noise_pct,
            topk_k=topk_k,
            hash_bits=hash_bits,
        )
        results.append(run_result)

        if run_result.success:
            k1 = run_result.metrics.get("1", {})
            k5 = run_result.metrics.get("5", {})
            print(f"  ✓ PASS  recall@1={k1.get('recall@k', '?'):.4f}  "
                  f"recall@5={k5.get('recall@k', '?'):.4f}  "
                  f"exact@5={k5.get('exact_match_rate', '?'):.4f}")
        else:
            print(f"  ✗ FAIL  {run_result.error_msg[:200]}")

    return results


def generate_summary(
    results: list[RunResult],
    dataset: str,
    num_rows: int,
    topk_k: int,
    hash_bits: int,
    output_path: Path,
) -> None:
    """Generate comprehensive markdown summary."""
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    n_total = len(results)
    n_pass = sum(1 for r in results if r.success)
    n_fail = sum(1 for r in results if not r.success)

    dataset_info = results[0].dataset_info if results else {}

    lines = [
        "# CAM Retrieval Benchmark — Noise Sweep Summary",
        "",
        f"**Generated:** {now}",
        "",
        "## Configuration",
        "",
        "| Parameter | Value |",
        "|---|---|",
        f"| Dataset path | `{dataset}` |",
        f"| NUM_ROWS | {num_rows} |",
        f"| TOPK_K | {topk_k} |",
        f"| HASH_BITS | {hash_bits} |",
        f"| Noise rates | {', '.join(f'{p}%' for p in DEFAULT_NOISE_RATES)} |",
        f"| Total runs | {n_total} |",
        f"| Passed | {n_pass} |",
        f"| Failed | {n_fail} |",
    ]

    if dataset_info:
        lines.extend([
            "",
            "### Dataset Details",
            "",
            "| Field | Value |",
            "|---|---|",
            f"| num_queries | {dataset_info.get('num_queries', '?')} |",
            f"| num_classes | {dataset_info.get('num_classes', '?')} |",
            f"| seed | {dataset_info.get('seed', '?')} |",
        ])

    lines.extend([
        "",
        "---",
        "",
        "## Results by Noise Rate",
        "",
    ])

    # Per-rate detailed table
    for k in (1, 5):
        # Hit@K (hit-rate), Precision@K (mean per-query tp/k),
        # Hit-F1@K (F1 from Hit@K × Precision@K),
        # Std-Recall@K (mean per-query tp/|relevant|),
        # Std-F1@K (F1 from Precision@K × Std-Recall@K),
        # Golden Match@K (exact match with reference)
        lines.extend([
            f"### k={k}",
            "",
            "| Noise (%) | WRITE_NOISE_EN | NUM/DEN | Hit@K | Precision@K | Hit-F1@K | Std-Recall@K | Std-F1@K | Golden Match@K | Status |",
            "|---|---:|---|---:|---|---:|---|---:|---|---:|---|",
        ])
        for r in results:
            metrics = r.metrics.get(str(k), {})
            noise_en = r.params.get("write_noise_en", 0)
            status = "✓" if r.success else "✗"
            if noise_en:
                rate_num = r.params.get("write_noise_rate_num", 0)
                rate_den = r.params.get("write_noise_rate_den", 100)
                rate_str = f"{rate_num}/{rate_den}"
            else:
                rate_str = "—"

            # Extract raw metrics
            hit_k = metrics.get('recall@k', 0)
            if isinstance(hit_k, str): hit_k = 0.0
            prec_k = metrics.get('macro_precision', 0)
            if isinstance(prec_k, str): prec_k = 0.0
            std_recall = metrics.get('retrieval_recall', 0)
            if isinstance(std_recall, str): std_recall = 0.0
            golden = metrics.get('exact_match_rate', 0)
            if isinstance(golden, str): golden = 0.0

            # Hit-F1@K = 2*Hit@K*Precision@K / (Hit@K + Precision@K)
            if prec_k + hit_k > 0:
                hit_f1 = (2.0 * prec_k * hit_k) / (prec_k + hit_k)
            else:
                hit_f1 = 0.0

            # Std-F1@K = 2*Precision@K*Std-Recall@K / (Precision@K + Std-Recall@K)
            if prec_k + std_recall > 0:
                std_f1 = (2.0 * prec_k * std_recall) / (prec_k + std_recall)
            else:
                std_f1 = 0.0

            lines.append(
                f"| {r.noise_pct:3d}% | {noise_en} | "
                f"{rate_str} | "
                f"{hit_k:.6f} | "
                f"{prec_k:.6f} | "
                f"{hit_f1:.6f} | "
                f"{std_recall:.6f} | "
                f"{std_f1:.6f} | "
                f"{golden:.6f} | "
                f"{status} |"
            )
        lines.append("")

    # Comparison across noise levels (using primary metric: Hit@K)
    lines.extend([
        "---",
        "",
        "## Cross-Noise Comparison (primary: Hit@K)",
        "",
        "| Noise (%) | Hit@1 | Hit@5 | Δ(Hit@1 vs 0%) | Δ(Hit@5 vs 0%) |",
        "|---|---:|---:|---:|---:|",
    ])

    # Find baseline (0% noise)
    zero_result = next((r for r in results if r.noise_pct == 0 and r.success), None)
    base_r1 = float(zero_result.metrics.get("1", {}).get("recall@k", 0)) if zero_result else 0.0
    base_r5 = float(zero_result.metrics.get("5", {}).get("recall@k", 0)) if zero_result else 0.0

    for r in results:
        r1 = float(r.metrics.get("1", {}).get("recall@k", 0)) if r.success else float("nan")
        r5 = float(r.metrics.get("5", {}).get("recall@k", 0)) if r.success else float("nan")
        d1 = f"{r1 - base_r1:+.6f}" if r.success and zero_result else "—"
        d5 = f"{r5 - base_r5:+.6f}" if r.success and zero_result else "—"
        r1_str = f"{r1:.6f}" if r.success else "FAIL"
        r5_str = f"{r5:.6f}" if r.success else "FAIL"
        lines.append(f"| {r.noise_pct:3d}% | {r1_str} | {r5_str} | {d1} | {d5} |")

    # Failures section
    failures = [r for r in results if not r.success]
    if failures:
        lines.extend([
            "",
            "---",
            "",
            "## Failed Runs",
            "",
        ])
        for r in failures:
            lines.extend([
                f"### Noise rate: {r.noise_pct}%",
                "",
                "```",
                r.error_msg.strip(),
                "```",
                "",
            ])

    lines.extend([
        "",
        "---",
        "",
        "## Metric Definitions",
        "",
        "- **Hit@K**: fraction of queries where at least one relevant item appears in Top-K results (primary metric).",
        "- **Precision@K**: mean per-query precision — averaged `tp/k` across all queries.",
        "- **Hit-F1@K**: `2 × Hit@K × Precision@K / (Hit@K + Precision@K)` — F1 using hit-rate recall.",
        "- **Std-Recall@K**: mean per-query standard retrieval recall — `tp / |relevant|` averaged across queries (supplementary).",
        "- **Std-F1@K**: `2 × Precision@K × Std-Recall@K / (Precision@K + Std-Recall@K)` — F1 using standard recall (supplementary).",
        "- **Golden Match@K**: fraction of queries where DUT Top-K exactly matches the reference golden Top-K.",
        "",
        "The paper uses Hit@K and Precision@K as primary metrics. Std-Recall@K and Std-F1@K are supplementary,",
        "included to show Top-K coverage against all relevant items in the database.",
        "",
        "*Results from Verilator/Cocotb simulation. Not measured on physical FPGA hardware.*",
        "",
    ])

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"\n  Summary written to: {output_path}")


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Run CAM retrieval benchmark noise sweep (0%–100%, step 10%)"
    )
    parser.add_argument(
        "--dataset", required=True,
        help="Path to prepared .npz dataset file",
    )
    parser.add_argument(
        "--num-rows", type=int, default=512,
        help="Number of CAM rows (must match dataset, default: 512)",
    )
    parser.add_argument(
        "--topk-k", type=int, default=5,
        help="TOPK_K parameter (default: 5)",
    )
    parser.add_argument(
        "--hash-bits", type=int, default=512,
        help="HASH_BITS parameter (default: 512)",
    )
    parser.add_argument(
        "--output", type=Path, default=PROJECT_ROOT / "docs" / "cam_retrieval_noise_sweep.md",
        help="Output summary markdown path",
    )
    parser.add_argument(
        "--noise-rates",
        type=str,
        default=None,
        help="Comma-separated noise rates (default: 0,10,20,...,100)",
    )
    args = parser.parse_args()

    # Resolve dataset path
    dataset_path = args.dataset
    if not os.path.isabs(dataset_path):
        dataset_path = str(PROJECT_ROOT / dataset_path)

    # Parse custom noise rates if provided
    global DEFAULT_NOISE_RATES
    if args.noise_rates:
        DEFAULT_NOISE_RATES = [int(x.strip()) for x in args.noise_rates.split(",")]

    print(f"Dataset:       {dataset_path}")
    print(f"NUM_ROWS:      {args.num_rows}")
    print(f"TOPK_K:        {args.topk_k}")
    print(f"HASH_BITS:     {args.hash_bits}")
    print(f"Noise rates:   {DEFAULT_NOISE_RATES}")
    print(f"Total runs:    {len(DEFAULT_NOISE_RATES)}")
    print(f"Output:        {args.output}")
    print()

    results = run_sweep(
        dataset=dataset_path,
        num_rows=args.num_rows,
        topk_k=args.topk_k,
        hash_bits=args.hash_bits,
    )

    generate_summary(
        results=results,
        dataset=dataset_path,
        num_rows=args.num_rows,
        topk_k=args.topk_k,
        hash_bits=args.hash_bits,
        output_path=args.output,
    )

    # Exit with non-zero if any run failed
    if any(not r.success for r in results):
        sys.exit(1)


if __name__ == "__main__":
    main()