#!/usr/bin/env bash
# Render a clean Optuna contour PNG for the current exam's single-obj study.
# Studies are "{study_key}:{shard}" where study_key fingerprints exam.toml+exam.txt
# and shard is which scalar (0=stylo+256 mean, 1=stylo+32 mean, 2=32 emb).
# Usage:
#   ./plot-contour.sh              → exam hash study :0 → contour-obj0.png
#   ./plot-contour.sh 1            → exam hash study :1 → contour-obj1.png
#   ./plot-contour.sh 0 out.png    → custom path
set -euo pipefail

ROOT="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}" 2>/dev/null || realpath "${BASH_SOURCE[0]}" 2>/dev/null || echo "${BASH_SOURCE[0]}")")" && pwd)"
cd "$ROOT"

SHARD="${1:-0}"
OUT="${2:-contour-obj${SHARD}.png}"
DB="$ROOT/presets.db"

if [[ ! -f "$DB" ]]; then
  echo "MISSING: presets.db" >&2
  exit 1
fi

VENV="$ROOT/.venv"
if [[ ! -d "$VENV" ]]; then
  echo "MISSING: .venv — run ./local-dashboard.sh or ./run.sh once first." >&2
  exit 1
fi
# shellcheck disable=SC1091
source "$VENV/bin/activate"

python -m pip install -q "matplotlib!=3.6.0" "scipy" >/dev/null

python - "$DB" "$SHARD" "$OUT" "$ROOT" <<'PY'
"""Cleaner contour than stock optuna.visualization.matplotlib.plot_contour.

Stock plot is ugly here mainly because:
  1. Every completed trial is drawn as a black marker (~900 dots → solid mud).
  2. Dense black contour lines sit on top of a dark fill.
  3. Passing target= flips reverse_scale even for maximize studies.
  4. Default 4×4 grid is tiny; labels/colorbar collide.
"""
from __future__ import annotations

import hashlib
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import optuna
from matplotlib import colors
from optuna.trial import TrialState
from scipy.interpolate import griddata

# Keep names in sync with novelai_preset_optimize.OBJECTIVE_NAMES.
OBJ_NAMES = (
    "mean(stylo, 256 emb)",
    "mean(stylo, 32 emb)",
    "32 emb",
)
GRID = 80
MAX_SCATTER = 0  # 0 = all complete trials; >0 = only top-N by objective
# Sequential colormap (low → high).
CMAP = "viridis"


def _current_study_key(root: Path) -> str:
    """Same fingerprint as novelai_preset_optimize.make_study_key / build_exam."""
    toml_path = root / "exam.toml"
    txt_path = root / "exam.txt"
    if not toml_path.is_file() or not txt_path.is_file():
        raise SystemExit(f"missing exam.toml or exam.txt under {root}")
    toml_hash = hashlib.sha256(toml_path.read_bytes()).hexdigest()
    exam_hash = hashlib.sha256(txt_path.read_bytes()).hexdigest()
    return hashlib.sha256(f"{toml_hash}:{exam_hash}".encode("utf-8")).hexdigest()


