"""Orchestrate the staged curriculum (see TRAINING.md's "Curriculum training"
section): run each stage, evaluate the result against a reference bot, and
either advance to the next stage or retry the same one.

State is persisted to curriculum_state.json (committed to git) so the script
is safe to Ctrl-C and re-run — it picks up exactly where it left off. Each
attempt reuses run_training.sh (pull, train, export, commit+push) so every
attempt's checkpoint, log, and exported policy is versioned like any other
run; this script additionally evaluates the result and commits the updated
eval_history.json + curriculum_state.json.

The gate is deliberately lenient ("block only on a clear regression," not
"require improvement") — see TRAINING.md. A 40-episode eval can call a real
improvement a regression on sample noise alone (this happened with run11:
it was the first model to deliberately score, but lost its head-to-head
evals). A strict improvement-required gate would have retried that stage
forever for the wrong reason. When a stage does fail MAX_RETRIES times in a
row, the script stops and asks for a human look rather than retrying
indefinitely or silently advancing past a bad stage.

This is generation 4 of the curriculum — a full redesign, not a patch.
Generations 1-3 (archived in curriculum_state_gen1.json/_gen2.json/_gen3.json)
all tried teaching full 3D flight by training grounded first and then
opening up vertical/pitch-roll authority (a hard 0/1 mask in gen 1/2, a
gradual float ramp in gen 3) on top of a continuous Gaussian action space.
All three failed: gen 1's hard mask went 25% -> 20% -> 15% win rate across 3
attempts; gen 2's single-flip retune landed at a stable 32%/28%/31%; gen 3's
gradual ramp landed at 29%/30%/24% — actually the worst of the three by its
final attempt. Every attempt showed the same signature regardless of
mechanism: PPO's Gaussian action-distribution std collapsed from ~0.30 to
~0.13-0.15 within the first ~10% of steps and never recovered. The root
cause: hovering this ship (mass 5.0, vertical_thrust 120, default gravity
9.8 — see ship.gd) requires *holding* thrust.y ~= 0.408 continuously; a
collapsed near-zero-mean Gaussian can brush that value but never sustain it
long enough to earn the reward gradient that would move the mean. No amount
of gating *when* the axis is allowed to act fixes a problem in *how* the
policy represents a decision on it.

Generation 4 (see TRAINING.md and Game/scripts/ship_action_codec.gd)
replaces the action space itself with per-axis MultiDiscrete bins instead of
a continuous Gaussian, trains the full action space from step 1 with no
grounded stage at all (no successful self-play RL bot in this problem class
gates control authority — see the RLGym/RLBot research cited in
TRAINING.md), and adds a state-setter "air drill" episode-start branch
(training_mode.gd's air_drill_chance) to force aerial practice instead of
relying on reward-driven exploration alone. All of generation 1-3's
checkpoint-lineage machinery (FOUNDATION_EXPERIMENT, locomotion-groundedness
tracking, resume/reference overrides for skipping a regressed branch) is
gone because there is nothing to resume from: every prior checkpoint is a
different, incompatible action/observation shape. The two strongest prior
artifacts are kept as fixed evaluation references instead (see
PROMOTED_EASY/PROMOTED_REFERENCE_GROUNDED below) — they remain playable
opponents forever via PolicyNetwork's format-versioned JSON even though
their own checkpoints and generation are gone.

Every experiment name this script generates is timestamped
(YYYYMMDD-HHMM-<name>, applied once in run_stage_attempt) so runs stay
unique across restarts/generations and sort chronologically in TensorBoard
and checkpoints/.

Usage:
    .venv/bin/python curriculum.py                  # run/resume the curriculum
    .venv/bin/python curriculum.py --seed-checkpoint checkpoints/some/final.zip
    .venv/bin/python curriculum.py --force-retry     # after fixing something, retry the blocked stage
    .venv/bin/python curriculum.py --skip-to-next-stage  # human judgment call: good enough, move on anyway

Typically started via curriculum.sh, which runs this in a detached tmux
session the way start_training.sh does for a single run.
"""

import argparse
import json
import pathlib
import subprocess
import sys
from datetime import datetime

TRAINING_DIR = pathlib.Path(__file__).resolve().parent
STATE_PATH = TRAINING_DIR / "curriculum_state.json"
EVAL_HISTORY_PATH = TRAINING_DIR / "eval_history.json"

