#!/usr/bin/env python3
"""Adaptive grid search over the simulator's `conf.toml` space.

For every point on a (view_dist x weapon-radius) grid this driver asks the
simulator for a *sequential* evaluation: seeds are run in batches and the run
stops as soon as the RED-BLU food-share difference is `--sigma` standard
errors away from zero, or `--max-samples` is hit. The per-point statistics are
appended to a JSONL file (resumable) and a flat CSV.

The grid is *adaptive*: after the coarse pass, cells whose four corners do not
agree on a winner are refined (centre + edge midpoints) with a stricter sigma,
so extra samples are spent where the phase boundary actually is.

Only one thing about the base config is required: red aggressive, blue
pacifist. `conf.toml` already satisfies that.

Examples
--------
    # coarse pass only
    python3 tools/sweep.py --view-dist 100:550:25 --radius 0:250:12.5

    # coarse + two refinement passes
    python3 tools/sweep.py --view-dist 100:550:50 --radius 0:250:25 --refine 2

    # see the grid without running anything
    python3 tools/sweep.py --dry-run
"""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import tempfile
import threading
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
DEFAULT_BINARY = REPO / "simulator" / "target" / "release" / "simulator"
DEFAULT_BASE = REPO / "simulator" / "conf.toml"


# --------------------------------------------------------------------------- #
# Grid specification
# --------------------------------------------------------------------------- #
def parse_axis(spec: str) -> list[float]:
    """Parse an axis spec into a sorted list of unique floats.

    Accepts either a comma/space separated list (`0,10,20`) or a
    `start:stop:step` range (`0:250:12.5`, inclusive of `stop` when it lands).
    """
    spec = spec.strip()
    if ":" in spec and "," not in spec:
        parts = [p.strip() for p in spec.split(":")]
        if len(parts) != 3:
            raise ValueError(f"range must be start:stop:step, got {spec!r}")
        start, stop, step = (float(p) for p in parts)
        if step == 0:
            raise ValueError("range step must be non-zero")
        out: list[float] = []
        v = start
        # Tolerate float drift when deciding whether `stop` is reached.
        eps = abs(step) * 1e-9
        if step > 0:
            while v <= stop + eps:
                out.append(round(v, 9))
                v += step
        else:
            while v >= stop - eps:
                out.append(round(v, 9))
                v += step
        return out
    return sorted({round(float(p), 9) for p in re.split(r"[,\s]+", spec) if p})


# --------------------------------------------------------------------------- #
# Config templating
# --------------------------------------------------------------------------- #
def render_config(base_text: str, overrides: dict[str, float | int | bool]) -> str:
    """Replace `key = value` lines in the TOML text for each override."""
    text = base_text
    for key, value in overrides.items():
        if isinstance(value, bool):
            rendered = "true" if value else "false"
        elif isinstance(value, int):
            rendered = str(value)
        else:
            rendered = repr(float(value))
        pattern = re.compile(rf"^(\s*){re.escape(key)}(\s*)=.*$", re.M)
        text, n = pattern.subn(rf"\g<1>{key}\g<2>= {rendered}", text)
        if n == 0:
            raise ValueError(f"key {key!r} not found in base config")
    return text


def parse_value(raw: str) -> float | int | bool:
    """Parse a `--set KEY=VALUE` value into a bool/int/float."""
    low = raw.strip().lower()
    if low in ("true", "false"):
        return low == "true"
    try:
        return int(raw)
    except ValueError:
        return float(raw)


def parse_set(items: list[str]) -> dict[str, float | int | bool]:
    out: dict[str, float | int | bool] = {}
    for item in items or []:
        if "=" not in item:
            raise ValueError(f"--set expects KEY=VALUE, got {item!r}")
        key, _, raw = item.partition("=")
        out[key.strip()] = parse_value(raw)
    return out


