run_hyperparam_study.py

#!/usr/bin/env python3
"""
Hyperparameter study for latent loop-detection models in RL-RTHS.

Uses Optuna (SQLite backend) for HPO, Weights & Biases for experiment
tracking, TensorBoard for per-trial training curves, and JSON logs for
raw results.

Setup:
  1. Put your W&B API key in .env:   WANDB_API_KEY=<key>
     Or run:  wandb login
  2. Run the study (see Usage below).

Monitoring:
  - W&B dashboard:  https://wandb.ai/<entity>/<project>
  - optuna-dashboard sqlite:///hyperparam_study/study.db
  - tensorboard --logdir hyperparam_study/tensorboard
  - cat hyperparam_study/study_summary.json

Usage:
  python run_hyperparam_study.py \\
      --game pacman \\
      --data games/pacman/data/latent_transitions.pkl \\
      --n-trials 50 --wandb

Resume an interrupted study (same --study-name / --output-dir):
  python run_hyperparam_study.py \\
      --game pacman \\
      --data games/pacman/data/latent_transitions.pkl \\
      --n-trials 100 --wandb \\
      --study-name my_study --output-dir hyperparam_study

Run without W&B (offline-only):
  python run_hyperparam_study.py \\
      --game pacman \\
      --data games/pacman/data/latent_transitions.pkl \\
      --n-trials 50
"""

from __future__ import annotations

import argparse
import json
import os
import pickle
import random
import sys
import time
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter

PROJECT_ROOT = Path(__file__).resolve().parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import optuna
from optuna.pruners import MedianPruner
from optuna.samplers import TPESampler

from rths.latent.encoder_data import sample_encoder_batch_indices
from rths.latent.losses import ForwardLoss, InverseLoss, MarginLoss
from rths.latent.models import Encoder, ForwardModel, InverseModel

# ---------------------------------------------------------------------------
# Wandb helpers
# ---------------------------------------------------------------------------

_wandb_available = False
try:
    import wandb
    _wandb_available = True
except ImportError:
    wandb = None  # type: ignore[assignment]


def _load_dotenv():
    """Load WANDB_API_KEY from .env if present."""
    env_file = PROJECT_ROOT / ".env"
    if not env_file.is_file():
        return
    for line in env_file.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, val = line.partition("=")
        key, val = key.strip(), val.strip().strip("\"'")
        if key and val:
            os.environ.setdefault(key, val)

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Hyperparameter study for latent loop-detection models",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--game", required=True, choices=["pacman", "amidar", "qbert"])
    p.add_argument("--data", required=True, help="Path to latent_transitions.pkl")

    grp = p.add_argument_group("Study settings")
    grp.add_argument("--n-trials", type=int, default=50, help="Total Optuna trials")
    grp.add_argument("--study-name", default=None, help="Optuna study name (default: auto)")
    grp.add_argument("--output-dir", default="hyperparam_study", help="Root output directory")
    grp.add_argument("--seed", type=int, default=42)
    grp.add_argument("--device", default="auto")
    grp.add_argument("--val-split", type=float, default=0.1, help="Fraction held out for eval")
    grp.add_argument("--epochs", type=int, default=15, help="Fixed number of epochs for every trial")
    grp.add_argument("--pruning", action="store_true", help="Enable median pruning of bad trials")
    grp.add_argument("--noop-action-idx", type=int, default=0)

    wb = p.add_argument_group("Weights & Biases")
    wb.add_argument("--wandb", action="store_true", help="Enable W&B logging")
    wb.add_argument("--wandb-project", default=None, help="W&B project name (default: rl-rths-<game>)")
    wb.add_argument("--wandb-entity", default=None, help="W&B entity / team (default: your user)")

    return p.parse_args()


# ---------------------------------------------------------------------------
# Data helpers
# ---------------------------------------------------------------------------

def load_dataset(path: Path) -> dict:
    if not path.is_file():
        sys.exit(
            f"Dataset not found: {path}\n"
            "Generate one first:\n"
            "  python generate_data.py --game <game> [--random-actions]"
        )
    with path.open("rb") as f:
        payload = pickle.load(f)
    for k in ("states", "next_states", "actions", "num_actions"):
        if k not in payload:
            sys.exit(f"Dataset missing key: {k}")
    return payload