# Fixed evaluation references — never touched by training scripts (see
# TRAINING.md's "Promoted bots" section) — kept forever as playable
# opponents via PolicyNetwork's format-versioned JSON even after their own
# checkpoints/generation are gone. The final report (not a gate) evaluates
# generation 4's result against both.
PROMOTED_EASY = TRAINING_DIR.parent / "Game" / "bots" / "promoted" / "easy.json"
PROMOTED_REFERENCE_GROUNDED = TRAINING_DIR.parent / "Game" / "bots" / "promoted" / "reference-grounded.json"

MAX_RETRIES = 2
EVAL_EPISODES = 100
# "Clear regression" = the reference beats the candidate by at least this
# many percentage points of win rate. Below this, noise in a 100-episode
# sample is a more likely explanation than the stage actually failing (see
# module docstring) — advance rather than retry.
REGRESSION_MARGIN = 0.15

# Standing flags applied to every attempt. Generation 3's "--reset-std 0.3"
# (a one-shot shock, and meaningless anyway under MultiDiscrete — there is
# no log_std) is gone; EntropyFloorCallback (see train.py) is a continuous
# controller instead, which every generation's TensorBoard data argues is
# what was actually needed (a single reset at attempt start reliably decayed
# away within ~10% of steps, every time). --ent-coef raised an order of
# magnitude from generation 3's 0.001: that value was tuned for a Gaussian's
# unbounded differential entropy, not MultiDiscrete's bounded (~10-nat)
# entropy.
STANDING_ARGS = ["--ent-coef", "0.01", "--entropy-floor"]

# Reward-shaping flags shared by every stage so the studied variables (state
# mix, opponent mode) stay isolated — carried forward unchanged from
# generation 2/3, which the reward-farmability analysis in TRAINING.md
# confirmed were never the actual problem. draw_penalty and airborne_penalty
# are deliberately NOT overridden here (both default to 0.0 in
# training_mode.gd/ship_ai_controller.gd): generation 3's draw_penalty=5 and
# airborne_penalty ramping up in lockstep with the unmask ramp were both
# grounded-era, anti-flight pressures that have no place in a curriculum
# whose entire point is teaching flight.
_SHARED_REWARD_FLAGS = [
    "--velocity-to-ball-weight", "0.08",
    "--ball-distance-penalty", "0.01",
    "--ball-touch-reward", "0.7",
    "--ball-velocity-to-goal-weight", "0.06",
    "--goal-reward", "80",
]

# A stage dict may additionally set "abort_if": {"metric": "rollout/airborne_
# fraction", "below": 0.05, "at_steps": N} to end that attempt early if a
# flight-telemetry metric (see train.py's FlightTelemetryCallback) hasn't
# cleared a bar by N *absolute* PPO timesteps (model.num_timesteps keeps
# accumulating across --resume, so N must account for whatever this stage
# inherits from its predecessor, not just this stage's own budget).
# Deliberately unset on every stage below for now — rung 5 of TRAINING.md's
# validation ladder (a short controlled A/B) should establish what a
# sensible threshold actually looks like before any stage bets a real 12h+
# budget on a guessed one.
STAGES = [
    {
        "name": "bootstrap",
        # Stage 1/3: empty-net finishing practice from a random policy — no
        # live opponent, so the full action space's first behaviour to
        # emerge is "fly to ball, push it toward the net" without a moving
        # target complicating credit assignment. Generation 1's own stage 1
        # (also inert-opponent, also empty-net) was the one stage across all
        # 3 prior generations that unambiguously passed on its first
        # attempt — reusing that shape here, just with the full action space
        # live instead of yaw-only.
        "flags": [
            "--opponent-mode", "inert",
            "--attack-goal-bias", "1.0",
            "--kickoff-chance", "0.10",
            "--near-goal-chance", "0.50",
            "--air-drill-chance", "0.20",
            *_SHARED_REWARD_FLAGS,
        ],
        "gated": False,  # ungated waypoint: trains, checkpoints, always advances — no eval
        "timesteps": 40_000_000,  # ~4h at the standing n-parallel/speedup
    },
    {
        "name": "selfplay",
        # Stage 2/3: this is where essentially all of the actual learning
        # happens. Self-play (not frozen) as the main regime — it's what
        # scales and what Necto/Nexto-class bots actually use; a frozen
        # target this early would cap skill at "exploits one specific bot"
        # instead of a moving, improving target. air_drill_chance stays on
        # at a constant rate throughout (not introduced as a later stage) —
        # gating *when* a skill is drilled reproduces the exact "gate what
        # the policy is allowed to do" pattern that failed 3 generations in
        # a row; only the state mix should vary between stages, never what
        # the policy can act on.
        "flags": [
            "--opponent-mode", "self_play",
            "--kickoff-chance", "0.15",
            "--near-goal-chance", "0.25",
            "--air-drill-chance", "0.25",
            *_SHARED_REWARD_FLAGS,
        ],
        "gated": True,
        "timesteps": 160_000_000,  # ~16h
    },
    {
        "name": "gauntlet",
        # Stage 3/3: a stationary opponent (this stage's own predecessor's
        # export) gives a low-variance measurement — important when the gate
        # is a 100-episode sample with a lenient 15-point margin — and
        # catches a self-play fixed point: a policy that only learned to
        # beat itself will look fine in stage 2 and stall here.
        # opponent_model_from_previous_stage resolves --opponent-model at
        # run time to whatever stage 2's own passing export turns out to be
        # (see run_stage_attempt) rather than a hardcoded name.
        "flags": [
            "--opponent-mode", "frozen",
            "--kickoff-chance", "0.15",
            "--near-goal-chance", "0.25",
            "--air-drill-chance", "0.25",
            *_SHARED_REWARD_FLAGS,
        ],
        "gated": True,
        "timesteps": 120_000_000,  # ~12h
        "opponent_model_from_previous_stage": True,
    },
]


