backend/optimization.py

"""In-process Optuna study manager (single- and multi-objective).

Each trial reuses the existing RunManager (so trials show up in the runs
history alongside manual runs).

Multi-objective:
  - The study config can carry a list of objectives, each with its own
    metric / reduce / direction.
  - Optuna's NSGAIISampler / NSGAIIISampler are used for >=2 objectives
    (TPE is allowed too, but NSGA is the typical choice).
  - When n_objectives >= 2, `best_value`/`best_params` are replaced by
    `pareto_trials` — the non-dominated set.

Persistence:
  /work/studies/<id>/_meta.json   our wrapper state
  /work/studies/<id>/study.db     Optuna's own SQLite storage
"""
from __future__ import annotations

import asyncio
import json
import os
import statistics
import time
import uuid
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Optional

import optuna
from optuna.samplers import (
    CmaEsSampler,
    NSGAIISampler,
    NSGAIIISampler,
    RandomSampler,
    TPESampler,
)

from .params import schema_for_example
from .simulator import EXAMPLES_DIR, manager as run_manager
from .thermo import parse_thermo

optuna.logging.set_verbosity(optuna.logging.WARNING)

STUDIES_DIR = Path(os.environ.get("STUDIES_DIR", "/work/studies"))
STUDIES_DIR.mkdir(parents=True, exist_ok=True)

REDUCERS = ("last", "abs_last", "mean", "abs_mean", "stability", "mae", "mape")
TARGET_REDUCERS = ("mae", "mape")  # require an explicit `target` value
DIRECTIONS = ("minimize", "maximize")
SAMPLERS = ("tpe", "random", "cmaes", "nsga2", "nsga3")


@dataclass
class Objective:
    metric: str
    reduce: str
    direction: str
    # Required for mae/mape: the value to compare the thermo column against.
    # Ignored by the other reducers.
    target: Optional[float] = None


@dataclass
class TrialResult:
    number: int
    run_id: str
    overrides: dict[str, Any]
    values: Optional[list[float]]  # one entry per objective; None if pruned/failed
    status: str  # "running" | "done" | "failed" | "cancelled"


@dataclass
class StudyState:
    id: str
    name: str
    example: str
    fixed: dict[str, Any]
    space: dict[str, dict[str, Any]]
    objectives: list[Objective]
    n_trials: int
    sampler: str
    status: str = "queued"
    started_at: float = 0.0
    finished_at: float = 0.0
    trials: list[TrialResult] = field(default_factory=list)
    pareto_trials: list[dict[str, Any]] = field(default_factory=list)
    cancel_requested: bool = False
    error: Optional[str] = None

    def to_public(self) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "example": self.example,
            "fixed": self.fixed,
            "space": self.space,
            "objectives": [asdict(o) for o in self.objectives],
            "n_trials": self.n_trials,
            "sampler": self.sampler,
            "status": self.status,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
            "trials": [asdict(t) for t in self.trials],
            "pareto_trials": self.pareto_trials,
            "n_objectives": len(self.objectives),
            "n_completed": sum(1 for t in self.trials if t.values is not None),
            "error": self.error,
        }


def _coerce_objectives(cfg: dict) -> list[Objective]:
    """Accept either the new list-of-objectives shape or legacy single-objective."""
    if "objectives" in cfg and cfg["objectives"]:
        objs = cfg["objectives"]
    else:
        objs = [
            {
                "metric": cfg.get("metric", "Press"),
                "reduce": cfg.get("reduce", "abs_last"),
                "direction": cfg.get("direction", "minimize"),
            }
        ]
    out: list[Objective] = []
    for o in objs:
        if o.get("reduce") not in REDUCERS:
            raise ValueError(f"reduce must be one of {REDUCERS}")
        if o.get("direction") not in DIRECTIONS:
            raise ValueError("direction must be minimize|maximize")
        if not o.get("metric"):
            raise ValueError("Each objective needs a metric (thermo column)")
        target_raw = o.get("target")
        target_val: Optional[float] = None
        if o["reduce"] in TARGET_REDUCERS:
            if target_raw is None or target_raw == "":
                raise ValueError(
                    f"reduce={o['reduce']!r} requires a 'target' value "
                    "(the reference the thermo column is compared against)."
                )
            try:
                target_val = float(target_raw)
            except (TypeError, ValueError):
                raise ValueError(f"target must be numeric, got {target_raw!r}")
            if o["reduce"] == "mape" and target_val == 0.0:
                raise ValueError(
                    "MAPE is undefined when target=0 (division by zero). "
                    "Use MAE instead, or pick a non-zero target."
                )
        elif target_raw not in (None, ""):
            try:
                target_val = float(target_raw)
            except (TypeError, ValueError):
                target_val = None
        out.append(
            Objective(
                metric=str(o["metric"]),
                reduce=str(o["reduce"]),
                direction=str(o["direction"]),
                target=target_val,
            )
        )
    return out