def train_val_split(
    payload: dict, val_frac: float, seed: int
) -> tuple[dict, dict]:
    n = payload["states"].shape[0]
    rng = np.random.RandomState(seed)
    perm = rng.permutation(n)
    n_val = max(1, int(n * val_frac))
    val_idx, train_idx = perm[:n_val], perm[n_val:]

    def _slice(idx):
        return {
            "states": payload["states"][idx],
            "next_states": payload["next_states"][idx],
            "actions": payload["actions"][idx],
            "rewards": payload["rewards"][idx] if "rewards" in payload else None,
            "num_actions": payload["num_actions"],
        }

    return _slice(train_idx), _slice(val_idx)


# ---------------------------------------------------------------------------
# Training a single trial
# ---------------------------------------------------------------------------

def train_one_config(
    train_data: dict,
    *,
    game: str,
    lr: float,
    margin: float,
    w_margin: float,
    w_inverse: float,
    w_forward: float,
    latent_dim: int,
    hidden_dim_inverse: int,
    hidden_dim_forward: int,
    batch_size: int,
    epochs: int,
    updates_per_epoch: int,
    noop_action_idx: int,
    device: torch.device,
    writer: SummaryWriter | None = None,
    trial: optuna.Trial | None = None,
    wb_run=None,
) -> tuple[Encoder, ForwardModel, InverseModel, dict]:
    """Train encoder + forward + inverse from scratch.  Returns models and per-epoch metrics."""

    num_actions = int(train_data["num_actions"])
    states = train_data["states"]
    next_states = train_data["next_states"]
    actions = train_data["actions"]
    n = states.shape[0]

    encoder = Encoder(latent_dim=latent_dim).to(device)
    fwd_model = ForwardModel(latent_dim=latent_dim, num_actions=num_actions, hidden_dim=hidden_dim_forward).to(device)
    inv_model = InverseModel(latent_dim=latent_dim, num_actions=num_actions, hidden_dim=hidden_dim_inverse).to(device)

    margin_noop_idx = None if game == "pacman" else noop_action_idx
    margin_loss_fn = MarginLoss(margin=margin, noop_action_idx=margin_noop_idx)
    inverse_loss_fn = InverseLoss()
    forward_loss_fn = ForwardLoss()

    optimizer = torch.optim.Adam(
        list(encoder.parameters()) + list(fwd_model.parameters()) + list(inv_model.parameters()),
        lr=lr,
    )

    history: dict[str, list[float]] = {
        "total_loss": [], "margin_loss": [], "inverse_loss": [], "forward_loss": [],
    }
    global_step = 0

    for epoch in range(1, epochs + 1):
        encoder.train(); fwd_model.train(); inv_model.train()
        epoch_total = epoch_margin = epoch_inv = epoch_fwd = 0.0

        for _ in range(updates_per_epoch):
            idx = sample_encoder_batch_indices(
                game=game,
                n_transitions=n,
                actions=actions,
                noop_action_idx=noop_action_idx,
                batch_size=batch_size,
            )
            s = torch.from_numpy(states[idx]).unsqueeze(1).float().to(device)
            ns = torch.from_numpy(next_states[idx]).unsqueeze(1).float().to(device)
            a = torch.from_numpy(actions[idx]).long().to(device)
            s = (s / 255.0 - 0.5) / 0.5
            ns = (ns / 255.0 - 0.5) / 0.5

            z_t = encoder(s)
            z_t1 = encoder(ns)

            m_loss = margin_loss_fn(z_t, z_t1, a)
            inv_logits = inv_model(z_t, z_t1)
            i_loss = inverse_loss_fn(inv_logits, a)
            f_loss = forward_loss_fn(fwd_model(z_t, a), z_t1.detach())

            loss = w_margin * m_loss + w_inverse * i_loss + w_forward * f_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            epoch_total += loss.item()
            epoch_margin += m_loss.item()
            epoch_inv += i_loss.item()
            epoch_fwd += f_loss.item()
            global_step += 1

        n_up = updates_per_epoch
        avg = {
            "total_loss": epoch_total / n_up,
            "margin_loss": epoch_margin / n_up,
            "inverse_loss": epoch_inv / n_up,
            "forward_loss": epoch_fwd / n_up,
        }
        for k, v in avg.items():
            history[k].append(v)
            if writer:
                writer.add_scalar(f"train/{k}", v, epoch)
            if wb_run and _wandb_available:
                wandb.log({f"train/{k}": v, "epoch": epoch})

        unweighted_loss = avg["margin_loss"] + avg["inverse_loss"] + avg["forward_loss"]
        if writer:
            writer.add_scalar("train/unweighted_total_loss", unweighted_loss, epoch)
        if wb_run and _wandb_available:
            wandb.log({"train/unweighted_total_loss": unweighted_loss, "epoch": epoch})
        if trial is not None:
            trial.report(unweighted_loss, epoch)
            if trial.should_prune():
                raise optuna.TrialPruned()

    return encoder, fwd_model, inv_model, history