def load_state() -> dict:
    if STATE_PATH.exists():
        return json.loads(STATE_PATH.read_text())
    return {"stage_index": 0, "attempt": 0, "status": "in_progress", "log": []}


def save_state(state: dict) -> None:
    STATE_PATH.write_text(json.dumps(state, indent=2) + "\n")


def experiment_name(stage_index: int, attempt: int) -> str:
    name = f"curric-s{stage_index + 1}-{STAGES[stage_index]['name']}"
    return name if attempt == 0 else f"{name}-retry{attempt}"


def _logged_experiment_name(stage_index: int, attempt: int) -> str:
    """The actual (timestamped) experiment name recorded when this attempt
    ran — needed anywhere a *past* attempt's real name matters, since
    experiment_name() alone no longer identifies a run on disk (see
    run_stage_attempt's timestamp prefix)."""
    state = load_state()
    for entry in state["log"]:
        if entry["stage_index"] == stage_index and entry["attempt"] == attempt:
            return entry["experiment"]
    raise RuntimeError(f"No logged experiment for stage {stage_index} attempt {attempt}")


def _passing_experiment_for_stage(stage_index: int) -> str:
    state = load_state()
    for entry in state["log"]:
        if entry["stage_index"] == stage_index and entry["decision"] == "pass":
            return entry["experiment"]
    raise RuntimeError(f"No passing attempt recorded for stage {stage_index} ({STAGES[stage_index]['name']})")


def resume_checkpoint(stage_index: int, attempt: int, seed_checkpoint: str | None) -> str | None:
    if attempt > 0:
        # Retry: keep training the same stage's own last attempt. No
        # per-stage "reset to a clean upstream checkpoint" override in
        # generation 4 (unlike generation 3's "unmask" stage) — nothing yet
        # suggests a generation-4 retry needs that; add one if a stage's
        # retries turn out to be drifting rather than converging.
        prev = _logged_experiment_name(stage_index, attempt - 1)
        return str(TRAINING_DIR / "checkpoints" / prev / "final.zip")
    if stage_index == 0:
        # Deliberately fresh unless --seed-checkpoint says otherwise: full
        # action space live from step 1, nothing to inherit — every prior
        # generation's checkpoints are a different, incompatible
        # action/observation shape (see module docstring).
        return seed_checkpoint
    prev_experiment = _passing_experiment_for_stage(stage_index - 1)
    return str(TRAINING_DIR / "checkpoints" / prev_experiment / "final.zip")


def reference_bot(stage_index: int) -> str:
    """Only called for gated stages (stage 0 is ungated) — the previous
    stage's own passing export, exactly like every prior generation's
    default chaining."""
    prev_experiment = _passing_experiment_for_stage(stage_index - 1)
    return str(TRAINING_DIR.parent / "Game" / "bots" / f"{prev_experiment}.json")