def _validate_config(cfg: dict) -> tuple[dict, list[Objective]]:
    for key in ("name", "example", "space", "n_trials"):
        if key not in cfg:
            raise ValueError(f"Missing field: {key!r}")
    objectives = _coerce_objectives(cfg)
    sampler = cfg.get("sampler", "tpe")
    if sampler not in SAMPLERS:
        raise ValueError(f"sampler must be one of {SAMPLERS}")
    if sampler in ("nsga2", "nsga3") and len(objectives) < 2:
        raise ValueError(
            f"{sampler} requires at least 2 objectives "
            "(use TPE/Random for single-objective studies)"
        )
    if sampler == "cmaes" and len(objectives) > 1:
        raise ValueError("CMA-ES does not support multi-objective optimization")
    if not isinstance(cfg["space"], dict) or not cfg["space"]:
        raise ValueError("space must be a non-empty dict")

    # Build the allowlist from the chosen example's input file — same
    # source-of-truth as the run-submit path.
    example_dir = EXAMPLES_DIR / cfg["example"]
    input_file = next(example_dir.glob("in.*"), None) if example_dir.is_dir() else None
    if input_file is None:
        raise ValueError(f"Unknown example or missing input file: {cfg['example']}")
    allowed = {p["name"]: p["type"] for p in schema_for_example(input_file)}

    for name, spec in cfg["space"].items():
        if name not in allowed:
            raise ValueError(
                f"Unknown parameter for example {cfg['example']!r}: {name!r}"
            )
        t = spec.get("type")
        if t in ("float", "int"):
            if "low" not in spec or "high" not in spec:
                raise ValueError(f"{name}: low/high required for {t} parameters")
            if spec["high"] <= spec["low"]:
                raise ValueError(f"{name}: high must be > low")
        elif t == "categorical":
            if not spec.get("choices"):
                raise ValueError(f"{name}: choices required for categorical")
        else:
            raise ValueError(f"{name}: type must be float|int|categorical")
    n = int(cfg["n_trials"])
    if n < 1 or n > 500:
        raise ValueError("n_trials must be between 1 and 500")
    return cfg, objectives


