"""Parameter optimization for LAMPSUI driven via the REST API.

The script doesn't know anything about LAMMPS internals — it submits runs,
polls until they finish, then reduces the run's thermo trace to a single
scalar that Optuna optimizes over.

Because the bridge is the REST API, you can swap Optuna for any other
optimizer (BoTorch, scikit-optimize, your own Bayesian loop) without
touching LAMPSUI itself.

Quick start:

    pip install -r scripts/requirements.txt

    # In one terminal: start LAMPSUI
    docker run --rm -p 7860:7860 -v "$(pwd)/work:/work" lampsui

    # In another terminal: run a 15-trial sweep
    python scripts/optimize.py \\
        --example test_100 \\
        --space scripts/space_example.json \\
        --metric Press --reduce abs_last \\
        --direction minimize --n-trials 15

The search space JSON has one entry per parameter to vary:

    {
      "dh":     {"type": "float", "low": 0.15, "high": 0.30},
      "sigmao": {"type": "float", "low": 0.5,  "high": 2.0,  "log": true},
      "F":      {"type": "int",   "low": 1,    "high": 3},
      "randPos":{"type": "categorical", "choices": [0, 1, 2]}
    }
"""
from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
from pathlib import Path

import optuna
import requests


def submit_and_wait(
    api: str, example: str, overrides: dict, poll_s: float = 2.0
) -> tuple[str, str]:
    """Submit a run, poll until it terminates. Returns (run_id, final_status)."""
    payload = {
        "example": example,
        "overrides": overrides,
        "image": {"enabled": False},  # images off — we don't need them for objective
    }
    r = requests.post(f"{api}/runs", json=payload, timeout=30)
    r.raise_for_status()
    rid = r.json()["id"]
    while True:
        s = requests.get(f"{api}/runs/{rid}", timeout=30).json()
        if s["status"] in ("done", "failed", "cancelled"):
            return rid, s["status"]
        time.sleep(poll_s)