def run_stage_attempt(stage_index: int, attempt: int, args) -> str:
    # Timestamped so names stay unique across restarts/generations and sort
    # chronologically in TensorBoard/checkpoints — see module docstring.
    exp = f"{datetime.now().strftime('%Y%m%d-%H%M')}-{experiment_name(stage_index, attempt)}"
    resume = resume_checkpoint(stage_index, attempt, args.seed_checkpoint)
    timesteps = STAGES[stage_index].get("timesteps", args.timesteps)
    cmd = [
        "./run_training.sh", exp,
        "--timesteps", str(timesteps),
        "--n-parallel", str(args.n_parallel),
        "--speedup", str(args.speedup),
        *STANDING_ARGS,
        *STAGES[stage_index]["flags"],
    ]
    if resume:
        cmd += ["--resume", resume]
    if STAGES[stage_index].get("opponent_model_from_previous_stage"):
        prev_experiment = _passing_experiment_for_stage(stage_index - 1)
        opponent_model = TRAINING_DIR.parent / "Game" / "bots" / f"{prev_experiment}.json"
        cmd += ["--opponent-model", str(opponent_model)]
    abort_if = STAGES[stage_index].get("abort_if")
    if abort_if:
        cmd += [
            "--abort-metric", abort_if["metric"],
            "--abort-below", str(abort_if["below"]),
            "--abort-at-steps", str(abort_if["at_steps"]),
        ]
    print(f"\n=== Stage {stage_index + 1}/{len(STAGES)} ({STAGES[stage_index]['name']}), "
          f"attempt {attempt + 1}/{MAX_RETRIES + 1}: {exp} ===")
    print(" ".join(cmd))
    subprocess.run(cmd, cwd=TRAINING_DIR, check=True)
    return exp


def evaluate_attempt(experiment: str, reference: str, episodes: int) -> dict:
    candidate = TRAINING_DIR.parent / "Game" / "bots" / f"{experiment}.json"
    cmd = [".venv/bin/python", "evaluate.py", str(candidate), reference, "--episodes", str(episodes)]
    print(" ".join(cmd))
    subprocess.run(cmd, cwd=TRAINING_DIR, check=True)
    history = json.loads(EVAL_HISTORY_PATH.read_text())
    return history[-1]


def decide(record: dict) -> str:
    win_rate_candidate = record["wins_a"] / record["episodes"]
    win_rate_reference = record["wins_b"] / record["episodes"]
    if win_rate_reference - win_rate_candidate >= REGRESSION_MARGIN:
        return "fail"
    return "pass"


def final_report(experiment: str) -> None:
    """Not a gate — the two numbers that actually answer "did generation 4
    work?" (see TRAINING.md). promoted/easy.json is the shipped bot;
    promoted/reference-grounded.json (a copy of generation 3's
    curric-s5-aggression, made before the flat Game/bots/ dump was scrapped)
    is the strongest grounded-era artifact and the yardstick generations 1-3
    were all measured against. reference-grounded.json was trained with the
    locomotion mask on, so needs --grounded-b; easy.json was itself promoted
    from a *failed* unmask stage (curric-s6-unmask) and is full 3D like
    every generation-4 candidate, so needs no flag."""
    candidate = TRAINING_DIR.parent / "Game" / "bots" / f"{experiment}.json"
    print("\n=== Curriculum complete — final report (informational, not a gate) ===")
    for label, reference, extra_flags in [
        ("promoted/easy.json (shipped bot)", PROMOTED_EASY, []),
        ("promoted/reference-grounded.json (strongest grounded-era bot)", PROMOTED_REFERENCE_GROUNDED, ["--grounded-b"]),
    ]:
        if not reference.exists():
            print(f"  vs {label}: skipped, file not found")
            continue
        cmd = [".venv/bin/python", "evaluate.py", str(candidate), str(reference), "--episodes", str(EVAL_EPISODES), *extra_flags]
        print(" ".join(cmd))
        subprocess.run(cmd, cwd=TRAINING_DIR, check=True)
        record = json.loads(EVAL_HISTORY_PATH.read_text())[-1]
        print(f"  vs {label}: {record['wins_a']}-{record['wins_b']} ({record['draws']} draws), "
              f"win rate {record['win_rate_a']:.0%}")


