menu.py

from code.scripts.metrics_utils import *
import subprocess
import webbrowser
import os
import sys
import platform
from pathlib import Path
import shutil
import time
from omegaconf import OmegaConf

try:
    import winsound
    HAS_WINSOUND = True
except ImportError:
    HAS_WINSOUND = False

DEFAULT_TB_ROOT = "mylogs"
MODELS_DIR = Path("models/")
CONF_ROOT = Path("code/conf")
GRID_CONFIG_PATH = CONF_ROOT / "grid.yaml"
CONF_GAME_DIR = CONF_ROOT / "game"
CONF_REWARD_DIR = CONF_ROOT / "reward"
CONF_ALGO_DIR = CONF_ROOT / "algo"
CONF_ROBUSTNESS_DIR = CONF_ROOT / "robustness"

CURRENT_ALGO = None
DEFAULT_TRAIN_STEPS = 10_000_000


# ─────────────────────────────────────────────────────────────
# CONFIG HELPERS
# ─────────────────────────────────────────────────────────────

def load_grid_config():
    if not GRID_CONFIG_PATH.exists():
        return None
    try:
        return OmegaConf.load(GRID_CONFIG_PATH)
    except Exception as e:
        print(f"Error loading grid.yaml: {e}")
        return None


def get_available_games():
    cfg = load_grid_config()
    if cfg is not None and "games" in cfg:
        return sorted([g for g in cfg.games if (CONF_GAME_DIR / f"{g}.yaml").exists()])
    if not CONF_GAME_DIR.exists():
        return []
    return sorted([f.stem for f in CONF_GAME_DIR.glob("*.yaml")])


def get_available_algos_from_grid():
    cfg = load_grid_config()
    if cfg is None or "models" not in cfg:
        return []
    return sorted([a for a in cfg.models if (CONF_ALGO_DIR / f"{a}.yaml").exists()])


def get_available_personas_from_grid():
    cfg = load_grid_config()
    if cfg is None or "personas" not in cfg:
        return []
    return sorted([p for p in cfg.personas if (CONF_REWARD_DIR / f"{p}.yaml").exists()])


def get_personas_for_game(game: str):
    all_personas = get_available_personas_from_grid()
    filtered = [p for p in all_personas if p.startswith(f"{game}_")]
    return filtered if filtered else all_personas


def get_model_folders():
    BEST_DIR = MODELS_DIR / "best"
    if not BEST_DIR.exists():
        return []
    return [f for f in BEST_DIR.iterdir() if f.is_dir() and (f / "best_model.zip").exists()]


def get_trained_models_count():
    return len(get_model_folders())


def get_trained_games_from_models_flat():
    folders = get_model_folders()
    games = set()
    for folder in folders:
        parts = folder.name.split("_")
        if parts:
            games.add(parts[0])
    return sorted(games)


def ensure_current_algo():
    global CURRENT_ALGO
    if CURRENT_ALGO is None:
        algos = get_available_algos_from_grid()
        if algos:
            CURRENT_ALGO = "ppo" if "ppo" in algos else algos[0]


# ─────────────────────────────────────────────────────────────
# UI HELPERS
# ─────────────────────────────────────────────────────────────

def ask_index(prompt, options, add_back=True, default=None):
    """Print numbered list, return selected item or None for back."""
    if not options:
        print("No options available.")
        return None

    print(prompt)
    for i, opt in enumerate(options, 1):
        flag = " (default)" if opt == default else ""
        print(f"  {i}. {opt}{flag}")

    back_idx = len(options) + 1
    if add_back:
        print(f"  {back_idx}. Back")

    prompt_text = f"Select (1-{back_idx if add_back else len(options)})"
    if default:
        prompt_text += f" or Enter for [{default}]"
    prompt_text += ": "

    raw = input(prompt_text).strip()
    if raw == "" and default:
        return default
    try:
        num = int(raw)
        if add_back and num == back_idx:
            return None
        if 1 <= num <= len(options):
            return options[num - 1]
    except ValueError:
        pass
    print("Invalid selection.")
    return None


def ask_difficulty():
    """
    Ask difficulty. Returns 'default'/'easy'/'hard'/'all'.
    'default' = same settings the model was trained on (no robustness overrides).
    'all'     = runs default + easy + hard, produces separate CSVs per difficulty
                (all picked up automatically by the bar graph generator).
    """
    print("\nSelect difficulty:")
    print("  1. default  (training conditions, no changes)")
    print("  2. easy     (forgiving settings)")
    print("  3. hard     (punishing settings)")
    print("  4. all      (runs all three - default + easy + hard)")
    raw = input("Select (1-4) [Enter = default]: ").strip()
    mapping = {"1": "default", "2": "easy", "3": "hard", "4": "all"}
    difficulty = mapping.get(raw, "default")
    print(f"  -> {difficulty}")
    return difficulty