def reduce_thermo(
    api: str, rid: str, column: str, reduce: str, target: float | None = None
) -> float:
    """Fetch thermo data and reduce one column to a scalar.

    For mae / mape, `target` must be supplied (the reference value to compare
    the column against). MAPE is returned as a percentage (e.g. 12.34 means
    12.34%).
    """
    t = requests.get(f"{api}/runs/{rid}/thermo", timeout=30).json()
    if not t["rows"]:
        raise RuntimeError(f"run {rid}: no thermo data")
    if column not in t["columns"]:
        raise RuntimeError(
            f"run {rid}: column {column!r} not in {t['columns']}"
        )
    idx = t["columns"].index(column)
    vals = [row[idx] for row in t["rows"]]

    if reduce == "last":
        return vals[-1]
    if reduce == "abs_last":
        return abs(vals[-1])
    if reduce == "mean":
        return sum(vals) / len(vals)
    if reduce == "abs_mean":
        return abs(sum(vals) / len(vals))
    if reduce == "stability":
        tail = vals[len(vals) // 2 :]
        return statistics.pstdev(tail) if len(tail) >= 2 else 0.0
    if reduce == "mae":
        if target is None:
            raise ValueError("reduce=mae requires --target")
        return sum(abs(v - target) for v in vals) / len(vals)
    if reduce == "mape":
        if target is None:
            raise ValueError("reduce=mape requires --target")
        if target == 0.0:
            raise ValueError("MAPE is undefined when target=0")
        return (
            100.0 * sum(abs(v - target) for v in vals) / len(vals) / abs(target)
        )
    raise ValueError(f"unknown reduce: {reduce}")


def suggest(trial: optuna.Trial, name: str, spec: dict):
    t = spec["type"]
    if t == "float":
        return trial.suggest_float(
            name, spec["low"], spec["high"], log=spec.get("log", False)
        )
    if t == "int":
        return trial.suggest_int(name, spec["low"], spec["high"])
    if t == "categorical":
        return trial.suggest_categorical(name, spec["choices"])
    raise ValueError(f"unsupported parameter type: {t!r}")


def main() -> int:
    p = argparse.ArgumentParser(
        description="Optuna parameter sweep driving LAMPSUI through its REST API.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--api", default="http://localhost:7860/api")
    p.add_argument("--example", required=True, help="Example name (e.g. test_100)")
    p.add_argument(
        "--space",
        required=True,
        help="Path to JSON file describing the search space",
    )
    p.add_argument(
        "--fixed",
        default="{}",
        help="JSON dict of fixed overrides applied to every trial",
    )
    p.add_argument("--n-trials", type=int, default=15)
    p.add_argument("--metric", default="Press", help="Thermo column to reduce")
    p.add_argument(
        "--reduce",
        default="abs_last",
        choices=["last", "abs_last", "mean", "abs_mean", "stability", "mae", "mape"],
        help="How to reduce the column to a single scalar",
    )
    p.add_argument(
        "--target",
        type=float,
        default=None,
        help="Reference value for mae / mape (required for those reducers)",
    )
    p.add_argument(
        "--direction", default="minimize", choices=["minimize", "maximize"]
    )
    p.add_argument(
        "--study",
        help="Study name. If set, results persist to <name>.db (SQLite) so the "
        "study can be resumed.",
    )
    p.add_argument(
        "--sampler",
        default="tpe",
        choices=["tpe", "random", "cmaes"],
        help="Optuna sampler to use",
    )
    p.add_argument(
        "--seed", type=int, default=None, help="Sampler seed for reproducibility"
    )
    args = p.parse_args()

    space_path = Path(args.space)
    if not space_path.is_file():
        print(f"error: space file not found: {space_path}", file=sys.stderr)
        return 2
    space = json.loads(space_path.read_text())
    fixed = json.loads(args.fixed)

    # Sanity: poke the API once
    try:
        requests.get(f"{args.api}/health", timeout=5).raise_for_status()
    except Exception as e:
        print(
            f"error: cannot reach LAMPSUI API at {args.api} ({e})\n"
            "Make sure the container is running on the same host.",
            file=sys.stderr,
        )
        return 2

    def objective(trial: optuna.Trial) -> float:
        overrides = dict(fixed)
        for name, spec in space.items():
            overrides[name] = suggest(trial, name, spec)
        rid, status = submit_and_wait(args.api, args.example, overrides)
        if status != "done":
            print(f"  trial {trial.number}: run {rid} → {status} (pruned)")
            raise optuna.TrialPruned()
        val = reduce_thermo(args.api, rid, args.metric, args.reduce, args.target)
        suffix = "%" if args.reduce == "mape" else ""
        print(
            f"  trial {trial.number:3d}: run {rid}  "
            f"{args.metric}/{args.reduce}={val:.6g}{suffix}  overrides={overrides}"
        )
        # Tag the LAMPSUI run with the trial number so it's easy to find later
        try:
            requests.patch(
                f"{args.api}/runs/{rid}",
                json={
                    "name": (
                        f"{args.study or 'opt'} #{trial.number} "
                        f"({args.metric}/{args.reduce}={val:.4g})"
                    )[:80]
                },
                timeout=10,
            )
        except Exception:
            pass
        return val

    samplers = {
        "tpe": optuna.samplers.TPESampler(seed=args.seed),
        "random": optuna.samplers.RandomSampler(seed=args.seed),
        "cmaes": optuna.samplers.CmaEsSampler(seed=args.seed),
    }
    storage = f"sqlite:///{args.study}.db" if args.study else None
    study = optuna.create_study(
        study_name=args.study,
        direction=args.direction,
        sampler=samplers[args.sampler],
        storage=storage,
        load_if_exists=True,
    )

    print(
        f"Optimising {args.example} via {args.api}\n"
        f"  metric:    {args.metric} reduced by {args.reduce} ({args.direction})\n"
        f"  search:    {list(space.keys())}\n"
        f"  fixed:     {fixed or '{}'}\n"
        f"  trials:    {args.n_trials}\n"
        f"  sampler:   {args.sampler}\n"
        + (f"  storage:   {storage}\n" if storage else "")
    )
    study.optimize(objective, n_trials=args.n_trials)

    print()
    print(f"Best {args.direction}d value: {study.best_value:.6g}")
    print(f"Best params: {study.best_params}")
    if storage:
        print(f"\nResume / inspect later with:\n"
              f"  optuna-dashboard {storage}")
    return 0


if __name__ == "__main__":
    sys.exit(main())