def commit_progress(experiment: str) -> None:
    subprocess.run(["git", "add", "curriculum_state.json", "eval_history.json"], cwd=TRAINING_DIR, check=True)
    result = subprocess.run(["git", "diff", "--cached", "--quiet"], cwd=TRAINING_DIR)
    if result.returncode == 0:
        return
    subprocess.run(
        ["git", "commit", "-m", f"chore(training): curriculum progress after {experiment}"],
        cwd=TRAINING_DIR, check=True,
    )
    subprocess.run(["git", "push"], cwd=TRAINING_DIR, check=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--timesteps", type=int, default=20_000_000)
    parser.add_argument("--n-parallel", type=int, default=14)
    parser.add_argument("--speedup", type=int, default=16)
    parser.add_argument(
        "--seed-checkpoint", default=None,
        help="Resume stage 1 from this checkpoint instead of training from scratch",
    )
    parser.add_argument("--force-retry", action="store_true", help="Retry a blocked stage after human review")
    parser.add_argument("--skip-to-next-stage", action="store_true", help="Human judgment call: treat the blocked stage as good enough, advance anyway")
    args = parser.parse_args()

    state = load_state()

    if state["status"] == "blocked":
        if args.skip_to_next_stage:
            print(f"Human override: advancing past stage {state['stage_index'] + 1} "
                  f"({STAGES[state['stage_index']]['name']}) despite exhausted retries.")
            skipped_experiment = _logged_experiment_name(state["stage_index"], state["attempt"])
            state["log"].append({
                "stage_index": state["stage_index"],
                "experiment": skipped_experiment,
                "attempt": state["attempt"], "decision": "pass", "override": "skip_to_next_stage",
            })
            state["stage_index"] += 1
            state["attempt"] = 0
            state["status"] = "in_progress"
            save_state(state)
            # If this was the last stage, the while loop below never runs
            # (stage_index now == len(STAGES)), so this override's state
            # change would otherwise never get committed/pushed.
            commit_progress(skipped_experiment)
        elif args.force_retry:
            print(f"Human override: retrying stage {state['stage_index'] + 1} "
                  f"({STAGES[state['stage_index']]['name']}) after review.")
            state["attempt"] += 1
            state["status"] = "in_progress"
            save_state(state)
        else:
            print(f"BLOCKED at stage {state['stage_index'] + 1} ({STAGES[state['stage_index']]['name']}) "
                  f"after {MAX_RETRIES + 1} attempts — see curriculum_state.json's log for eval results.")
            print("Re-run with --force-retry (after adjusting flags/timesteps) or "
                  "--skip-to-next-stage (advance anyway) once you've looked at why.")
            sys.exit(1)

    last_experiment = None
    while state["stage_index"] < len(STAGES):
        stage_index = state["stage_index"]
        attempt = state["attempt"]

        experiment = run_stage_attempt(stage_index, attempt, args)
        last_experiment = experiment

        if not STAGES[stage_index].get("gated", True):
            # Ungated waypoint (see the bootstrap stage): trains,
            # checkpoints, and always advances — no eval, no regression
            # gate, nothing to retry against.
            print(f"{experiment}: ungated waypoint — skipping eval, advancing unconditionally")
            state["log"].append({
                "stage_index": stage_index, "experiment": experiment, "attempt": attempt,
                "decision": "pass", "note": "ungated waypoint (no eval)",
            })
            state["stage_index"] += 1
            state["attempt"] = 0
            state["status"] = "in_progress"
            save_state(state)
            commit_progress(experiment)
            continue

        reference = reference_bot(stage_index)
        record = evaluate_attempt(experiment, reference, EVAL_EPISODES)
        decision = decide(record)

        print(f"{experiment}: candidate {record['wins_a']}-{record['wins_b']} reference "
              f"({record['draws']} draws) over {record['episodes']} episodes -> {decision}")

        state["log"].append({
            "stage_index": stage_index, "experiment": experiment, "attempt": attempt,
            "eval": record, "decision": decision,
        })

        if decision == "pass":
            state["stage_index"] += 1
            state["attempt"] = 0
            state["status"] = "in_progress"
            save_state(state)
            commit_progress(experiment)
            continue

        if attempt >= MAX_RETRIES:
            state["status"] = "blocked"
            save_state(state)
            commit_progress(experiment)
            print(f"\nBLOCKED: stage {stage_index + 1} ({STAGES[stage_index]['name']}) failed "
                  f"{MAX_RETRIES + 1} attempts in a row. Stopping for human review — see "
                  f"curriculum_state.json. Re-run with --force-retry or --skip-to-next-stage.")
            sys.exit(1)

        state["attempt"] += 1
        save_state(state)
        commit_progress(experiment)

    print("\nCurriculum complete — all stages passed.")
    state["status"] = "done"
    save_state(state)
    if last_experiment is not None:
        # None only if the loop above never ran at all (e.g. re-invoking
        # after the curriculum was already "done") — nothing new to commit
        # in that case.
        final_report(last_experiment)
        commit_progress(last_experiment)


if __name__ == "__main__":
    main()