def ask_episodes(default=100):
    raw = input(f"Episodes? [{default}]: ").strip()
    try:
        v = int(raw)
        if v > 0:
            return v
    except ValueError:
        pass
    return default


def ask_timesteps(default_steps=DEFAULT_TRAIN_STEPS):
    s = input(f"Timesteps [default {default_steps}]: ").strip()
    if not s:
        return int(default_steps)
    try:
        v = int(s)
        if v > 0:
            return v
    except ValueError:
        pass
    print("Invalid, using default.")
    return int(default_steps)


def open_browser(url):
    try:
        webbrowser.open(url)
    except Exception:
        pass
    print(f"If browser did not open: {url}")


# ─────────────────────────────────────────────────────────────
# TRAINING
# ─────────────────────────────────────────────────────────────

def execute_training_run(game, algo, persona, tb_root=DEFAULT_TB_ROOT, timesteps=None):
    cmd = [
        sys.executable, "-m", "code.scripts.train",
        f"game={game}", f"model={algo}", f"persona={persona}", f"tb_root={tb_root}",
    ]
    if timesteps is not None:
        cmd += ["skill=Custom", f"++skills.Custom={timesteps}"]

    print(">>> " + " ".join(cmd) + "\n")
    try:
        subprocess.run(cmd, check=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"Failed: {game}|{algo}|{persona} (exit {e.returncode})")
        return False
    except KeyboardInterrupt:
        raise


def print_training_summary(total, successful, failed):
    print("\n" + "=" * 60)
    print("TRAINING COMPLETE")
    print("=" * 60)
    print(f"Successful: {successful}/{total}")
    if failed > 0:
        print(f"Failed: {failed}/{total}")
    if HAS_WINSOUND and failed == 0:
        try:
            winsound.PlaySound("chime.wav", winsound.SND_FILENAME)
        except Exception:
            pass


def run_training():
    global CURRENT_ALGO
    print("\n=== Training ===")

    games = get_available_games()
    if not games:
        print("No game configs found.")
        return

    game = ask_index("Available games:", games)
    if game is None:
        return

    algos = get_available_algos_from_grid()
    if not algos:
        print("No algorithm configs found.")
        return

    ensure_current_algo()
    default_algo = CURRENT_ALGO if CURRENT_ALGO in algos else algos[0]
    algo_choice = ask_index("Available algorithms:", algos, default=default_algo)
    if algo_choice is None:
        return
    CURRENT_ALGO = algo_choice

    personas = get_personas_for_game(game)
    if not personas:
        print(f"No personas found for game='{game}'.")
        return

    persona_choice = ask_index("Available personas:", personas)
    if persona_choice is None:
        return

    tb_root = input(f"TensorBoard log root [{DEFAULT_TB_ROOT}]: ").strip() or DEFAULT_TB_ROOT
    timesteps = ask_timesteps(DEFAULT_TRAIN_STEPS)

    print(f"\n>>> Training {game} | {algo_choice} | {persona_choice} | {timesteps} steps\n")
    success = execute_training_run(game, algo_choice, persona_choice, tb_root=tb_root, timesteps=timesteps)
    if HAS_WINSOUND and success:
        try:
            winsound.PlaySound("chime.wav", winsound.SND_FILENAME)
        except Exception:
            pass
    print("\nTraining completed.\n")


def train_all_models_for_game():
    global CURRENT_ALGO
    print("\n" + "=" * 60)
    print("TRAIN ALL MODELS FOR ONE GAME")
    print("=" * 60)

    games = get_available_games()
    if not games:
        print("No game configs found.")
        return

    game = ask_index("Available games:", games)
    if game is None:
        return

    algos = get_available_algos_from_grid()
    if not algos:
        print("No algorithm configs found.")
        return

    ensure_current_algo()
    print("\nSelect algorithms to train:")
    print("  1. All algorithms")
    for i, algo in enumerate(algos, 2):
        flag = " (current)" if algo == CURRENT_ALGO else ""
        print(f"  {i}. {algo}{flag}")
    print(f"  {len(algos) + 2}. Back")

    choice = input(f"Select (1-{len(algos) + 2}): ").strip()
    try:
        num = int(choice)
        if num == len(algos) + 2:
            return
        elif num == 1:
            selected_algos = algos
        elif 2 <= num <= len(algos) + 1:
            selected_algos = [algos[num - 2]]
            CURRENT_ALGO = selected_algos[0]
        else:
            print("Invalid selection.")
            return
    except ValueError:
        print("Invalid selection.")
        return

    personas = get_personas_for_game(game)
    if not personas:
        print(f"No personas for '{game}'.")
        return

    timesteps = ask_timesteps(DEFAULT_TRAIN_STEPS)
    total_runs = len(selected_algos) * len(personas)

    print(f"\nTraining {total_runs} model(s) for '{game}'.")
    confirm = input("Proceed? [y/N]: ").strip().lower()
    if confirm not in ("y", "yes"):
        print("Aborted.")
        return

    completed = failed = 0
    try:
        for algo in selected_algos:
            for persona in personas:
                completed += 1
                print(f"\n[{completed}/{total_runs}] {game} | {algo} | {persona}")
                if not execute_training_run(game=game, algo=algo, persona=persona, timesteps=timesteps):
                    failed += 1
    except KeyboardInterrupt:
        print(f"\nInterrupted. Completed {completed - 1}/{total_runs}")
        return

    print_training_summary(total_runs, completed - failed, failed)


