"""Micro-benchmark for zen_fronts.selection_core.mc_rank_select_detailed.

Run:
  python examples/bench_mc_rank.py --out out/bench_mc_rank.csv

This benchmarks the *selection core* only (Cython if installed), not CriteriaStore.

Scaling intuition
-----------------
The Pareto core uses a domination matrix and checks all pairs, so one rating is
~O(N^2*M). Monte Carlo repeats it n_samples times, so overall ~O(S*N^2*M).
"""

from __future__ import annotations

import argparse
import csv
import os
import time
from dataclasses import asdict, dataclass

import numpy as np

from zen_fronts.selection_core import mc_rank_select_detailed


@dataclass
class BenchRow:
    N: int
    M: int
    n_samples: int
    percentile: float
    mode: str
    quantiles_mode_i: int
    quantiles_budget: int
    repeats: int
    warmup: int
    sigma2_kind: str
    seconds_median: float
    seconds_min: float
    seconds_max: float


def _parse_int_list(s: str) -> list[int]:
    return [int(x.strip()) for x in s.split(",") if x.strip()]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--Ns", default="128,256,512", help="Comma-separated N sizes")
    ap.add_argument("--Ms", default="2,4,8", help="Comma-separated M sizes")
    ap.add_argument("--Ss", default="64,128,256", help="Comma-separated n_samples sizes")
    ap.add_argument("--percentile", type=float, default=0.2)
    ap.add_argument("--mode", choices=["values", "ranks"], default="ranks")
    ap.add_argument("--quantiles-modes", default="0,1,2", help="Comma-separated {0,1,2}")
    ap.add_argument("--quantiles-budget", type=int, default=2_500_000)
    ap.add_argument("--sigma2", choices=["random", "zero"], default="random")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--warmup", type=int, default=1)
    ap.add_argument("--repeats", type=int, default=3)
    ap.add_argument("--out", default="", help="CSV output path (optional)")

    args = ap.parse_args()

    Ns = _parse_int_list(args.Ns)
    Ms = _parse_int_list(args.Ms)
    Ss = _parse_int_list(args.Ss)
    qms = _parse_int_list(args.quantiles_modes)

    rng = np.random.default_rng(args.seed)

    mode_i = 0 if args.mode == "values" else 1

    rows: list[BenchRow] = []

    for N in Ns:
        for M in Ms:
            mu = rng.normal(size=(N, M)).astype(np.float64)
            if args.sigma2 == "zero":
                sigma2 = np.zeros((N, M), dtype=np.float64)
            else:
                sigma2 = (0.01 + rng.random((N, M))).astype(np.float64)

            for S in Ss:
                for qm in qms:
                    # Warmups
                    for _ in range(max(0, args.warmup)):
                        mc_rank_select_detailed(
                            mu,
                            sigma2,
                            n_samples=S,
                            percentile=args.percentile,
                            seed=123,
                            mode_i=mode_i,
                            quantiles_mode_i=qm,
                            quantiles_budget=args.quantiles_budget,
                        )

                    times: list[float] = []
                    for r in range(args.repeats):
                        t0 = time.perf_counter()
                        mc_rank_select_detailed(
                            mu,
                            sigma2,
                            n_samples=S,
                            percentile=args.percentile,
                            seed=123 + r,
                            mode_i=mode_i,
                            quantiles_mode_i=qm,
                            quantiles_budget=args.quantiles_budget,
                        )
                        t1 = time.perf_counter()
                        times.append(t1 - t0)

                    times_sorted = sorted(times)
                    med = times_sorted[len(times_sorted) // 2]

                    row = BenchRow(
                        N=N,
                        M=M,
                        n_samples=S,
                        percentile=float(args.percentile),
                        mode=args.mode,
                        quantiles_mode_i=int(qm),
                        quantiles_budget=int(args.quantiles_budget),
                        repeats=int(args.repeats),
                        warmup=int(args.warmup),
                        sigma2_kind=args.sigma2,
                        seconds_median=float(med),
                        seconds_min=float(min(times)),
                        seconds_max=float(max(times)),
                    )
                    rows.append(row)

                    # concise console output
                    est_ops = S * (N * N) * M
                    print(
                        f"N={N:5d} M={M:2d} S={S:5d} qm={qm}  "
                        f"median={med:8.4f}s  (min={min(times):.4f}, max={max(times):.4f})  "
                        f"~{est_ops/1e9:.2f}e9 pair-ops"
                    )

    if args.out:
        os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
        with open(args.out, "w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=list(asdict(rows[0]).keys()))
            w.writeheader()
            for row in rows:
                w.writerow(asdict(row))
        print(f"\nWrote CSV: {args.out}")


if __name__ == "__main__":
    main()
