"""Benchmark a realistic ZenFronts loop: update_crits -> refresh -> replace losers.

This benchmark aims to measure *end-to-end* ZenFronts overhead in a scenario
similar to "PBT-style" population maintenance:

  1) Update criteria for all active points
  2) refresh() (build matrices + MC ranking)
  3) For each loser: choose_parent() -> perform_new(remove_loser=True)
  4) Update criteria for each new child (so it's ready next epoch)

It records a CSV with per-epoch timings and also prints a short summary.

Notes
-----
* For timing internals without modifying the library, we monkeypatch:
    - CriteriaStore.build_matrices (instance)
    - mc_rank_select_detailed symbol imported by zen_fronts.selection.engine

* By default we update criteria via *external stats* (fast, deterministic).
  You can switch to raw updates (owrf) with --crits-source raw.

Run
---
  python examples/bench_refresh.py --out out/bench_refresh.csv

Typical (your target)
---------------------
  python examples/bench_refresh.py --Ns 128,256 --Ms 3 --Ss 64,128,256,512 --epochs 30
"""

from __future__ import annotations

import argparse
import csv
import math
import os
import platform
import sys
import time
from dataclasses import dataclass
from statistics import median
from time import perf_counter
from typing import Any, Callable, Iterable

import numpy as np

from zen_fronts import ZenFronts
from zen_fronts.selection import SelectionConfig
from zen_fronts.utils.boundary import clip01, reflect01

import zen_fronts.selection.engine as selection_engine_mod


# ------------------------------ toy objective (3 criteria) ------------------------------


def toy_objectives_3(params: dict[str, Any]) -> tuple[float, float, float]:
    """A cheap deterministic 3-objective function on [0,1]^2.

    f1, f2 are ZDT1-like; f3 is an extra smooth penalty.
    All are to be minimized.
    """

    x1 = float(params["x1"])
    x2 = float(params["x2"])

    # ZDT1-ish
    f1 = x1
    g = 1.0 + 9.0 * x2
    f2 = g * (1.0 - math.sqrt(max(f1 / g, 0.0)))

    # Third objective: smooth "bowl" near (0.25, 0.75)
    f3 = 0.15 * (x1 - 0.25) ** 2 + 0.85 * (x2 - 0.75) ** 2

    return f1, f2, f3


def make_stats_payload(
    values: tuple[float, float, float],
    *,
    sigma2: float,
    noise_scale: float,
    rng: np.random.Generator,
) -> dict[str, dict[str, float]]:
    """Return external-stats payload for update_crits()."""

    # We treat 'mu' as the observed value plus some jitter (optional).
    # sigma2 is the uncertainty passed into the selection core.
    out: dict[str, dict[str, float]] = {}
    for name, v in zip(("f1", "f2", "f3"), values, strict=True):
        mu = float(v) + float(noise_scale) * float(rng.normal())
        out[name] = {"mu": mu, "trend": 0.0, "sigma2": float(sigma2)}
    return out


def make_raw_payload(
    values: tuple[float, float, float],
    *,
    noise_scale: float,
    rng: np.random.Generator,
) -> dict[str, float]:
    """Return raw payload (float) for update_crits()."""

    out: dict[str, float] = {}
    for name, v in zip(("f1", "f2", "f3"), values, strict=True):
        out[name] = float(v) + float(noise_scale) * float(rng.normal())
    return out


# ------------------------------ timing hooks ------------------------------


@dataclass
class Hooks:
    build_s: float = 0.0
    core_s: float = 0.0


def _wrap_build_matrices(orig: Callable[..., Any], hooks: Hooks) -> Callable[..., Any]:
    def wrapped(point_ids: Iterable[int], *, now: float | None = None):
        t0 = perf_counter()
        out = orig(point_ids, now=now)
        hooks.build_s += perf_counter() - t0
        return out

    return wrapped


