delete_model.py

import argparse
import csv
import json
import os
import shutil

from profiles import format_profile_label, get_profile_paths, list_profile_names


LEDGE_REGION_BOUNDS = {
    "bottom_right": {
        "x_min": 0.90,
        "x_max": 0.98,
        "y_min": 0.74,
        "y_max": 0.84,
    },
    "left_large": {
        "x_min": 0.75,
        "x_max": 0.89,
        "y_min": 0.80,
        "y_max": 0.88,
    },
}


def _delete_path(path):
    if os.path.isdir(path):
        shutil.rmtree(path)
        print(f"Deleted folder: {path}")
    elif os.path.isfile(path):
        os.remove(path)
        print(f"Deleted file: {path}")


def _delete_profile(profile):
    paths = get_profile_paths(profile)
    targets = [
        paths["logs_dir"],
        paths["models_dir"],
        paths["captures_dir"],
        paths["config_path"],
    ]
    if profile != "default":
        targets.append(paths["swf_path"])
        targets.append(paths["run_root"])
    for path in targets:
        if os.path.exists(path):
            _delete_path(path)


def _load_episode_rows(profile_name):
    csv_path = get_profile_paths(profile_name)["episodes_csv"]
    if not os.path.exists(csv_path):
        return []
    with open(csv_path, "r", encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


def _episode_num(row):
    try:
        return int(float(row.get("Episode", 0)))
    except (TypeError, ValueError):
        return 0


def _gems_num(row):
    try:
        return int(float(row.get("Gems", 0)))
    except (TypeError, ValueError):
        return 0


def _replay_payload(profile_name, episode_number):
    replay_path = os.path.join(get_profile_paths(profile_name)["replays_dir"], f"episode_{episode_number}.json")
    if not os.path.exists(replay_path):
        return None
    try:
        with open(replay_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def _character_positions_from_payload(payload, character_name):
    if not payload:
        return []
    positions = []
    for step in payload.get("steps", []):
        character = step.get(character_name) or {}
        x = character.get("x")
        y = character.get("y")
        if x is None or y is None:
            continue
        try:
            positions.append((float(x), float(y)))
        except (TypeError, ValueError):
            continue
    return positions


def _payload_hits_any_ledge(payload):
    for bounds in LEDGE_REGION_BOUNDS.values():
        for character_name in ("fireboy", "watergirl"):
            for x, y in _character_positions_from_payload(payload, character_name):
                if bounds["x_min"] <= x <= bounds["x_max"] and bounds["y_min"] <= y <= bounds["y_max"]:
                    return True
    return False


def _profile_summary(profile_name):
    rows = _load_episode_rows(profile_name)
    total_episodes = len(rows)
    best_gems = max((_gems_num(row) for row in rows), default=0)
    best_gem_episodes = sum(1 for row in rows if _gems_num(row) == best_gems) if rows else 0
    ledge_episodes = 0
    for row in rows:
        episode_number = _episode_num(row)
        if episode_number <= 0:
            continue
        if _payload_hits_any_ledge(_replay_payload(profile_name, episode_number)):
            ledge_episodes += 1
    return {
        "total_episodes": total_episodes,
        "best_gems": best_gems,
        "best_gem_episodes": best_gem_episodes,
        "ledge_episodes": ledge_episodes,
    }


def _parse_profile_selection(choice, profile_summaries):
    if not choice:
        return []
    selected_profiles = []
    seen_profiles = set()
    by_index = {str(index): profile for index, (profile, _summary) in enumerate(profile_summaries, start=1)}
    by_name = {profile: profile for profile, _summary in profile_summaries}
    tokens = [token.strip() for token in choice.split(",")]
    for token in tokens:
        if not token:
            continue
        profile = by_index.get(token) or by_name.get(token)
        if profile is None:
            return None
        if profile not in seen_profiles:
            seen_profiles.add(profile)
            selected_profiles.append(profile)
    return selected_profiles


def _prompt_for_profiles():
    profiles = [name for name in list_profile_names() if name != "default"]
    if not profiles:
        print("No model profiles found.")
        return None
    print("--- Fireboy AI Delete Model Utility ---")
    profile_summaries = []
    for profile in profiles:
        summary = _profile_summary(profile)
        profile_summaries.append((profile, summary))

    print("Available profiles:")
    for index, (profile, summary) in enumerate(profile_summaries, start=1):
        print(
            f"  {index}. {format_profile_label(profile)}: "
            f"most_gems={summary['best_gems']}, "
            f"most_gem_episodes={summary['best_gem_episodes']}, "
            f"ledge_eps={summary['ledge_episodes']}, "
            f"total_eps={summary['total_episodes']}"
        )

    choice = input("Which model(s) do you want to delete? Use numbers or names, comma-separated: ").strip()
    selected_profiles = _parse_profile_selection(choice, profile_summaries)
    if selected_profiles:
        return selected_profiles
    print("Invalid selection.")
    return None


def delete_model(profile=None):
    if isinstance(profile, str):
        available_profiles = [(name, None) for name in list_profile_names() if name != "default"]
        selected_profiles = _parse_profile_selection(profile, available_profiles)
    elif isinstance(profile, (list, tuple, set)):
        selected_profiles = list(profile)
    else:
        selected_profiles = _prompt_for_profiles()

    if not selected_profiles:
        return

    selected_label = selected_profiles[0] if len(selected_profiles) == 1 else ", ".join(selected_profiles)
    confirm = input(
        f"This will permanently DELETE {len(selected_profiles)} model profile(s) ({selected_label}), "
        f"including settings, SWF copies, logs, captures, and weights. Continue? (y/n): "
    ).strip().lower()
    if confirm != "y":
        print("Delete cancelled.")
        return

    for selected in selected_profiles:
        _delete_profile(selected)
    print(f"\nDeleted {len(selected_profiles)} model profile(s).")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--profile", help="Profile name to delete")
    args = parser.parse_args()
    delete_model(args.profile)