def train_complete_grid():
    print("\n" + "=" * 60)
    print("TRAIN COMPLETE GRID (All Games x Algos x Personas)")
    print("=" * 60)

    games = get_available_games()
    algos = get_available_algos_from_grid()

    if not games or not algos:
        print("No game or algorithm configs found.")
        return

    timesteps = ask_timesteps(DEFAULT_TRAIN_STEPS)
    total_runs = 0
    breakdown = []

    for game in games:
        personas = get_personas_for_game(game)
        if not personas:
            continue
        runs = len(algos) * len(personas)
        total_runs += runs
        breakdown.append(f"  {game}: {len(personas)} persona(s) x {len(algos)} algo(s) = {runs} runs")

    if total_runs == 0:
        print("No valid training configurations found.")
        return

    print(f"\nThis will train {total_runs} total model(s) at {timesteps} steps each.")
    for line in breakdown:
        print(line)

    confirm = input("\nProceed? [y/N]: ").strip().lower()
    if confirm not in ("y", "yes"):
        print("Aborted.")
        return

    completed = failed = 0
    try:
        for game in games:
            personas = get_personas_for_game(game)
            if not personas:
                continue
            for algo in algos:
                for persona in personas:
                    completed += 1
                    print(f"\n[{completed}/{total_runs}] {game} | {algo} | {persona}")
                    if not execute_training_run(game=game, algo=algo, persona=persona, timesteps=timesteps):
                        failed += 1
    except KeyboardInterrupt:
        print_training_summary(total_runs, completed - failed, failed)
        return

    print_training_summary(total_runs, completed - failed, failed)


# ─────────────────────────────────────────────────────────────
# EVALUATION
# ─────────────────────────────────────────────────────────────

def _run_eval_for_models(model_dirs, difficulty, episodes):
    """
    Core eval runner. Loops over model dirs, calls evaluate.py for each.
    When difficulty='all', evaluate.py internally runs all three and writes
    separate CSVs (e.g. name_eval_easy.csv / name_eval_hard.csv / name_eval_default.csv).
    All of these match the '*_eval.csv' glob so bar graphs pick them up automatically.
    """
    for model_dir in model_dirs:
        model_zip = model_dir / "best_model.zip"
        model_name = model_dir.name
        parts = model_name.split("_")
        game = parts[0]
        algo = parts[1] if len(parts) >= 2 else "ppo"

        out_csv = MODELS_DIR / f"{model_name}_eval.csv"
        # dk_balance uses lowercase 'dk' prefix; others use Title case
        _metrics_cls_name = {
            "dk": "dkBalanceStats",
        }.get(game, f"{game.capitalize()}BalanceStats")
        metrics_class = f"code.metrics.{game}_balance.{_metrics_cls_name}"

        cmd = [
            sys.executable, "-m", "code.scripts.evaluate",
            "--game", game,
            "--algo", algo,
            "--model", str(model_zip),
            "--episodes", str(episodes),
            "--difficulty", difficulty,
            "--render", "none",
            "--out", str(out_csv),
            "--metrics", metrics_class,
        ]
        print(">>>", " ".join(cmd))
        subprocess.run(cmd)



# Default kill-switch windows (minutes) per game
# Default max minutes per game for unlimited eval
_UNLIMITED_DEFAULTS = {"flappy": 30, "dk": 30, "pong": 30, "snake": 30}

# Reference FPS per game used to convert minutes -> steps
# DK display FPS is 30 but physics REFERENCE_FPS is 60; snake uses max_fps 26
_GAME_REF_FPS = {"flappy": 60, "dk": 60, "pong": 60, "snake": 26}


def _minutes_to_steps(game: str, minutes: float) -> int:
    fps = _GAME_REF_FPS.get(game, 60)
    return int(minutes * 60 * fps)


def ask_unlimited_times(games: list) -> dict:
    """
    Ask the user how long (in minutes) each game should run before the
    kill switch fires.  If only one game, ask once.  If multiple, ask per game.
    Returns dict: {game: steps}
    """
    print("\n=== Unlimited eval  kill-switch step limits ===")
    times = {}
    if len(games) == 1:
        game = games[0]
        default = _UNLIMITED_DEFAULTS.get(game, 30)
        raw = input(f"  Max time for {game} in minutes [{default}]: ").strip()
        try:
            minutes = float(raw) if raw else default
        except ValueError:
            minutes = default
        times[game] = _minutes_to_steps(game, minutes)
    else:
        print("  Enter max time in minutes for each game (Enter = default):")
        for game in games:
            default = _UNLIMITED_DEFAULTS.get(game, 30)
            raw = input(f"    {game} [{default} min]: ").strip()
            try:
                minutes = float(raw) if raw else default
            except ValueError:
                minutes = default
            times[game] = _minutes_to_steps(game, minutes)
    return times