def _wrap_mc_core(orig: Callable[..., Any], hooks: Hooks) -> Callable[..., Any]:
    def wrapped(*args: Any, **kwargs: Any):
        t0 = perf_counter()
        out = orig(*args, **kwargs)
        hooks.core_s += perf_counter() - t0
        return out

    return wrapped


# ------------------------------ benchmark runner ------------------------------


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


def parse_float_list(s: str) -> list[float]:
    s = s.strip()
    if not s:
        return []
    return [float(x.strip()) for x in s.split(",") if x.strip()]


def run_one(
    *,
    N: int,
    S: int,
    percentile: float,
    epochs: int,
    seed: int,
    quantiles_mode_i: int,
    quantiles_budget: int,
    mode_i: int,
    within_mode_i: int,
    sigma2: float,
    noise_scale: float,
    crits_source: str,
    boundary: str,
) -> list[dict[str, Any]]:
    """Run one benchmark scenario and return per-epoch rows."""

    N = int(N)
    S = int(S)

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

    # --- sampler/mutator ---
    def sampler(rng_: np.random.Generator) -> dict[str, float]:
        return {"x1": float(rng_.random()), "x2": float(rng_.random())}

    if boundary == "reflect":
        bound = reflect01
    elif boundary == "clip":
        bound = clip01
    else:
        raise ValueError(f"unknown boundary mode: {boundary!r}")

    def mutator(parent: dict[str, Any], rng_: np.random.Generator, **_: Any) -> dict[str, float]:
        # Tiny local Gaussian perturbation
        sigma = 0.10
        x1 = bound(float(parent["x1"]) + float(rng_.normal(0.0, sigma)))
        x2 = bound(float(parent["x2"]) + float(rng_.normal(0.0, sigma)))
        return {"x1": x1, "x2": x2}, {}

    zf = ZenFronts(
        crits={"f1": "min", "f2": "min", "f3": "min"},
        selection=SelectionConfig(
            n_samples=int(S),
            percentile=float(percentile),
            seed=seed,
            mode_i=int(mode_i),
            within_mode_i=int(within_mode_i),
            quantiles_mode_i=int(quantiles_mode_i),
            quantiles_budget=int(quantiles_budget),
        ),
        sampler=sampler,
        mutator=mutator,
        config=None,
    )

    # --- install timing hooks ---
    hooks = Hooks()

    orig_build = zf._crits.build_matrices  # type: ignore[attr-defined]
    zf._crits.build_matrices = _wrap_build_matrices(orig_build, hooks)  # type: ignore[attr-defined]

    orig_core = selection_engine_mod.mc_rank_select_detailed
    selection_engine_mod.mc_rank_select_detailed = _wrap_mc_core(orig_core, hooks)  # type: ignore[assignment]

    rows: list[dict[str, Any]] = []

    try:
        # --- init population ---
        ids = zf.add_random_point(N)

        # Initial criteria update (so all are ready)
        t0 = perf_counter()
        for pid in ids:
            params = zf.params(pid)
            vals = toy_objectives_3(params)
            if crits_source == "stats":
                payload = make_stats_payload(vals, sigma2=sigma2, noise_scale=noise_scale, rng=rng)
            elif crits_source == "raw":
                payload = make_raw_payload(vals, noise_scale=noise_scale, rng=rng)
            else:
                raise ValueError(f"unknown crits_source: {crits_source!r}")
            zf.update_crits(pid, payload, t=0.0)
        init_update_s = perf_counter() - t0

        # --- epochs ---
        for ep in range(1, int(epochs) + 1):
            now = float(ep)

            # Update criteria for all active points (simulate evaluation loop)
            t0 = perf_counter()
            active_now = list(zf.points.keys())
            for pid in active_now:
                params = zf.params(pid)
                vals = toy_objectives_3(params)
                if crits_source == "stats":
                    payload = make_stats_payload(vals, sigma2=sigma2, noise_scale=noise_scale, rng=rng)
                else:
                    payload = make_raw_payload(vals, noise_scale=noise_scale, rng=rng)
                zf.update_crits(pid, payload, t=now)
            update_all_s = perf_counter() - t0

            # Refresh: time engine vs persist separately
            hooks.build_s = 0.0
            hooks.core_s = 0.0

            t0 = perf_counter()
            losers = zf._selector.refresh(now=now, seed=seed)  # type: ignore[attr-defined]
            engine_s = perf_counter() - t0

            rr = zf._selector.last_result  # type: ignore[attr-defined]
            if rr is None:
                raise RuntimeError("unexpected: last_result is None")

            t0 = perf_counter()
            zf._persist_refresh_stats(rr)  # type: ignore[attr-defined]
            persist_s = perf_counter() - t0

            refresh_s = engine_s + persist_s

            # Replace losers: choose parent + perform_new + update child crits
            choose_parent_s = 0.0
            perform_new_s = 0.0
            update_children_s = 0.0

            for loser in losers:
                t0 = perf_counter()
                parent = zf.choose_parent(loser, seed=seed)
                choose_parent_s += perf_counter() - t0

                t0 = perf_counter()
                child_id, _child_params = zf.perform_new(parent, looser=loser, remove_looser=True, seed=seed)
                perform_new_s += perf_counter() - t0

                # Child must become ready to participate next epoch
                t0 = perf_counter()
                vals = toy_objectives_3(zf.params(child_id))
                if crits_source == "stats":
                    payload = make_stats_payload(vals, sigma2=sigma2, noise_scale=noise_scale, rng=rng)
                else:
                    payload = make_raw_payload(vals, noise_scale=noise_scale, rng=rng)
                zf.update_crits(child_id, payload, t=now)
                update_children_s += perf_counter() - t0

            total_s = update_all_s + refresh_s + choose_parent_s + perform_new_s + update_children_s

            rows.append(
                {
                    "N": int(N),
                    "M": 3,
                    "n_samples": int(S),
                    "percentile": float(percentile),
                    "quantiles_mode_i": int(quantiles_mode_i),
                    "quantiles_budget": int(quantiles_budget),
                    "mode_i": int(mode_i),
                    "within_mode_i": int(within_mode_i),
                    "crits_source": str(crits_source),
                    "boundary": str(boundary),
                    "seed": int(seed),
                    "epoch": int(ep),
                    "n_active": int(len(rr.active_ids)),
                    "n_losers": int(len(losers)),
                    "init_update_ms": 1000.0 * float(init_update_s) if ep == 1 else 0.0,
                    "update_all_ms": 1000.0 * float(update_all_s),
                    "build_ms": 1000.0 * float(hooks.build_s),
                    "core_ms": 1000.0 * float(hooks.core_s),
                    "engine_ms": 1000.0 * float(engine_s),
                    "persist_ms": 1000.0 * float(persist_s),
                    "refresh_ms": 1000.0 * float(refresh_s),
                    "choose_parent_ms": 1000.0 * float(choose_parent_s),
                    "perform_new_ms": 1000.0 * float(perform_new_s),
                    "update_children_ms": 1000.0 * float(update_children_s),
                    "total_ms": 1000.0 * float(total_s),
                }
            )

    finally:
        # restore patched symbols
        zf._crits.build_matrices = orig_build  # type: ignore[attr-defined]
        selection_engine_mod.mc_rank_select_detailed = orig_core  # type: ignore[assignment]

    return rows