# --------------------------------------------------------------------------- #
# Sequential evaluation of a single grid point
# --------------------------------------------------------------------------- #
class Evaluator:
    def __init__(
        self,
        binary: Path,
        base_text: str,
        rayon_threads: int,
        timeout: float,
        extra: dict[str, float | int | bool] | None = None,
    ) -> None:
        self.binary = binary
        self.base_text = base_text
        self.rayon_threads = rayon_threads
        self.timeout = timeout
        self.extra = extra or {}

    def __call__(self, point: tuple[float, float, dict]) -> dict:
        view_dist, radius, opts = point
        overrides = {"view_dist": view_dist, "radius": radius, **self.extra}
        text = render_config(self.base_text, overrides)
        # Unique temp file per call so concurrent threads never race.
        fd, path = tempfile.mkstemp(suffix=".toml", dir=tempfile.gettempdir())
        try:
            with os.fdopen(fd, "w") as f:
                f.write(text)
            env = dict(os.environ)
            env["RAYON_NUM_THREADS"] = str(self.rayon_threads)
            cmd = [
                str(self.binary),
                "--sequential",
                "--config", path,
                "--sigma", str(opts["sigma"]),
                "--min-samples", str(opts["min_samples"]),
                "--max-samples", str(opts["max_samples"]),
                "--batch", str(opts["batch"]),
            ]
            proc = subprocess.run(
                cmd, capture_output=True, text=True, env=env, timeout=self.timeout
            )
        finally:
            try:
                os.unlink(path)
            except OSError:
                pass

        if proc.returncode != 0:
            raise RuntimeError(
                f"simulator failed for view_dist={view_dist} radius={radius}: "
                f"{proc.stderr.strip() or proc.stdout.strip()}"
            )

        line = next(
            (ln for ln in reversed(proc.stdout.splitlines()) if ln.strip().startswith("{")),
            None,
        )
        if line is None:
            raise RuntimeError(f"no JSON in simulator output:\n{proc.stdout}")
        result = json.loads(line)
        # Normalise to the grid coordinates we asked for (JSON round-trip of
        # the floats should match, but this makes the key exact).
        result["view_dist"] = view_dist
        result["weapon_radius"] = radius
        return result


# --------------------------------------------------------------------------- #
# Persistence
# --------------------------------------------------------------------------- #
def result_key(r: dict) -> tuple[float, float]:
    return (round(float(r["view_dist"]), 9), round(float(r["weapon_radius"]), 9))


def load_results(path: Path) -> dict[tuple[float, float], dict]:
    results: dict[tuple[float, float], dict] = {}
    if path.exists():
        with path.open() as f:
            for line in f:
                line = line.strip()
                if line:
                    r = json.loads(line)
                    results[result_key(r)] = r
    return results


def write_csv(path: Path, results: dict[tuple[float, float], dict]) -> None:
    if not results:
        return
    cols = [
        "view_dist", "weapon_radius", "n", "sigma", "winner", "stopped",
        "mean_red", "mean_blu", "mean_diff", "sd_diff", "se_diff",
        "red_wins", "blu_wins", "ties", "weapon_kills_user", "food_amount",
    ]
    rows = sorted(results.values(), key=result_key)
    with path.open("w", newline="") as f:
        import csv
        w = csv.DictWriter(f, fieldnames=cols, extrasaction="ignore")
        w.writeheader()
        for r in rows:
            w.writerow(r)


# --------------------------------------------------------------------------- #
# Adaptive refinement
# --------------------------------------------------------------------------- #
def outcome(r: dict) -> str:
    """Collapse a result to the label used for refinement decisions.

    A point that hit the sample cap without reaching sigma is reported as
    `tie`, so the boundary around genuinely undecided regions is refined too.
    """
    return "tie" if r.get("stopped") != "sigma" else r["winner"]