def _run_unlimited_eval_for_models(model_dirs, difficulty, episodes, time_limits: dict = None):
    """
    Like _run_eval_for_models but calls evaluate_unlimited.py.
    CSV is saved as <model_name>_unlimited_eval.csv.
    time_limits: {game: seconds}  per-game kill-switch window.
    """
    for model_dir in model_dirs:
        model_zip  = model_dir / "best_model.zip"
        model_name = model_dir.name
        parts = model_name.split("_")
        game  = parts[0]
        algo  = parts[1] if len(parts) >= 2 else "ppo"

        out_csv = MODELS_DIR / f"{model_name}_unlimited_eval.csv"
        _metrics_cls_name = {
            "dk": "dkBalanceStats",
        }.get(game, f"{game.capitalize()}BalanceStats")
        metrics_class = f"code.metrics.{game}_balance.{_metrics_cls_name}"

        # Look up this game's step limit; fall back to default if not specified
        game_steps = (time_limits or {}).get(game, _minutes_to_steps(game, _UNLIMITED_DEFAULTS.get(game, 30)))

        cmd = [
            sys.executable, "-m", "code.scripts.evaluate_unlimited",
            "--game",       game,
            "--algo",       algo,
            "--model",      str(model_zip),
            "--episodes",   str(episodes),
            "--difficulty", difficulty,
            "--render",     "none",
            "--out",        str(out_csv),
            "--metrics",    metrics_class,
            "--step_limit", str(game_steps),
        ]
        print(">>>", " ".join(cmd))
        subprocess.run(cmd)


def run_evaluation():
    """
    Evaluate trained models.
    Sub-menu: Standard Evaluation or Unlimited Evaluation.
    """
    print("\n=== Evaluation ===")
    print("  1. Standard Evaluation   (normal time limits)")
    print("  2. Unlimited Evaluation  (generous per-game kill switches)")
    mode_raw = input("Select (1-2) [Enter = standard]: ").strip()
    unlimited = (mode_raw == "2")

    BEST_DIR = MODELS_DIR / "best"
    if not BEST_DIR.exists():
        print("[!] models/best/ does not exist  train some models first.")
        return

    model_folders = get_model_folders()
    if not model_folders:
        print("No best_model.zip files found in models/best/.")
        return

    games = sorted(set(f.name.split("_")[0] for f in model_folders if "_" in f.name))
    if not games:
        print("No recognizable games found.")
        return

    # Game selection  includes "All games" option
    game_options = games + ["All games"]
    game_choice = ask_index("Select game to evaluate:", game_options)
    if game_choice is None:
        return

    if game_choice == "All games":
        selected_model_dirs = model_folders
        print(f"\nEvaluating all {len(selected_model_dirs)} model(s) across all games.")
    else:
        selected_model_dirs = [f for f in model_folders if f.name.startswith(f"{game_choice}_")]
        if not selected_model_dirs:
            print(f"No models found for '{game_choice}'.")
            return
        print(f"\nFound {len(selected_model_dirs)} model(s) for '{game_choice}'.")

    difficulty = ask_difficulty()
    default_eps = 10 if unlimited else 100
    episodes = ask_episodes(default=default_eps)

    time_limits = None
    if unlimited:
        # Determine which games are in the selection
        selected_games = sorted(set(f.name.split("_")[0] for f in selected_model_dirs if "_" in f.name))
        time_limits = ask_unlimited_times(selected_games)

    mode_label = "UNLIMITED" if unlimited else "STANDARD"
    print(f"\nRunning [{mode_label}]: {len(selected_model_dirs)} model(s) | difficulty={difficulty} | episodes={episodes}\n")

    if unlimited:
        _run_unlimited_eval_for_models(selected_model_dirs, difficulty, episodes, time_limits=time_limits)
    else:
        _run_eval_for_models(selected_model_dirs, difficulty, episodes)

    print("\nEvaluation complete.\n")


# ─────────────────────────────────────────────────────────────
# WATCH AGENT
# ─────────────────────────────────────────────────────────────