# ---------------------------------------------------------------------------
# Evaluation — loop-detection quality
# ---------------------------------------------------------------------------

@torch.no_grad()
def evaluate_loop_detection(
    encoder: Encoder,
    fwd_model: ForwardModel,
    inv_model: InverseModel,
    val_data: dict,
    *,
    noop_action_idx: int,
    device: torch.device,
    eval_batch: int = 512,
    threshold_sweep: np.ndarray | None = None,
) -> dict:
    """
    Compute metrics that proxy loop-detection quality on held-out data.

    Returns dict with:
      - forward_mse          : mean forward prediction error
      - inverse_accuracy     : fraction of correctly predicted actions
      - noop_mean_dist       : mean ||z_t - z_{t+1}|| for NOOP transitions
      - neg_mean_dist        : mean distance for random negative pairs
      - best_threshold       : threshold maximising F1 on NOOP-vs-random classification
      - best_f1              : F1 at best_threshold
      - best_recall          : recall at best_threshold  (1 - false negative rate)
      - best_precision       : precision at best_threshold
      - threshold_curve      : list of {threshold, precision, recall, f1}
      - composite_score      : single scalar for Optuna to maximize
    """
    encoder.eval(); fwd_model.eval(); inv_model.eval()

    states = val_data["states"]
    next_states = val_data["next_states"]
    actions_np = val_data["actions"]
    num_actions = int(val_data["num_actions"])
    n = states.shape[0]

    if threshold_sweep is None:
        threshold_sweep = np.concatenate([
            np.arange(0.01, 0.10, 0.01),
            np.arange(0.10, 0.50, 0.02),
            np.arange(0.50, 1.01, 0.05),
        ])

    all_z_t, all_z_t1, all_pred_z, all_inv_logits = [], [], [], []

    for start in range(0, n, eval_batch):
        end = min(start + eval_batch, n)
        s = torch.from_numpy(states[start:end]).unsqueeze(1).float().to(device)
        ns = torch.from_numpy(next_states[start:end]).unsqueeze(1).float().to(device)
        a = torch.from_numpy(actions_np[start:end]).long().to(device)
        s = (s / 255.0 - 0.5) / 0.5
        ns = (ns / 255.0 - 0.5) / 0.5

        z_t = encoder(s)
        z_t1 = encoder(ns)
        pred_z = fwd_model(z_t, a)
        inv_logits = inv_model(z_t, z_t1)

        all_z_t.append(z_t.cpu())
        all_z_t1.append(z_t1.cpu())
        all_pred_z.append(pred_z.cpu())
        all_inv_logits.append(inv_logits.cpu())

    z_t_all = torch.cat(all_z_t)
    z_t1_all = torch.cat(all_z_t1)
    pred_z_all = torch.cat(all_pred_z)
    inv_logits_all = torch.cat(all_inv_logits)

    # Forward MSE
    forward_mse = float(F.mse_loss(pred_z_all, z_t1_all).item())

    # Inverse accuracy
    inv_preds = inv_logits_all.argmax(dim=1).numpy()
    inverse_accuracy = float((inv_preds == actions_np[:n]).mean())

    # Pairwise distances for positive pairs (same transition)
    pos_dists = torch.sqrt(torch.sum((z_t_all - z_t1_all) ** 2, dim=1) + 1e-8).numpy()

    # NOOP subset — strongest positive signal (same state, distractors only)
    noop_mask = actions_np[:n] == noop_action_idx
    noop_dists = pos_dists[noop_mask] if noop_mask.any() else np.array([])
    noop_mean_dist = float(noop_dists.mean()) if noop_dists.size > 0 else float("nan")

    # Negative pairs: sample random pairs that are likely different states
    rng = np.random.RandomState(0)
    n_neg = min(len(noop_dists) * 5, n) if noop_dists.size > 0 else min(5000, n)
    idx_a = rng.randint(0, n, size=n_neg)
    idx_b = rng.randint(0, n, size=n_neg)
    different = idx_a != idx_b
    idx_a, idx_b = idx_a[different], idx_b[different]
    neg_dists = torch.sqrt(
        torch.sum((z_t_all[idx_a] - z_t_all[idx_b]) ** 2, dim=1) + 1e-8
    ).numpy()
    neg_mean_dist = float(neg_dists.mean())

    # Threshold sweep — precision / recall / F1 for reidentification
    # Positive = "same state" (NOOP pairs), Negative = random pairs
    threshold_curve = []
    best_f1, best_thresh = 0.0, 0.1
    best_prec, best_recall = 0.0, 0.0

    if noop_dists.size > 0:
        for thr in threshold_sweep:
            tp = int((noop_dists < thr).sum())
            fp = int((neg_dists < thr).sum())
            fn = int((noop_dists >= thr).sum())
            prec = tp / (tp + fp + 1e-8)
            rec = tp / (tp + fn + 1e-8)
            f1 = 2 * prec * rec / (prec + rec + 1e-8)
            threshold_curve.append({
                "threshold": round(float(thr), 4),
                "precision": round(prec, 4),
                "recall": round(rec, 4),
                "f1": round(f1, 4),
            })
            if f1 > best_f1:
                best_f1 = f1
                best_thresh = float(thr)
                best_prec = prec
                best_recall = rec

    # Margin satisfaction: fraction of negative pairs above the margin
    # (uses the training margin, but we don't have it here — use best_thresh as proxy)
    margin_sat = float((neg_dists > best_thresh).mean()) if neg_dists.size > 0 else 0.0

    # Composite score: F1-centric to balance false positives and false negatives.
    # Margin separation is weighted strongly because poor separation is the
    # root cause of false positives (different states collapsing together).
    fwd_quality = 1.0 / (1.0 + forward_mse * 10)
    composite = (
        0.35 * best_f1
        + 0.15 * best_prec
        + 0.10 * best_recall
        + 0.10 * inverse_accuracy
        + 0.10 * fwd_quality
        + 0.20 * margin_sat
    )

    return {
        "forward_mse": round(forward_mse, 6),
        "inverse_accuracy": round(inverse_accuracy, 4),
        "noop_mean_dist": round(noop_mean_dist, 6) if not np.isnan(noop_mean_dist) else None,
        "neg_mean_dist": round(neg_mean_dist, 6),
        "best_threshold": round(best_thresh, 4),
        "best_f1": round(best_f1, 4),
        "best_recall": round(best_recall, 4),
        "best_precision": round(best_prec, 4),
        "margin_satisfaction": round(margin_sat, 4),
        "fwd_quality": round(fwd_quality, 4),
        "composite_score": round(composite, 6),
        "threshold_curve": threshold_curve,
        "n_noop_pairs": int(noop_dists.size),
        "n_neg_pairs": int(neg_dists.size),
    }