def main() -> None:
    db, shard_s, out, root_s = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
    shard = int(shard_s)
    root = Path(root_s)
    if shard < 0 or shard >= len(OBJ_NAMES):
        raise SystemExit(f"shard index {shard} out of range 0..{len(OBJ_NAMES) - 1}")

    storage = f"sqlite:///{Path(db).resolve()}"
    study_key = _current_study_key(root)
    study_name = f"{study_key}:{shard}"
    names = {s.study_name for s in optuna.get_all_study_summaries(storage)}
    if study_name not in names:
        raise SystemExit(
            f"study not in presets.db: {study_name[:16]}…:{shard}\n"
            f"  study_key={study_key[:16]}… (from current exam.toml+exam.txt)\n"
            f"  known: {sorted(n[-4:] if ':' in n else n[:12] for n in names) or '(none)'}\n"
            f"Run ./run.sh once so the exam-hash studies exist."
        )
    study = optuna.load_study(study_name=study_name, storage=storage)
    n_obj = len(study.directions)
    # Single-obj studies store one value; multi-obj legacy uses values[shard].
    value_index = 0 if n_obj == 1 else shard
    if value_index < 0 or value_index >= n_obj:
        raise SystemExit(
            f"study {study_name[:16]}…:{shard} has {n_obj} direction(s); "
            f"cannot read value index {value_index}"
        )

    target_name = OBJ_NAMES[shard]
    print(
        f"study={study_name[:12]}…:{shard} ({target_name})  "
        f"directions={n_obj}  value_index={value_index}"
    )

    trials = [
        t
        for t in study.trials
        if t.state == TrialState.COMPLETE and t.values is not None and t.params
    ]
    if len(trials) < 5:
        raise SystemExit(f"need ≥5 complete trials, got {len(trials)}")

    # Stable param order (numeric-ish first).
    param_names = sorted({k for t in trials for k in t.params})
    # Prefer a fixed readable order when this project’s space is present.
    preferred = ["temperature", "top_p", "top_k", "min_p"]
    if set(param_names) == set(preferred):
        param_names = preferred

    # Matrix of raw param values + objective (always real units, not log10).
    raw = {p: np.array([float(t.params[p]) for t in trials], dtype=float) for p in param_names}
    y = np.array([float(t.values[value_index]) for t in trials], dtype=float)

    # Log scale in display space (matplotlib log axes), but tick labels stay real:
    # temperature 0.8…25, top_k 1…max — not log10(·).
    log_axes: set[str] = set()
    axis_lim: dict[str, tuple[float, float]] = {}
    for p in param_names:
        dist = trials[0].distributions.get(p)
        if dist is not None and getattr(dist, "log", False):
            log_axes.add(p)
            low = float(dist.low)
            high = float(dist.high)
            # log axes cannot include 0
            low = max(low, 1e-12)
            axis_lim[p] = (low, high)
        elif dist is not None and hasattr(dist, "low") and hasattr(dist, "high"):
            axis_lim[p] = (float(dist.low), float(dist.high))
        else:
            v = raw[p]
            axis_lim[p] = (float(v.min()), float(v.max()))
        # Heuristic: huge positive span → log display even without dist flag
        v = raw[p]
        if p not in log_axes and v.min() > 0 and (v.max() / max(v.min(), 1e-12)) > 50:
            log_axes.add(p)
            lo, hi = axis_lim[p]
            axis_lim[p] = (max(lo, float(v.min()), 1e-12), max(hi, float(v.max())))

    def lin_or_log_space(p: str, n: int) -> np.ndarray:
        lo, hi = axis_lim[p]
        if p in log_axes:
            return np.logspace(np.log10(lo), np.log10(hi), n)
        return np.linspace(lo, hi, n)

    def hist_bin_edges(p: str, n_bins: int = 28) -> np.ndarray:
        """Bin edges that keep exact search-space endpoints inside a bar.

        np.logspace(..., hi) often lands a hair *below* hi (float), so every
        trial at top_k=151523 / temperature=25 falls outside the last bin and
        the end bar vanishes. Force edges[0]=lo and edges[-1] > hi.
        """
        lo, hi = axis_lim[p]
        if p in log_axes and lo > 0:
            edges = np.logspace(np.log10(lo), np.log10(hi), n_bins + 1)
            edges[0] = lo
            edges[-1] = float(np.nextafter(hi, np.inf)) if hi > 0 else hi + 1.0
        else:
            edges = np.linspace(lo, hi, n_bins + 1)
            edges[0] = lo
            edges[-1] = float(np.nextafter(hi, np.inf)) if hi != 0 else 1e-15
        # Guarantee strictly increasing (nextafter can collapse on ints).
        for i in range(1, len(edges)):
            if edges[i] <= edges[i - 1]:
                edges[i] = np.nextafter(edges[i - 1], np.inf)
        return edges

    # Nice tick sets for this project's search space (real units).
    NICE_TICKS: dict[str, list[float]] = {
        "temperature": [0.8, 1, 2, 5, 10, 25],
        "top_k": [1, 10, 100, 1_000, 10_000, 100_000, 151_523],
        "top_p": [0.05, 0.2, 0.4, 0.6, 0.8, 1.0],
        "min_p": [0.0, 0.2, 0.4, 0.6, 0.8, 0.95],
    }

    def _fmt_tick(v: float) -> str:
        if v == 0:
            return "0"
        if abs(v) >= 1000:
            exp = int(round(np.log10(v)))
            if abs(v - 10**exp) / 10**exp < 1e-6:
                return f"$10^{{{exp}}}$"
            return f"{v:g}"
        if float(v).is_integer():
            return str(int(v))
        return f"{v:g}"

    def apply_axis_scale(ax, which: str, p: str) -> None:
        lo, hi = axis_lim[p]
        ticks = NICE_TICKS.get(p)
        if ticks is not None:
            ticks = [t for t in ticks if lo - 1e-12 <= t <= hi + 1e-12]
            if not ticks:
                ticks = None
        if which == "x":
            if p in log_axes:
                ax.set_xscale("log")
            ax.set_xlim(lo, hi)
            if ticks:
                ax.set_xticks(ticks)
                ax.set_xticklabels([_fmt_tick(t) for t in ticks])
        else:
            if p in log_axes:
                ax.set_yscale("log")
            ax.set_ylim(lo, hi)
            if ticks:
                ax.set_yticks(ticks)
                ax.set_yticklabels([_fmt_tick(t) for t in ticks])

    n = len(param_names)
    # Extra width reserved for a dedicated colorbar column (does not overlay panels).
    fig_w = 2.8 * n + 1.6
    fig_h = 2.8 * n + 0.4
    fig = plt.figure(figsize=(fig_w, fig_h), dpi=140)
    # n×n plot grid + thin colorbar column on the right
    gs = fig.add_gridspec(
        n,
        n + 1,
        width_ratios=[1.0] * n + [0.055],
        left=0.08,
        right=0.96,
        top=0.92,
        bottom=0.08,
        # Extra gap so per-panel tick numbers don't collide.
        wspace=0.32,
        hspace=0.32,
    )
    axes = np.empty((n, n), dtype=object)
    for yi in range(n):
        for xi in range(n):
            axes[yi, xi] = fig.add_subplot(gs[yi, xi])
    cax = fig.add_subplot(gs[:, n])

    vmin, vmax = float(np.nanpercentile(y, 2)), float(np.nanpercentile(y, 98))
    if vmin >= vmax:
        vmin, vmax = float(np.nanmin(y)), float(np.nanmax(y))
    norm = colors.Normalize(vmin=vmin, vmax=vmax)
    cmap = CMAP

    # Scatter: all trials (MAX_SCATTER=0) or top-N by objective.
    if MAX_SCATTER and MAX_SCATTER > 0:
        scatter_idx = np.argsort(y)[-min(MAX_SCATTER, len(y)) :]
    else:
        scatter_idx = np.arange(len(y))

    mesh = None
    for yi, py in enumerate(param_names):
        for xi, px in enumerate(param_names):
            ax = axes[yi, xi]
            xv = raw[px]
            if xi == yi:
                # Diagonal: param marginal (real units; log bins when needed).
                # Clip to [lo, hi] then use endpoint-safe edges so bound piles
                # (k=1, k=151523, T=0.8, T=25, …) actually form end bars.
                lo, hi = axis_lim[px]
                xv_h = np.clip(xv, lo, hi)
                edges = hist_bin_edges(px, 28)
                ax.hist(
                    xv_h,
                    bins=edges,
                    color="#9ecae1",
                    edgecolor="white",
                    linewidth=0.4,
                    alpha=0.9,
                )
                apply_axis_scale(ax, "x", px)
                ax.set_yticks([])
            else:
                yv = raw[py]
                # Build grid in real units (log-spaced when the param is log).
                x_lin = lin_or_log_space(px, GRID)
                y_lin = lin_or_log_space(py, GRID)
                gx, gy = np.meshgrid(x_lin, y_lin)
                # Interpolate in log space for log params so the surface is even.
                def to_interp(p: str, vals: np.ndarray) -> np.ndarray:
                    return np.log10(vals) if p in log_axes else vals

                pts = np.column_stack([to_interp(px, xv), to_interp(py, yv)])
                grid_pts = np.column_stack(
                    [to_interp(px, gx.ravel()), to_interp(py, gy.ravel())]
                )
                gz = griddata(pts, y, grid_pts, method="linear").reshape(gx.shape)
                if np.isnan(gz).any():
                    gn = griddata(pts, y, grid_pts, method="nearest").reshape(gx.shape)
                    gz = np.where(np.isnan(gz), gn, gz)

                mesh = ax.contourf(gx, gy, gz, levels=18, cmap=cmap, norm=norm, extend="both")
                ax.contour(
                    gx,
                    gy,
                    gz,
                    levels=9,
                    colors="k",
                    linewidths=0.35,
                    alpha=0.22,
                )
                ax.scatter(
                    xv[scatter_idx],
                    yv[scatter_idx],
                    c=y[scatter_idx],
                    cmap=cmap,
                    norm=norm,
                    s=8 if len(scatter_idx) > 200 else 12,
                    edgecolors="none",
                    linewidths=0,
                    alpha=0.45 if len(scatter_idx) > 200 else 0.9,
                    zorder=3,
                )
                apply_axis_scale(ax, "x", px)
                apply_axis_scale(ax, "y", py)

            # Axis name labels only on outer edges; numeric ticks on EVERY panel.
            # Real units (0.8…25, 1…top_k_max) — not log10(·).
            if yi == n - 1:
                ax.set_xlabel(px, fontsize=9)
            if xi == 0:
                ax.set_ylabel(py, fontsize=9)

            ax.tick_params(
                axis="both",
                which="both",
                labelbottom=True,
                labelleft=True,
                labelsize=6,
                length=3,
            )
            # Diagonal is a count histogram — keep y ticks off (counts not comparable
            # to param scales); still show x numbers.
            if xi == yi:
                ax.tick_params(labelleft=False, left=False)
            ax.grid(False)

    fig.suptitle(
        f"Contour — {study_name[:12]}…:{shard}: {target_name}",
        fontsize=13,
        y=0.98,
    )
    if mesh is not None:
        cbar = fig.colorbar(mesh, cax=cax)
        cbar.set_label(target_name, fontsize=10)
        cbar.ax.tick_params(labelsize=8)

    fig.savefig(out, dpi=160, bbox_inches="tight", facecolor="white", pad_inches=0.15)
    plt.close(fig)
    sc_label = "all" if not MAX_SCATTER else f"top {len(scatter_idx)}"
    print(
        f"Wrote {out}  (study :{shard}: {target_name}, "
        f"trials={len(trials)}, scatter={sc_label})"
    )


if __name__ == "__main__":
    main()
PY

echo "Open: $OUT"