def watch_trained_agent():
    """Watch a trained agent play  pick game -> model -> difficulty."""
    print("\n=== Watch Trained Agent Play ===")

    model_folders = get_model_folders()
    if not model_folders:
        print("No best_model.zip files found in models/best/.")
        return

    # Step 1: pick game
    games = sorted(set(f.name.split("_")[0] for f in model_folders if "_" in f.name))
    game = ask_index("Pick a game:", games)
    if game is None:
        return

    # Step 2: filter and list models for that game
    game_models = [f for f in model_folders if f.name.startswith(f"{game}_")]
    if not game_models:
        print(f"No trained models found for '{game}'.")
        return

    labels = []
    for folder in game_models:
        parts = folder.name.split("_")
        if len(parts) >= 3:
            algo = parts[1]
            skill = parts[-1]
            persona = "_".join(parts[2:-1]) if len(parts) > 3 else parts[2]
            labels.append(f"{algo:<6} | {persona:<20} | {skill}")
        else:
            labels.append(folder.name)

    selected_label = ask_index(f"Select model for '{game}':", labels)
    if selected_label is None:
        return

    model_idx = labels.index(selected_label)
    selected_folder = game_models[model_idx]
    folder_parts = selected_folder.name.split("_")
    algo = folder_parts[1] if len(folder_parts) >= 2 else "ppo"
    skill = folder_parts[-1] if len(folder_parts) >= 4 else None
    persona = "_".join(folder_parts[2:-1]) if len(folder_parts) > 3 else None

    # Step 3: difficulty
    difficulty = ask_difficulty()

    episodes = ask_episodes(default=5)
    fps_input = input("FPS? [30]: ").strip()
    fps = int(fps_input) if fps_input.isdigit() and int(fps_input) > 0 else 30

    cmd = [
        sys.executable, "-m", "code.scripts.play",
        "--game", game,
        "--model", algo,
        "--render",
        "--fps", str(fps),
        "--episodes", str(episodes),
        "--difficulty", difficulty,
        "--deterministic",
    ]
    if persona:
        cmd += ["--persona", persona]
    if skill:
        cmd += ["--skill", skill]

    print("\n>>>", " ".join(cmd), "\n")
    subprocess.run(cmd)
    print("\nVisualization complete.\n")


# ─────────────────────────────────────────────────────────────
# MANUAL PLAY
# ─────────────────────────────────────────────────────────────

def run_manual_play():
    """Human keyboard play  pick game then difficulty."""
    print("\n=== Manual Play ===")

    available_games = get_available_games()
    if not available_games:
        print("No game configs found.")
        return

    game = ask_index("Select game to play:", available_games)
    if game is None:
        return

    # Always ask difficulty so player can try different settings
    difficulty = ask_difficulty()

    controls = {
        "flappy": "SPACE = flap, ESC = quit",
        "snake":  "W/Up, S/Down, A/Left, D/Right, ESC = quit",
        "pong":   "W/Up = up, S/Down = down, ESC = quit",
        "dk":     "W/A/S/D = move, SPACE = jump, ESC = quit",
    }
    print(f"\n=== Playing {game} ({difficulty}) ===")
    print(f"Controls: {controls.get(game, 'Use game keys, ESC to quit')}")

    cmd = [
        sys.executable, "-m", "code.scripts.manual_play",
        "--game", game,
        "--fps", "30",
        "--difficulty", difficulty,
        "--episodes", "999",
    ]
    print(">>>", " ".join(cmd))
    subprocess.run(cmd)


# ─────────────────────────────────────────────────────────────
# RANDOM BASELINE
# ─────────────────────────────────────────────────────────────

def run_random_baseline_eval():
    """Random baseline  pick game, difficulty, episodes."""
    print("\n=== Random Baseline Evaluation ===")

    available_games = get_available_games()
    if not available_games:
        print("No game configs found.")
        return

    game_choice = ask_index("Select game:", available_games + ["All games"])
    if game_choice is None:
        return

    selected_games = available_games if game_choice == "All games" else [game_choice]
    difficulty = ask_difficulty()
    episodes = ask_episodes(default=100)

    difficulties = ["default", "easy", "hard"] if difficulty == "all" else [difficulty]

    for game in selected_games:
        for diff in difficulties:
            print(f"\n=== Random baseline: {game} | {diff} | {episodes} episodes ===")
            cmd = [
                sys.executable, "-m", "code.scripts.random_eval",
                "--game", game,
                "--episodes", str(episodes),
                "--difficulty", diff,
            ]
            print(">>>", " ".join(cmd))
            subprocess.run(cmd)

    print("\nRandom baseline evaluation complete.\n")


# ─────────────────────────────────────────────────────────────
# TENSORBOARD
# ─────────────────────────────────────────────────────────────

def run_tensorboard():
    print("\n=== TensorBoard ===")
    root = Path(DEFAULT_TB_ROOT)
    if not root.exists():
        print(f"No '{DEFAULT_TB_ROOT}/' folder found yet. Train first.")
        return

    for port in range(6006, 6011):
        cmd = [sys.executable, "-m", "tensorboard.main", "--logdir", str(root), "--port", str(port)]
        print(f"\nLaunching TensorBoard on http://localhost:{port}/")
        print(">>> " + " ".join(cmd))
        print("Press Ctrl+C to stop.\n")
        try:
            tb_proc = subprocess.Popen(cmd)
            time.sleep(3)
            open_browser(f"http://localhost:{port}/")
            tb_proc.wait()
            return
        except KeyboardInterrupt:
            print("\nTensorBoard stopped.")
            return
        except FileNotFoundError:
            print("TensorBoard not found. Run: pip install tensorboard")
            return
        except Exception as e:
            print(f"Error on port {port}: {e}")

    print("Unable to start TensorBoard on ports 6006-6010.")