# ---------------------------------------------------------------------------
# Optuna objective
# ---------------------------------------------------------------------------

def make_objective(
    train_data: dict,
    val_data: dict,
    args: argparse.Namespace,
    device: torch.device,
    output_dir: Path,
):
    tb_dir = output_dir / "tensorboard"
    results_dir = output_dir / "trial_results"
    models_dir = output_dir / "trial_models"
    results_dir.mkdir(parents=True, exist_ok=True)

    use_wandb = args.wandb and _wandb_available
    wandb_group = args.study_name or f"loop_detect_{args.game}"

    def objective(trial: optuna.Trial) -> float:
        # --- Sample hyperparameters ---
        lr = trial.suggest_float("lr", 1e-5, 5e-3, log=True)
        margin = trial.suggest_float("margin", 0.05, 2.0)
        w_margin = trial.suggest_float("w_margin", 0.1, 10.0, log=True)
        w_inverse = trial.suggest_float("w_inverse", 0.1, 10.0, log=True)
        w_forward = trial.suggest_float("w_forward", 0.1, 10.0, log=True)
        latent_dim = trial.suggest_categorical("latent_dim", [8, 16, 32])
        hidden_dim_inverse = trial.suggest_categorical("hidden_dim_inverse", [16, 32, 64, 128])
        hidden_dim_forward = trial.suggest_categorical("hidden_dim_forward", [128, 256, 512])
        batch_size = trial.suggest_categorical("batch_size", [64, 128, 256])
        updates_per_epoch = trial.suggest_categorical("updates_per_epoch", [100, 250, 500])
        epochs = args.epochs

        trial_tag = f"trial_{trial.number}"
        hparams = {k: v for k, v in trial.params.items()}
        hparams["epochs"] = epochs

        writer = SummaryWriter(log_dir=str(tb_dir / trial_tag))
        writer.add_text("hparams", json.dumps(hparams, indent=2), 0)

        # --- W&B run for this trial ---
        wb_run = None
        if use_wandb:
            wb_run = wandb.init(
                project=args.wandb_project,
                entity=args.wandb_entity,
                group=wandb_group,
                name=trial_tag,
                config={"game": args.game, **hparams},
                reinit=True,
                tags=[args.game, "hparam-study"],
            )

        if device.type == "cuda":
            torch.cuda.empty_cache()

        t0 = time.time()
        try:
            encoder, fwd_model, inv_model, history = train_one_config(
                train_data,
                game=args.game,
                lr=lr,
                margin=margin,
                w_margin=w_margin,
                w_inverse=w_inverse,
                w_forward=w_forward,
                latent_dim=latent_dim,
                hidden_dim_inverse=hidden_dim_inverse,
                hidden_dim_forward=hidden_dim_forward,
                batch_size=batch_size,
                epochs=epochs,
                updates_per_epoch=updates_per_epoch,
                noop_action_idx=args.noop_action_idx,
                device=device,
                writer=writer,
                trial=trial,
                wb_run=wb_run,
            )
        except optuna.TrialPruned:
            writer.close()
            if wb_run:
                wandb.finish(quiet=True)
            raise
        except torch.cuda.OutOfMemoryError:
            writer.close()
            if wb_run:
                wandb.finish(quiet=True)
            if device.type == "cuda":
                torch.cuda.empty_cache()
            raise optuna.TrialPruned(f"CUDA OOM with batch_size={batch_size}, latent_dim={latent_dim}")

        train_time = time.time() - t0

        # --- Evaluate ---
        metrics = evaluate_loop_detection(
            encoder, fwd_model, inv_model, val_data,
            noop_action_idx=args.noop_action_idx,
            device=device,
        )

        # Save model checkpoints while they're still alive
        trial_model_dir = models_dir / trial_tag
        trial_model_dir.mkdir(parents=True, exist_ok=True)
        torch.save(encoder.state_dict(), trial_model_dir / "encoder.pth")
        torch.save(fwd_model.state_dict(), trial_model_dir / "forward.pth")
        torch.save(inv_model.state_dict(), trial_model_dir / "inverse.pth")

        del encoder, fwd_model, inv_model
        if device.type == "cuda":
            torch.cuda.empty_cache()

        # --- Log eval metrics: TensorBoard ---
        for k in ("forward_mse", "inverse_accuracy", "best_f1", "best_recall",
                   "best_precision", "composite_score", "margin_satisfaction"):
            writer.add_scalar(f"eval/{k}", metrics[k], epochs)

        for pt in metrics.get("threshold_curve", []):
            writer.add_scalars(
                "eval/threshold_sweep",
                {"precision": pt["precision"], "recall": pt["recall"], "f1": pt["f1"]},
                global_step=int(pt["threshold"] * 1000),
            )
        writer.close()

        # --- Log eval metrics: W&B ---
        if wb_run:
            wb_summary = {
                f"eval/{k}": metrics[k]
                for k in ("forward_mse", "inverse_accuracy", "best_f1",
                           "best_recall", "best_precision", "composite_score",
                           "margin_satisfaction", "best_threshold",
                           "noop_mean_dist", "neg_mean_dist")
                if metrics.get(k) is not None
            }
            wb_summary["train_time_s"] = round(train_time, 1)
            wandb.log(wb_summary)

            thr_table = wandb.Table(
                columns=["threshold", "precision", "recall", "f1"],
                data=[[pt["threshold"], pt["precision"], pt["recall"], pt["f1"]]
                      for pt in metrics.get("threshold_curve", [])],
            )
            wandb.log({"eval/threshold_curve": thr_table})
            wandb.finish(quiet=True)

        # --- Optuna attrs ---
        trial.set_user_attr("best_threshold", metrics["best_threshold"])
        trial.set_user_attr("best_recall", metrics["best_recall"])
        trial.set_user_attr("best_f1", metrics["best_f1"])
        trial.set_user_attr("inverse_accuracy", metrics["inverse_accuracy"])
        trial.set_user_attr("forward_mse", metrics["forward_mse"])
        trial.set_user_attr("train_time_s", round(train_time, 1))

        # --- JSON log ---
        result = {
            "trial": trial.number,
            "params": hparams,
            "metrics": {k: v for k, v in metrics.items() if k != "threshold_curve"},
            "training_history": history,
            "train_time_s": round(train_time, 1),
        }
        with (results_dir / f"{trial_tag}.json").open("w") as f:
            json.dump(result, f, indent=2)

        score = metrics["composite_score"]
        print(
            f"  Trial {trial.number:>3d} | "
            f"score={score:.4f} recall={metrics['best_recall']:.3f} "
            f"prec={metrics['best_precision']:.3f} f1={metrics['best_f1']:.3f} "
            f"inv_acc={metrics['inverse_accuracy']:.3f} fwd_mse={metrics['forward_mse']:.5f} "
            f"thr={metrics['best_threshold']:.3f} | {train_time:.0f}s"
        )
        return score

    return objective