def _summarize(rows: list[dict[str, Any]], *, skip_first: int = 1) -> dict[str, float]:
    if not rows:
        return {}

    rows2 = rows[skip_first:] if len(rows) > skip_first else rows
    def col(name: str) -> list[float]:
        return [float(r[name]) for r in rows2]

    return {
        "engine_ms_med": median(col("engine_ms")),
        "build_ms_med": median(col("build_ms")),
        "core_ms_med": median(col("core_ms")),
        "persist_ms_med": median(col("persist_ms")),
        "refresh_ms_med": median(col("refresh_ms")),
        "total_ms_med": median(col("total_ms")),
        "n_losers_med": median(col("n_losers")),
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=str, default="out/bench_refresh.csv")

    ap.add_argument("--Ns", type=str, default="128,256", help="Comma-separated population sizes")
    ap.add_argument("--Ss", type=str, default="64,128,256,512", help="Comma-separated n_samples values")
    ap.add_argument("--epochs", type=int, default=25)
    ap.add_argument("--reps", type=int, default=3, help="Independent repetitions per setting")

    ap.add_argument("--percentile", type=float, default=0.2)
    ap.add_argument("--seed", type=int, default=0)

    ap.add_argument("--mode-i", type=int, default=1)
    ap.add_argument("--within-mode-i", type=int, default=0)
    ap.add_argument("--quantiles-mode-i", type=int, default=2)
    ap.add_argument("--quantiles-budget", type=int, default=2_500_000)

    ap.add_argument("--sigma2", type=float, default=0.01, help="Per-criterion variance for external stats")
    ap.add_argument("--noise-scale", type=float, default=0.0, help="Noise added to mu each update")

    ap.add_argument(
        "--crits-source",
        choices=["stats", "raw"],
        default="stats",
        help="Update criteria via external stats (fast) or raw OWRF (slower)",
    )
    ap.add_argument("--boundary", choices=["reflect", "clip"], default="reflect")

    ap.add_argument(
        "--skip-first",
        type=int,
        default=1,
        help="How many initial epochs to exclude from medians (warmup)",
    )

    args = ap.parse_args()

    Ns = parse_int_list(args.Ns)
    Ss = parse_int_list(args.Ss)

    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)

    print(
        "\n".join(
            [
                "ZenFronts bench_refresh",
                f"Python: {sys.version.split()[0]}",
                f"NumPy: {np.__version__}",
                f"Platform: {platform.platform()}",
                f"Ns={Ns}, Ss={Ss}, epochs={args.epochs}, reps={args.reps}",
                f"percentile={args.percentile}, crits_source={args.crits_source}, boundary={args.boundary}",
                "",
            ]
        )
    )

    all_rows: list[dict[str, Any]] = []

    for N in Ns:
        for S in Ss:
            summaries: list[dict[str, float]] = []
            for rep in range(int(args.reps)):
                rows = run_one(
                    N=N,
                    S=S,
                    percentile=float(args.percentile),
                    epochs=int(args.epochs),
                    seed=int(args.seed + rep),
                    quantiles_mode_i=int(args.quantiles_mode_i),
                    quantiles_budget=int(args.quantiles_budget),
                    mode_i=int(args.mode_i),
                    within_mode_i=int(args.within_mode_i),
                    sigma2=float(args.sigma2),
                    noise_scale=float(args.noise_scale),
                    crits_source=str(args.crits_source),
                    boundary=str(args.boundary),
                )
                for r in rows:
                    r["rep"] = int(rep)
                all_rows.extend(rows)
                summaries.append(_summarize(rows, skip_first=int(args.skip_first)))

            # print a compact summary per setting
            def sm(name: str) -> float:
                xs = [s[name] for s in summaries if name in s]
                return float(median(xs)) if xs else float("nan")

            print(
                f"N={N:4d} S={S:4d} | engine {sm('engine_ms_med'):8.2f} ms"
                f" (build {sm('build_ms_med'):7.2f} / core {sm('core_ms_med'):7.2f})"
                f" | refresh {sm('refresh_ms_med'):8.2f} | total {sm('total_ms_med'):8.2f}"
                f" | losers~{sm('n_losers_med'):.0f}"
            )

    # write csv
    if all_rows:
        fieldnames = list(all_rows[0].keys())
        with open(args.out, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=fieldnames)
            w.writeheader()
            w.writerows(all_rows)

    print(f"\nWrote: {args.out}")


if __name__ == "__main__":
    main()