# ─────────────────────────────────────────────────────────────
# METRICS / GRAPHS
# ─────────────────────────────────────────────────────────────

def metrics_menu(models_dir: Path) -> None:
    while True:
        print("\n=== Metrics / Graphs Menu ===")
        print("--- Training ---")
        print("1. All training graphs (per game)")
        print("--- Standard Eval ---")
        print("2. All standard eval bar graphs (per game)")
        print("3. Snake only  (training + standard eval)")
        print("4. Flappy only (training + standard eval)")
        print("5. Pong only   (training + standard eval)")
        print("6. DK only     (training + standard eval)")
        print("--- Unlimited Eval ---")
        print("7. All unlimited eval bar graphs (per game)")
        print("8. Snake only  (unlimited eval)")
        print("9. Flappy only (unlimited eval)")
        print("10. Pong only  (unlimited eval)")
        print("11. DK only    (unlimited eval)")
        print("0. Back")

        choice = input("Select: ").strip()

        if choice == "1":
            generate_all_metrics_graphs(models_dir)
        elif choice == "2":
            generate_all_eval_bar_graphs(models_dir)
        elif choice == "3":
            generate_all_snake_graphs(models_dir)
            generate_eval_bar_for_game(models_dir, "snake")
        elif choice == "4":
            generate_all_flappy_graphs(models_dir)
            generate_eval_bar_for_game(models_dir, "flappy")
        elif choice == "5":
            generate_all_pong_graphs(models_dir)
            generate_eval_bar_for_game(models_dir, "pong")
        elif choice == "6":
            generate_all_dk_graphs(models_dir)
            generate_eval_bar_for_game(models_dir, "dk")
        elif choice == "7":
            generate_all_unlimited_eval_bar_graphs(models_dir)
        elif choice == "8":
            generate_unlimited_eval_bar_for_game(models_dir, "snake")
        elif choice == "9":
            generate_unlimited_eval_bar_for_game(models_dir, "flappy")
        elif choice == "10":
            generate_unlimited_eval_bar_for_game(models_dir, "pong")
        elif choice == "11":
            generate_unlimited_eval_bar_for_game(models_dir, "dk")
        elif choice == "0":
            break
        else:
            print("Invalid option.")


# ─────────────────────────────────────────────────────────────
# PROJECT STATUS
# ─────────────────────────────────────────────────────────────

def show_project_status():
    print("\n=== Project Status ===")
    games = get_available_games()
    algos = get_available_algos_from_grid()
    trained = get_trained_games_from_models_flat()

    print(f"Game configs: {len(games)}")
    for g in games:
        flag = "trained" if g in trained else "not trained"
        print(f"  {g}: {flag}")

    print(f"\nAlgorithm configs: {len(algos)}")
    for a in algos:
        flag = " (current)" if a == CURRENT_ALGO else ""
        print(f"  {a}{flag}")

    model_folders = get_model_folders()
    print(f"\nTrained models: {len(model_folders)}")
    algo_counts = {}
    for f in model_folders:
        parts = f.name.split("_")
        if len(parts) >= 2:
            algo_counts[parts[1]] = algo_counts.get(parts[1], 0) + 1
    for algo, count in sorted(algo_counts.items()):
        print(f"  {algo}: {count} model(s)")

    if CONF_ROBUSTNESS_DIR.exists():
        rob_files = list(CONF_ROBUSTNESS_DIR.glob("*.yaml"))
        print(f"\nRobustness configs: {len(rob_files)}")
        for f in sorted(rob_files):
            print(f"  {f.stem}")
    else:
        print("\nRobustness configs: not found")
    print()


# ─────────────────────────────────────────────────────────────
# DELETE LOGS
# ─────────────────────────────────────────────────────────────

def delete_logs_and_models():
    print("\n=== Delete Logs / Models ===")
    print("WARNING: This is irreversible.")
    print("  1. Delete TensorBoard logs only (mylogs/)")
    print("  2. Delete models only (models/)")
    print("  3. Delete BOTH")
    print("  4. Back")

    choice = input("Select (1-4): ").strip()
    if choice == "4" or choice == "":
        return

    targets = []
    if choice in ("1", "3"):
        targets.append(Path(DEFAULT_TB_ROOT))
    if choice in ("2", "3"):
        targets.append(MODELS_DIR)

    if not targets:
        print("Invalid selection.")
        return

    confirm = input(f"Delete {[str(t) for t in targets]}? Type YES to confirm: ").strip()
    if confirm != "YES":
        print("Aborted.")
        return

    for t in targets:
        if t.exists():
            shutil.rmtree(t)
            print(f"Deleted: {t}")
        else:
            print(f"Not found (skipped): {t}")
    print("Done.\n")