# ---------------------------------------------------------------------------
# Post-study summary
# ---------------------------------------------------------------------------

def write_summary(study: optuna.Study, output_dir: Path, args: argparse.Namespace):
    best = study.best_trial
    summary = {
        "study_name": study.study_name,
        "game": args.game,
        "dataset": str(Path(args.data).resolve()),
        "n_trials_completed": len(study.trials),
        "best_trial": best.number,
        "best_composite_score": best.value,
        "best_params": best.params,
        "best_detection_threshold": best.user_attrs.get("best_threshold"),
        "best_recall": best.user_attrs.get("best_recall"),
        "best_f1": best.user_attrs.get("best_f1"),
        "best_inverse_accuracy": best.user_attrs.get("inverse_accuracy"),
        "best_forward_mse": best.user_attrs.get("forward_mse"),
        "timestamp": datetime.now().isoformat(),
        "monitoring": {
            "optuna_dashboard": f"optuna-dashboard sqlite:///{output_dir / 'study.db'}",
            "tensorboard": f"tensorboard --logdir {output_dir / 'tensorboard'}",
            "trial_results": str(output_dir / "trial_results/"),
        },
    }

    # Also write best config as a standalone file for easy consumption
    best_config = {
        "game": args.game,
        **best.params,
        "detection_threshold": best.user_attrs.get("best_threshold"),
        "noop_action_idx": args.noop_action_idx,
    }

    (output_dir / "study_summary.json").write_text(json.dumps(summary, indent=2))
    (output_dir / "best_config.json").write_text(json.dumps(best_config, indent=2))

    # Copy best trial's models to best_models/
    best_model_src = output_dir / "trial_models" / f"trial_{best.number}"
    best_model_dst = output_dir / "best_models"
    best_model_dst.mkdir(parents=True, exist_ok=True)
    if best_model_src.exists():
        for f in best_model_src.iterdir():
            (best_model_dst / f.name).write_bytes(f.read_bytes())

    return summary