def _reduce_metric(
    log_lines: list[str],
    column: str,
    reduce: str,
    target: Optional[float] = None,
) -> Optional[float]:
    t = parse_thermo(log_lines)
    if not t["rows"] or column not in t["columns"]:
        return None
    idx = t["columns"].index(column)
    vals = [r[idx] for r 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:
            return None
        return sum(abs(v - target) for v in vals) / len(vals)
    if reduce == "mape":
        if target is None or target == 0.0:
            return None
        # Returned as a percentage so 12.34 means "12.34% mean absolute error".
        return 100.0 * sum(abs(v - target) for v in vals) / len(vals) / abs(target)
    return None


def _build_sampler(name: str, n_obj: int):
    if name == "tpe":
        # TPE works for both single and multi-objective
        return TPESampler(multivariate=n_obj >= 2)
    if name == "random":
        return RandomSampler()
    if name == "cmaes":
        return CmaEsSampler()
    if name == "nsga2":
        return NSGAIISampler()
    if name == "nsga3":
        return NSGAIIISampler()
    raise ValueError(f"unknown sampler: {name}")


class StudyManager:
    def __init__(self) -> None:
        self.studies: dict[str, StudyState] = {}
        self._load_existing()

    def _load_existing(self) -> None:
        if not STUDIES_DIR.exists():
            return
        for d in sorted(STUDIES_DIR.iterdir()):
            if not d.is_dir():
                continue
            meta = d / "_meta.json"
            if not meta.is_file():
                continue
            try:
                m = json.loads(meta.read_text())
            except (json.JSONDecodeError, OSError):
                continue
            status = m.get("status", "failed")
            if status in ("queued", "running"):
                status = "failed"
            try:
                # Back-compat: old _meta.json had metric/reduce/direction at top
                if "objectives" in m and m["objectives"]:
                    objs = [Objective(**o) for o in m["objectives"]]
                else:
                    objs = [
                        Objective(
                            metric=m.get("metric", "Press"),
                            reduce=m.get("reduce", "abs_last"),
                            direction=m.get("direction", "minimize"),
                        )
                    ]
                trials = []
                for t in m.get("trials", []):
                    # old shape: 'value': float | None
                    vals = t.get("values")
                    if vals is None and "value" in t:
                        vals = None if t["value"] is None else [t["value"]]
                    trials.append(
                        TrialResult(
                            number=t["number"],
                            run_id=t.get("run_id", ""),
                            overrides=t.get("overrides", {}),
                            values=vals,
                            status=t.get("status", "failed"),
                        )
                    )
                state = StudyState(
                    id=m["id"],
                    name=m.get("name", m["id"]),
                    example=m.get("example", ""),
                    fixed=m.get("fixed") or {},
                    space=m.get("space") or {},
                    objectives=objs,
                    n_trials=int(m.get("n_trials", 0)),
                    sampler=m.get("sampler", "tpe"),
                    status=status,
                    started_at=float(m.get("started_at", 0.0)),
                    finished_at=float(m.get("finished_at", 0.0)),
                    trials=trials,
                    pareto_trials=m.get("pareto_trials") or [],
                    error=m.get("error"),
                )
                self.studies[state.id] = state
            except (KeyError, TypeError):
                continue

    def list(self) -> list[StudyState]:
        return sorted(self.studies.values(), key=lambda s: -s.started_at)

    def get(self, sid: str) -> Optional[StudyState]:
        return self.studies.get(sid)

    def is_anything_running(self) -> bool:
        return any(s.status == "running" for s in self.studies.values())

    def cancel(self, sid: str) -> bool:
        s = self.studies.get(sid)
        if s is None or s.status != "running":
            return False
        s.cancel_requested = True
        return True

    def write_meta(self, state: StudyState) -> None:
        d = STUDIES_DIR / state.id
        d.mkdir(parents=True, exist_ok=True)
        try:
            (d / "_meta.json").write_text(json.dumps(state.to_public()))
        except OSError:
            pass

    async def start(self, raw_config: dict) -> StudyState:
        cfg, objectives = _validate_config(raw_config)
        if self.is_anything_running():
            raise ValueError("Another optimization study is already running")
        sid = uuid.uuid4().hex[:8]
        state = StudyState(
            id=sid,
            name=str(cfg["name"]).strip()[:80] or f"study-{sid}",
            example=cfg["example"],
            fixed=cfg.get("fixed") or {},
            space=cfg["space"],
            objectives=objectives,
            n_trials=int(cfg["n_trials"]),
            sampler=cfg.get("sampler", "tpe"),
            status="running",
            started_at=time.time(),
        )
        self.studies[sid] = state
        self.write_meta(state)
        asyncio.create_task(self._run_study(state))
        return state

    async def _run_study(self, state: StudyState) -> None:
        try:
            n_obj = len(state.objectives)
            sampler = _build_sampler(state.sampler, n_obj)
            db_path = STUDIES_DIR / state.id / "study.db"
            db_path.parent.mkdir(parents=True, exist_ok=True)
            create_kwargs: dict[str, Any] = {
                "study_name": state.name,
                "sampler": sampler,
                "storage": f"sqlite:///{db_path}",
                "load_if_exists": True,
            }
            if n_obj == 1:
                create_kwargs["direction"] = state.objectives[0].direction
            else:
                create_kwargs["directions"] = [o.direction for o in state.objectives]
            ostudy = optuna.create_study(**create_kwargs)

            for n in range(state.n_trials):
                if state.cancel_requested:
                    break
                trial = ostudy.ask()
                overrides: dict[str, Any] = dict(state.fixed)
                for name, spec in state.space.items():
                    t = spec["type"]
                    step = spec.get("step")
                    if t == "float":
                        log = bool(spec.get("log", False))
                        # Optuna disallows step + log together; drop step in that case.
                        kw: dict[str, Any] = {"log": log}
                        if step is not None and not log and float(step) > 0:
                            kw["step"] = float(step)
                        overrides[name] = trial.suggest_float(
                            name, float(spec["low"]), float(spec["high"]), **kw
                        )
                    elif t == "int":
                        kw_i: dict[str, Any] = {}
                        if step is not None and int(step) > 0:
                            kw_i["step"] = int(step)
                        overrides[name] = trial.suggest_int(
                            name, int(spec["low"]), int(spec["high"]), **kw_i
                        )
                    elif t == "categorical":
                        overrides[name] = trial.suggest_categorical(
                            name, spec["choices"]
                        )

                tr = TrialResult(
                    number=n,
                    run_id="",
                    overrides=overrides,
                    values=None,
                    status="running",
                )
                state.trials.append(tr)
                self.write_meta(state)

                try:
                    run = run_manager.submit(
                        state.example, overrides=overrides, image={"enabled": False}
                    )
                    tr.run_id = run.id
                    run.name = f"{state.name} #{n}"
                    run.write_meta()
                    await run_manager.start(run)
                    while run.status in ("queued", "running"):
                        if state.cancel_requested and run.status == "running":
                            run_manager.cancel(run)
                        await asyncio.sleep(1)

                    tr.status = run.status
                    if run.status == "done":
                        # Reduce each objective's metric → list of floats
                        values: list[float] = []
                        all_ok = True
                        for o in state.objectives:
                            v = _reduce_metric(run.log, o.metric, o.reduce, o.target)
                            if v is None:
                                all_ok = False
                                break
                            values.append(v)
                        if not all_ok:
                            ostudy.tell(
                                trial,
                                state=optuna.trial.TrialState.PRUNED,
                            )
                        else:
                            tr.values = values
                            ostudy.tell(trial, values if n_obj > 1 else values[0])
                            tag_parts = [
                                f"{o.metric}/{o.reduce}={v:.4g}"
                                for o, v in zip(state.objectives, values)
                            ]
                            run.name = (
                                f"{state.name} #{n} (" + ", ".join(tag_parts) + ")"
                            )[:80]
                            run.write_meta()
                    else:
                        ostudy.tell(trial, state=optuna.trial.TrialState.PRUNED)
                except Exception as e:
                    tr.status = "failed"
                    state.error = str(e)
                    try:
                        ostudy.tell(trial, state=optuna.trial.TrialState.FAIL)
                    except Exception:
                        pass

                # Update pareto / best tracking
                state.pareto_trials = self._collect_pareto(ostudy, state)
                self.write_meta(state)

            state.status = "cancelled" if state.cancel_requested else "done"
        except Exception as e:
            state.status = "failed"
            state.error = str(e)
        finally:
            state.finished_at = time.time()
            state.pareto_trials = self._collect_pareto_safe(state)
            self.write_meta(state)

    def _collect_pareto_safe(self, state: StudyState) -> list[dict[str, Any]]:
        try:
            db_path = STUDIES_DIR / state.id / "study.db"
            if not db_path.is_file():
                return state.pareto_trials
            ostudy = optuna.load_study(
                study_name=state.name, storage=f"sqlite:///{db_path}"
            )
            return self._collect_pareto(ostudy, state)
        except Exception:
            return state.pareto_trials

    def get_report(self, sid: str) -> Optional[dict[str, Any]]:
        """Build an analysis report for a study: importance, history, slices, correlations.

        Returns None if the study doesn't exist. Returns a report with empty
        sections if there aren't enough completed trials yet.
        """
        state = self.studies.get(sid)
        if state is None:
            return None
        n_obj = len(state.objectives)
        empty = {
            "n_completed": 0,
            "importances": [],
            "history": [],
            "correlations": [],
            "slices": {},
        }
        db_path = STUDIES_DIR / state.id / "study.db"
        if not db_path.is_file():
            return empty
        try:
            ostudy = optuna.load_study(
                study_name=state.name, storage=f"sqlite:///{db_path}"
            )
        except Exception:
            return empty

        completed = [
            t
            for t in ostudy.trials
            if t.state == optuna.trial.TrialState.COMPLETE and t.values is not None
        ]
        n = len(completed)
        if n < 2:
            return {**empty, "n_completed": n}

        param_names = list(state.space.keys())

        def value_for(t: "optuna.trial.FrozenTrial", obj_idx: int) -> Optional[float]:
            try:
                if n_obj == 1:
                    return float(t.value) if t.value is not None else None
                return float(t.values[obj_idx])
            except (TypeError, IndexError, ValueError):
                return None

        # ── Parameter importance per objective ──
        importances: list[dict[str, Any]] = []
        for i, obj in enumerate(state.objectives):
            try:
                imp_kwargs: dict[str, Any] = {}
                if n_obj > 1:
                    imp_kwargs["target"] = lambda t, ix=i: t.values[ix]
                imp = optuna.importance.get_param_importances(
                    ostudy, **imp_kwargs
                )
                importances.append(
                    {
                        "objective": f"{obj.metric}/{obj.reduce}",
                        "direction": obj.direction,
                        "scores": [
                            {"param": k, "score": float(v)} for k, v in imp.items()
                        ],
                    }
                )
            except Exception as e:  # importance can fail with too few trials / variants
                importances.append(
                    {
                        "objective": f"{obj.metric}/{obj.reduce}",
                        "direction": obj.direction,
                        "scores": [],
                        "error": str(e),
                    }
                )

        # ── Best-so-far history per objective ──
        history: list[dict[str, Any]] = []
        for i, obj in enumerate(state.objectives):
            is_min = obj.direction == "minimize"
            best = float("inf") if is_min else float("-inf")
            points: list[dict[str, Any]] = []
            for t in sorted(completed, key=lambda x: x.number):
                v = value_for(t, i)
                if v is None:
                    continue
                if (is_min and v < best) or ((not is_min) and v > best):
                    best = v
                points.append({"trial": t.number, "value": v, "best": best})
            history.append(
                {
                    "objective": f"{obj.metric}/{obj.reduce}",
                    "direction": obj.direction,
                    "points": points,
                }
            )

        # ── Pearson correlation between each (param, objective) ──
        correlations: list[dict[str, Any]] = []
        for i, obj in enumerate(state.objectives):
            for pname in param_names:
                xs: list[float] = []
                ys: list[float] = []
                for t in completed:
                    v = value_for(t, i)
                    p = t.params.get(pname)
                    if v is None or p is None:
                        continue
                    try:
                        xs.append(float(p))
                        ys.append(float(v))
                    except (TypeError, ValueError):
                        continue
                if len(xs) < 3:
                    continue
                r = _pearson(xs, ys)
                correlations.append(
                    {
                        "objective": f"{obj.metric}/{obj.reduce}",
                        "param": pname,
                        "r": r,
                        "n": len(xs),
                    }
                )

        # ── Slice scatter data per param per objective ──
        slices: dict[str, list[dict[str, Any]]] = {}
        for pname in param_names:
            per_obj: list[dict[str, Any]] = []
            for i, obj in enumerate(state.objectives):
                points = []
                for t in completed:
                    v = value_for(t, i)
                    p = t.params.get(pname)
                    if v is None or p is None:
                        continue
                    try:
                        points.append(
                            {"x": float(p), "y": v, "trial": t.number}
                        )
                    except (TypeError, ValueError):
                        continue
                per_obj.append(
                    {
                        "objective": f"{obj.metric}/{obj.reduce}",
                        "direction": obj.direction,
                        "points": points,
                    }
                )
            slices[pname] = per_obj

        return {
            "n_completed": n,
            "importances": importances,
            "history": history,
            "correlations": correlations,
            "slices": slices,
        }

    def _collect_pareto(
        self, ostudy: optuna.Study, state: StudyState
    ) -> list[dict[str, Any]]:
        try:
            if len(state.objectives) == 1:
                bt = ostudy.best_trial
                if bt is None or bt.value is None:
                    return []
                return [
                    {
                        "number": bt.number,
                        "values": [float(bt.value)],
                        "params": dict(bt.params),
                    }
                ]
            best_trials = ostudy.best_trials  # Pareto front
            return [
                {
                    "number": bt.number,
                    "values": [float(v) for v in bt.values],
                    "params": dict(bt.params),
                }
                for bt in best_trials
                if bt.values is not None
            ]
        except (ValueError, AttributeError):
            return []


def _pearson(xs: list[float], ys: list[float]) -> float:
    n = len(xs)
    if n < 2:
        return 0.0
    mx = sum(xs) / n
    my = sum(ys) / n
    sxy = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    sxx = sum((x - mx) ** 2 for x in xs)
    syy = sum((y - my) ** 2 for y in ys)
    if sxx == 0 or syy == 0:
        return 0.0
    return sxy / (sxx * syy) ** 0.5


study_manager = StudyManager()