# ─────────────────────────────────────────────────────────────
# RECORD GAMEPLAY
# ─────────────────────────────────────────────────────────────

def _build_record_cmd(game, difficulty, agent_type, model_path, duration, fps, fmt):
    """Build the record_gameplay subprocess command."""
    # Derive a distinct output filename so batch clips never overwrite each other.
    # random  → <game>_<diff>_random.<ext>
    # human   → <game>_<diff>_human.<ext>
    # model   → <game>_<diff>_<algo>.<ext>  e.g. dk_default_ppo.mp4
    out_dir = Path("outputs") / "recordings"
    if agent_type == "model" and model_path:
        folder_stem = Path(model_path).parent.name  # e.g. "dk_ppo_dk_baseline_custom"
        stem_parts = folder_stem.split("_")
        algo = stem_parts[1] if len(stem_parts) > 1 else folder_stem  # e.g. "ppo"
        out_name = f"{game}_{difficulty}_{algo}.{fmt}"
    else:
        out_name = f"{game}_{difficulty}_{agent_type}.{fmt}"
    output = str(out_dir / out_name)

    cmd = [
        sys.executable, "-m", "code.scripts.record_gameplay",
        "--game",       game,
        "--difficulty", difficulty,
        "--agent",      agent_type,
        "--duration",   str(duration),
        "--fps",        str(fps),
        "--format",     fmt,
        "--output",     output,
    ]
    if model_path:
        cmd += ["--model", model_path]
    return cmd


def _ask_record_settings():
    """Ask duration, fps, format. Returns (duration, fps, fmt)."""
    dur_raw = input("Duration per clip in seconds? [30]: ").strip()
    try:
        duration = max(1, int(dur_raw) if dur_raw else 30)
    except ValueError:
        duration = 30

    fps_raw = input("FPS? [60]: ").strip()
    try:
        fps = max(1, int(fps_raw) if fps_raw else 60)
    except ValueError:
        fps = 60

    print("\nOutput format:")
    print("  1. mp4  (recommended)")
    print("  2. gif")
    fmt_raw = input("Select (1-2) [Enter = mp4]: ").strip()
    fmt = "gif" if fmt_raw == "2" else "mp4"

    return duration, fps, fmt