def find_refinement_points(
    results: dict[tuple[float, float], dict],
    vds: list[float],
    rads: list[float],
) -> set[tuple[float, float]]:
    """Midpoints of cells whose four corners disagree on the outcome."""
    def label(vd: float, rad: float) -> str | None:
        r = results.get((round(vd, 9), round(rad, 9)))
        return outcome(r) if r else None

    new: set[tuple[float, float]] = set()
    for i in range(len(vds) - 1):
        for j in range(len(rads) - 1):
            corners = [
                label(vds[i], rads[j]),
                label(vds[i + 1], rads[j]),
                label(vds[i], rads[j + 1]),
                label(vds[i + 1], rads[j + 1]),
            ]
            if None in corners or len(set(corners)) == 1:
                continue
            mid_vd = round((vds[i] + vds[i + 1]) / 2, 9)
            mid_rad = round((rads[j] + rads[j + 1]) / 2, 9)
            new.update({
                (mid_vd, rads[j]),
                (mid_vd, rads[j + 1]),
                (vds[i], mid_rad),
                (vds[i + 1], mid_rad),
                (mid_vd, mid_rad),
            })
    return new


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--binary", type=Path, default=DEFAULT_BINARY)
    ap.add_argument("--base", type=Path, default=DEFAULT_BASE,
                    help="base config; red must be aggressive and blue pacifist")
    ap.add_argument("--view-dist", default="100:550:25", help="axis spec (default 100:550:25)")
    ap.add_argument("--radius", default="0:250:12.5", help="axis spec (default 0:250:12.5)")
    ap.add_argument("--out", type=Path, default=Path("sweep_out/results.jsonl"))
    ap.add_argument("--sigma", type=float, default=3.0)
    ap.add_argument("--min-samples", type=int, default=32)
    ap.add_argument("--max-samples", type=int, default=2000)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--refine", type=int, default=1, help="number of adaptive refinement passes")
    ap.add_argument("--refine-sigma", type=float, default=3.0)
    ap.add_argument("--refine-max-samples", type=int, default=2500)
    ap.add_argument("--refine-batch", type=int, default=128)
    ap.add_argument("--seed-budget", type=int, default=4_000_000,
                    help="stop scheduling work once this many seeds have run (safety valve)")
    ap.add_argument("--workers", type=int, default=os.cpu_count() or 4)
    ap.add_argument("--rayon-threads", type=int, default=1,
                    help="threads the simulator may use per point (keep at 1: we parallelise over points)")
    ap.add_argument("--timeout", type=float, default=600.0)
    ap.add_argument("--set", action="append", default=[], metavar="KEY=VALUE",
                    help="pin any other conf.toml field (repeatable), e.g. --set kills_user=false")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    if not args.binary.exists():
        ap.error(f"binary not found: {args.binary} (build with `cargo build --release`)")
    base_text = args.base.read_text()

    vds = parse_axis(args.view_dist)
    rads = parse_axis(args.radius)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    csv_path = args.out.with_suffix(".csv")

    results = load_results(args.out)
    if results:
        print(f"resuming: {len(results)} points already in {args.out}", file=sys.stderr)

    if args.dry_run:
        print(f"view_dist ({len(vds)}): {vds}")
        print(f"radius    ({len(rads)}): {rads}")
        print(f"coarse grid: {len(vds) * len(rads)} points; "
              f"{args.refine} refinement pass(es); out={args.out}")
        return 0

    ev = Evaluator(args.binary, base_text, args.rayon_threads, args.timeout,
                   extra=parse_set(args.set))
    base_opts = dict(sigma=args.sigma, min_samples=args.min_samples,
                     max_samples=args.max_samples, batch=args.batch)

    out_lock = threading.Lock()
    out_fh = args.out.open("a")
    state = {"seeds": sum(int(r["n"]) for r in results.values())}
    if state["seeds"]:
        print(f"resuming: {state['seeds']} seeds already spent", file=sys.stderr)

    def evaluate(points: list[tuple[float, float]], opts: dict, label: str) -> bool:
        """Evaluate `points`, appending results as they finish.

        Returns False if the seed budget ran out (caller should stop refining).
        """
        todo = [
            (vd, rad, opts) for vd, rad in points
            if (round(vd, 9), round(rad, 9)) not in results
        ]
        if not todo:
            print(f"[{label}] nothing new to evaluate", flush=True)
            return True
        print(f"[{label}] evaluating {len(todo)} point(s) "
              f"(sigma={opts['sigma']}, max={opts['max_samples']}, "
              f"seeds={state['seeds']}/{args.seed_budget})", flush=True)
        chunk = max(args.workers * 2, 8)
        done = 0
        with ThreadPoolExecutor(max_workers=args.workers) as ex:
            for i in range(0, len(todo), chunk):
                if state["seeds"] >= args.seed_budget:
                    print(f"[{label}] seed budget exhausted; stopping", flush=True)
                    return False
                for r in ex.map(ev, todo[i:i + chunk]):
                    key = result_key(r)
                    with out_lock:
                        results[key] = r
                        out_fh.write(json.dumps(r) + "\n")
                        out_fh.flush()
                        state["seeds"] += int(r["n"])
                    done += 1
                    print(f"  [{done}/{len(todo)}] vd={r['view_dist']:>6} "
                          f"rad={r['weapon_radius']:>6} n={r['n']:>5} "
                          f"sigma={r['sigma']:.2f} -> {outcome(r)} ({r['stopped']})",
                          flush=True)
        return True

    # Coarse pass.
    evaluate([(vd, rad) for vd in vds for rad in rads], base_opts, "coarse")

    # Adaptive refinement passes.
    for p in range(args.refine):
        vds = sorted({k[0] for k in results})
        rads = sorted({k[1] for k in results})
        new = find_refinement_points(results, vds, rads)
        if not new:
            print("no boundary cells left to refine")
            break
        opts = dict(sigma=args.refine_sigma, min_samples=args.min_samples,
                    max_samples=args.refine_max_samples, batch=args.refine_batch)
        if not evaluate(sorted(new), opts, f"refine {p + 1}"):
            break

    out_fh.close()
    write_csv(csv_path, results)

    red = sum(1 for r in results.values() if outcome(r) == "red")
    blu = sum(1 for r in results.values() if outcome(r) == "blu")
    tie = sum(1 for r in results.values() if outcome(r) == "tie")
    print(f"\ndone: {len(results)} points -> red={red} blu={blu} tie={tie}")
    print(f"wrote {args.out} and {csv_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())