#!/usr/bin/env python3
"""Local Optuna Dashboard with multi-objective plot fixes.

Stock optuna-dashboard calls optuna.visualization.plot_* without ``target``.
Multi-objective studies (this project: 3 maximize objectives) raise:

  ValueError: If the `study` is being used for multi-objective optimization,
  please specify the `target`.

This launcher monkey-patches those plot helpers so, when target is omitted on a
multi-objective study, objective index 0 is used (override with env
OPTUNA_DASHBOARD_PLOT_OBJECTIVE, default 0).
"""
from __future__ import annotations

import argparse
import functools
import os
import sys
from collections.abc import Callable, Sequence
from socketserver import ThreadingMixIn
from typing import Any, TypeVar
from wsgiref.simple_server import WSGIServer, make_server

import optuna
from optuna import Study
from optuna.trial import FrozenTrial
from optuna_dashboard import wsgi

F = TypeVar("F", bound=Callable[..., Any])

# Matches novelai_preset_optimize objectives (all maximize):
#   0: mean(stylo, 256 emb)
#   1: mean(stylo, 32 emb)
#   2: 32 emb only
_DEFAULT_OBJECTIVE_NAMES = (
    "mean(stylo, 256 emb)",
    "mean(stylo, 32 emb)",
    "32 emb",
)


def _plot_objective_index() -> int:
    raw = os.environ.get("OPTUNA_DASHBOARD_PLOT_OBJECTIVE", "0")
    try:
        return max(0, int(raw))
    except ValueError:
        return 0


def _as_studies(study: Study | Sequence[Study]) -> list[Study]:
    if isinstance(study, Study):
        return [study]
    return list(study)


def _needs_target(study: Study | Sequence[Study], target: Any) -> bool:
    if target is not None:
        return False
    return any(s._is_multi_objective() for s in _as_studies(study))


def _target_for(study: Study | Sequence[Study], idx: int) -> Callable[[FrozenTrial], float]:
    n = min(len(s.directions) for s in _as_studies(study))
    if n <= 0:
        idx = 0
    else:
        idx = min(idx, n - 1)

    def target(trial: FrozenTrial, _i: int = idx) -> float:
        if trial.values is None or _i >= len(trial.values):
            return float("nan")
        return float(trial.values[_i])

    return target


def _target_name_for(study: Study | Sequence[Study], idx: int, current: str) -> str:
    if current != "Objective Value":
        return current
    studies = _as_studies(study)
    n = min(len(s.directions) for s in studies)
    idx = min(idx, max(0, n - 1))
    if n == len(_DEFAULT_OBJECTIVE_NAMES):
        return _DEFAULT_OBJECTIVE_NAMES[idx]
    return f"Objective {idx}"


def _patch_mo_plot(plot_fn: F) -> F:
    """Default target=values[obj] for multi-objective when target is omitted."""

    @functools.wraps(plot_fn)
    def wrapper(
        study: Study | Sequence[Study],
        *args: Any,
        target: Callable[[FrozenTrial], float] | None = None,
        target_name: str = "Objective Value",
        **kwargs: Any,
    ) -> Any:
        if _needs_target(study, target):
            idx = _plot_objective_index()
            target = _target_for(study, idx)
            target_name = _target_name_for(study, idx, target_name)
        return plot_fn(study, *args, target=target, target_name=target_name, **kwargs)

    return wrapper  # type: ignore[return-value]


def patch_multi_objective_plots() -> None:
    """Patch optuna.visualization plot helpers used by optuna-dashboard."""
    import optuna.visualization as viz

    for name in (
        "plot_edf",
        "plot_contour",
        "plot_slice",
        "plot_parallel_coordinate",
        "plot_rank",
        "plot_param_importances",
    ):
        fn = getattr(viz, name, None)
        if fn is None:
            continue
        setattr(viz, name, _patch_mo_plot(fn))


class ThreadedWSGIServer(ThreadingMixIn, WSGIServer):
    daemon_threads = True


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Local Optuna Dashboard (multi-objective plot-safe)."
    )
    parser.add_argument(
        "storage",
        help="Storage URL (e.g. sqlite:////path/to/presets.db)",
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8080)
    args = parser.parse_args(argv)

    patch_multi_objective_plots()
    app = wsgi(args.storage)

    obj = _plot_objective_index()
    label = (
        _DEFAULT_OBJECTIVE_NAMES[obj]
        if 0 <= obj < len(_DEFAULT_OBJECTIVE_NAMES)
        else f"Objective {obj}"
    )
    print(
        f"=== Local dashboard (multi-obj plots → {label!r}, "
        f"idx={obj}; set OPTUNA_DASHBOARD_PLOT_OBJECTIVE to change) ===",
        file=sys.stderr,
    )
    print(f"Listening on http://{args.host}:{args.port}/", file=sys.stderr)

    httpd = make_server(args.host, args.port, app, server_class=ThreadedWSGIServer)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped.", file=sys.stderr)
        return 0
    return 0


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