def run_record_gameplay():
    """Record gameplay to MP4 or GIF  single clip or full batch (all games × difficulties × agents)."""
    print("\n=== Record Gameplay ===")

    available_games = get_available_games()
    if not available_games:
        print("No game configs found.")
        return

    print("\nRecord mode:")
    print("  1. Single clip  (pick game, agent, difficulty)")
    print("  2. All          (every game × every difficulty × random + all trained models)")
    print("  3. Back")
    mode_raw = input("Select (1-3): ").strip()

    if mode_raw == "3" or mode_raw == "":
        return

    # ── BATCH MODE ──────────────────────────────────────────
    if mode_raw == "2":
        difficulties = ["default", "easy", "hard"]
        model_folders = get_model_folders()

        duration, fps, fmt = _ask_record_settings()

        # Build all jobs: (game, difficulty, agent_type, model_path)
        jobs = []
        for game in available_games:
            for diff in difficulties:
                jobs.append((game, diff, "random", None))
                game_models = [f for f in model_folders if f.name.startswith(f"{game}_")]
                for folder in game_models:
                    jobs.append((game, diff, "model", str(folder / "best_model.zip")))

        if not jobs:
            print("No jobs to run (no game configs or trained models found).")
            return

        print(f"\nBatch will record {len(jobs)} clip(s):")
        print(f"  Games      : {', '.join(available_games)}")
        print(f"  Difficulties: default, easy, hard")
        print(f"  Agents     : random + {len(model_folders)} trained model(s)")
        print(f"  Settings   : {duration}s @ {fps}fps → .{fmt}")
        print(f"  Output     : outputs/recordings/\n")
        confirm = input("Proceed? [y/N]: ").strip().lower()
        if confirm not in ("y", "yes"):
            print("Aborted.")
            return

        completed = failed = 0
        for i, (game, diff, agent_type, model_path) in enumerate(jobs, 1):
            label = f"{game} | {diff} | {agent_type}"
            if model_path:
                label += f" | {Path(model_path).parent.name}"
            print(f"\n[{i}/{len(jobs)}] {label}")
            cmd = _build_record_cmd(game, diff, agent_type, model_path, duration, fps, fmt)
            print(">>>", " ".join(cmd))
            result = subprocess.run(cmd)
            if result.returncode == 0:
                completed += 1
            else:
                failed += 1
                print(f"  [!] Failed (exit {result.returncode})")

        print(f"\n{'=' * 60}")
        print(f"Batch complete: {completed}/{len(jobs)} succeeded, {failed} failed.")
        print(f"Output → outputs/recordings/")
        print(f"{'=' * 60}\n")
        return

    # ── SINGLE MODE ─────────────────────────────────────────
    game = ask_index("Select game to record:", available_games)
    if game is None:
        return

    print("\nSelect agent type:")
    print("  1. random  (uniform random actions)")
    print("  2. model   (trained SB3 model)")
    print("  3. human   (you play via keyboard)")
    print("  4. Back")
    agent_raw = input("Select (1-4): ").strip()
    agent_map = {"1": "random", "2": "model", "3": "human"}
    if agent_raw == "4" or agent_raw not in agent_map:
        return
    agent_type = agent_map[agent_raw]

    model_path = None
    if agent_type == "model":
        model_folders = get_model_folders()
        game_models = [f for f in model_folders if f.name.startswith(f"{game}_")]
        if not game_models:
            print(f"No trained models found for '{game}'. Train first or choose a different agent.")
            return
        labels = []
        for folder in game_models:
            parts = folder.name.split("_")
            if len(parts) >= 3:
                algo = parts[1]
                skill = parts[-1]
                persona = "_".join(parts[2:-1]) if len(parts) > 3 else parts[2]
                labels.append(f"{algo:<6} | {persona:<20} | {skill}")
            else:
                labels.append(folder.name)
        selected_label = ask_index(f"Select model for '{game}':", labels)
        if selected_label is None:
            return
        model_idx = labels.index(selected_label)
        model_path = str(game_models[model_idx] / "best_model.zip")

    difficulty = ask_difficulty()
    duration, fps, fmt = _ask_record_settings()

    print(f"\n[record] Starting: {game} | {difficulty} | {agent_type} | {duration}s @ {fps}fps → .{fmt}")
    if agent_type == "human":
        controls = {
            "flappy": "SPACE/UP = flap",
            "snake":  "W/A/S/D or arrow keys",
            "pong":   "W/Up = up,  S/Down = down",
            "dk":     "Arrows = move,  Z/SPACE = jump",
        }
        print(f"Controls: {controls.get(game, 'Arrow keys / WASD')}")
    print("Output → outputs/recordings/\n")

    cmd = _build_record_cmd(game, difficulty, agent_type, model_path, duration, fps, fmt)
    print(">>>", " ".join(cmd))
    subprocess.run(cmd)
    print("\nRecording complete.\n")


# ─────────────────────────────────────────────────────────────
# MAIN MENU
# ─────────────────────────────────────────────────────────────

def main():
    import warnings
    warnings.filterwarnings("ignore", message=r".*pkg_resources is deprecated as an API.*")

    global CURRENT_ALGO

    while True:
        games = get_available_games()
        algos = get_available_algos_from_grid()
        trained_games = get_trained_games_from_models_flat()
        trained_models = get_trained_models_count()
        ensure_current_algo()

        algo_display = f" | Current algo: {CURRENT_ALGO}" if CURRENT_ALGO else ""

        print("=" * 60)
        print("MULTI-GAME RL TRAINING & EVALUATION MENU")
        print("=" * 60)
        print(f"Games: {len(games)} | Algos: {len(algos)} | Trained games: {len(trained_games)} | Models: {trained_models}{algo_display}")
        print("------------------------------------------------------------")
        print("1.  Run Training")
        print("2.  Run Evaluation     (pick game or all, pick difficulty)")
        print("3.  View TensorBoard Logs")
        print("4.  Play Game Manually (keyboard, pick difficulty)")
        print("5.  Show Project Status")
        print("6.  Run Random Baseline (pick game or all, pick difficulty)")
        print("7.  Watch Trained Agent (pick game + model + difficulty)")
        print("8.  Train All Models for One Game")
        print("9.  Train Complete Grid")
        print("10. Metrics / Graphs")
        print("11. Record Gameplay    (save MP4 or GIF)")
        print("12. Delete Logs / Models")
        print("13. Exit")
        print("=" * 60)

        choice = input("Select option (1-13): ").strip()

        if choice == "1":
            run_training()
        elif choice == "2":
            run_evaluation()
        elif choice == "3":
            run_tensorboard()
        elif choice == "4":
            run_manual_play()
        elif choice == "5":
            show_project_status()
        elif choice == "6":
            run_random_baseline_eval()
        elif choice == "7":
            watch_trained_agent()
        elif choice == "8":
            train_all_models_for_game()
        elif choice == "9":
            train_complete_grid()
        elif choice == "10":
            metrics_menu(MODELS_DIR)
        elif choice == "11":
            run_record_gameplay()
        elif choice == "12":
            delete_logs_and_models()
        elif choice == "13":
            print("Exiting. Happy training!")
            break
        else:
            print("Invalid selection, try again.")


if __name__ == "__main__":
    main()