def print_report(summary: dict):
    print("\n" + "=" * 72)
    print("HYPERPARAMETER STUDY COMPLETE")
    print("=" * 72)
    print(f"  Game:            {summary['game']}")
    print(f"  Trials:          {summary['n_trials_completed']}")
    print(f"  Best trial:      #{summary['best_trial']}")
    print(f"  Composite score: {summary['best_composite_score']:.4f}")
    print(f"  Recall:          {summary['best_recall']:.4f}  (1 = no false negatives)")
    print(f"  F1:              {summary['best_f1']:.4f}")
    print(f"  Inv accuracy:    {summary['best_inverse_accuracy']:.4f}")
    print(f"  Fwd MSE:         {summary['best_forward_mse']:.6f}")
    print(f"  Det. threshold:  {summary['best_detection_threshold']:.4f}")
    print()
    print("Best hyperparameters:")
    for k, v in summary["best_params"].items():
        print(f"  {k:>25s} = {v}")
    print()
    print("Monitoring:")
    for label, cmd in summary["monitoring"].items():
        print(f"  {label:>20s}: {cmd}")
    print("=" * 72)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    _load_dotenv()
    args = parse_args()
    if args.wandb_project is None:
        args.wandb_project = f"rl-rths-{args.game}"

    device = (
        torch.device("cuda" if torch.cuda.is_available() else "cpu")
        if args.device == "auto"
        else torch.device(args.device)
    )
    print(f"Device: {device}")

    if args.wandb:
        if not _wandb_available:
            sys.exit("--wandb requested but `wandb` is not installed.  Run: uv pip install wandb")
        if not os.environ.get("WANDB_API_KEY"):
            sys.exit(
                "WANDB_API_KEY not set.\n"
                "Either run `wandb login` or add it to .env:\n"
                "  echo 'WANDB_API_KEY=<your-key>' >> .env"
            )
        print(f"W&B: project={args.wandb_project}  entity={args.wandb_entity or '(default)'}")
    else:
        print("W&B: disabled (pass --wandb to enable)")

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # --- Load data ---
    dataset_path = Path(args.data).resolve()
    payload = load_dataset(dataset_path)
    n_total = payload["states"].shape[0]
    n_noop = int((payload["actions"] == args.noop_action_idx).sum())
    print(f"Dataset: {dataset_path}  ({n_total} transitions, {n_noop} NOOP)")

    train_data, val_data = train_val_split(payload, args.val_split, args.seed)
    n_train = train_data["states"].shape[0]
    n_val = val_data["states"].shape[0]
    print(f"Split: {n_train} train / {n_val} val")

    # --- Setup output ---
    output_dir = Path(args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    storage_uri = f"sqlite:///{output_dir / 'study.db'}"

    study_name = args.study_name or f"loop_detect_{args.game}_{datetime.now():%Y%m%d_%H%M%S}"

    # --- Create / load Optuna study ---
    sampler = TPESampler(seed=args.seed)
    pruner = MedianPruner(n_startup_trials=5, n_warmup_steps=3) if args.pruning else optuna.pruners.NopPruner()

    study = optuna.create_study(
        study_name=study_name,
        storage=storage_uri,
        direction="maximize",
        sampler=sampler,
        pruner=pruner,
        load_if_exists=True,
    )

    existing = len(study.trials)
    remaining = max(0, args.n_trials - existing)
    if existing:
        print(f"Resuming study '{study_name}' — {existing} trials already done, running {remaining} more")
    else:
        print(f"Starting study '{study_name}' — {args.n_trials} trials")

    print(f"\nMonitor live:")
    if args.wandb and _wandb_available:
        entity_str = f"{args.wandb_entity}/" if args.wandb_entity else ""
        print(f"  wandb:            https://wandb.ai/{entity_str}{args.wandb_project}")
    print(f"  optuna-dashboard {storage_uri}")
    print(f"  tensorboard --logdir {output_dir / 'tensorboard'}")
    print()

    # --- Run ---
    objective_fn = make_objective(train_data, val_data, args, device, output_dir)
    study.optimize(
        objective_fn,
        n_trials=remaining,
        show_progress_bar=True,
        catch=(torch.cuda.OutOfMemoryError, RuntimeError),
    )

    # --- Summary ---
    summary = write_summary(study, output_dir, args)
    print_report(summary)


if __name__ == "__main__":
    main()