"""Train the Cosmic Clash self-play PPO policy.

Example (smoke run):
    .venv/bin/python train.py --experiment smoke --timesteps 100000

Long run on the Linux/CUDA box:
    GODOT_BIN=~/godot/Godot_v4.7.1-stable_linux.x86_64 \
        .venv/bin/python train.py --experiment run01 --timesteps 20000000 \
        --n-parallel 6 --speedup 16

See TRAINING.md at the repo root for the full workflow.
"""

import argparse
import os
import pathlib

from gymnasium import spaces
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback, CheckpointCallback
from stable_baselines3.common.utils import safe_mean
from stable_baselines3.common.vec_env.vec_monitor import VecMonitor

from cosmic_env import CosmicClashVecEnv

TRAINING_DIR = pathlib.Path(__file__).resolve().parent
DEFAULT_GODOT_MACOS = "/Applications/Godot.app/Contents/MacOS/Godot"

# Must match Game/scripts/ship_action_codec.gd's HEADS order exactly (both
# are independently the gymnasium-sorted key order of the same 7 names) —
# training/test_action_space.py's rung-0 check asserts this. Used only for
# per-head entropy logging/reset-logits head selection below.
ACTION_HEAD_NAMES = ["rot_x", "rot_y", "rot_z", "thrust_x", "thrust_y", "thrust_z", "turbo"]


class GoalRateCallback(BaseCallback):
    """Logs rollout/goal_rate: the fraction of completed episodes in the
    current ep_info_buffer that ended in an actual goal, vs. timing out as a
    draw. rollout/ep_rew_mean mixes dense reward-shaping (ball chasing/
    touching) with the sparse terminal goal reward, so it can trend up from
    better shaping alone without the policy finishing more episodes by
    actually scoring — this isolates that. Requires VecMonitor(...,
    info_keywords=("goal_scored",)), which copies ShipAIController.get_info()
    into each completed episode's info["episode"] dict (see
    training_mode.gd's _on_goal_scored / timeout branch)."""

    def _on_step(self) -> bool:
        return True

    def _on_rollout_end(self) -> None:
        if len(self.model.ep_info_buffer) == 0:
            return
        # The vendored godot_rl sync bridge (Game/addons/godot_rl_agents/sync.gd,
        # _training_process) snapshots each agent's info dict once per tick and
        # has its own "NEEDS REFACTOR" comment on the reset-timing path, so an
        # episode's terminal info entry can arrive without "goal_scored" at all
        # (observed crashing a run after 2026-07-28). Skip those rather than
        # crash training over a monitoring-only metric.
        rates = [ep_info["goal_scored"] for ep_info in self.model.ep_info_buffer if "goal_scored" in ep_info]
        if rates:
            self.logger.record("rollout/goal_rate", safe_mean(rates))


class FlightTelemetryCallback(BaseCallback):
    """Logs flight and handling telemetry — leading indicators for curriculum generation 4's
    core hypothesis (a discrete action space lets the policy actually hold a
    sustained vertical set-point, e.g. hovering), visible from the very
    first rollout instead of only in a win-rate number measured a full
    24h+ run later, which is what made every past generation's failure mode
    expensive to diagnose. Requires VecMonitor(..., info_keywords=(...,
    "airborne_fraction", "mean_altitude", "air_touch_fraction",
    "vertical_thrust_mean")) — see ShipAIController.get_info."""

    _KEYS = (
        "airborne_fraction",
        "mean_altitude",
        "air_touch_fraction",
        "vertical_thrust_mean",
        "productive_air_touch_fraction",
        "upright_fraction",
        "forward_motion_fraction",
    )

    def _on_step(self) -> bool:
        return True

    def _on_rollout_end(self) -> None:
        if len(self.model.ep_info_buffer) == 0:
            return
        for key in self._KEYS:
            values = [ep_info[key] for ep_info in self.model.ep_info_buffer if key in ep_info]
            if values:
                self.logger.record(f"rollout/{key}", safe_mean(values))


class EntropyFloorCallback(BaseCallback):
    """Replaces the old one-shot `--reset-std` shock (meaningless under
    MultiDiscrete — there is no log_std) with a persistent controller.
    Three curriculum generations' TensorBoard runs all show the same
    signature: exploration (train/std, under the previous continuous
    Gaussian) collapsing within the first ~10% of steps and never
    recovering from a single reset applied at attempt start. A controller
    that responds every rollout instead of once should not have that decay-
    and-stay-collapsed failure mode.

    Reads mean policy entropy each rollout (recomputed from a fresh
    minibatch via the same RolloutBuffer.get() plumbing PPO's own train()
    uses, since _on_rollout_end fires before that iteration's train() call)
    and nudges model.ent_coef multiplicatively toward a target that decays
    linearly from target_start_frac to target_end_frac of the action
    space's maximum possible entropy (sum of ln(n) over each MultiDiscrete
    head) over the run. PPO reads self.ent_coef fresh inside train() each
    update, so mutating it here from a callback takes effect on the very
    next update with no subclassing needed. No-ops (does nothing) for a
    non-MultiDiscrete action space, e.g. a continuous-action A/B run.

    Also logs train/entropy_head_<name> per action head — the direct
    replacement for the old aggregate train/std scalar, and strictly more
    useful: it identifies *which* axis is collapsing instead of one number
    for all seven.
    """

    def __init__(
        self,
        total_timesteps: int,
        target_start_frac: float = 0.55,
        target_end_frac: float = 0.20,
        adjust_rate: float = 1.02,
        ent_coef_bounds: tuple[float, float] = (1e-4, 0.05),
    ):
        super().__init__()
        self.total_timesteps = total_timesteps
        self.target_start_frac = target_start_frac
        self.target_end_frac = target_end_frac
        self.adjust_rate = adjust_rate
        self.ent_coef_bounds = ent_coef_bounds
        self._is_multi_discrete = False
        self._h_max = 0.0
        self._start_timesteps = 0

    def _on_training_start(self) -> None:
        import numpy as np

        self._is_multi_discrete = isinstance(self.model.action_space, spaces.MultiDiscrete)
        if self._is_multi_discrete:
            self._h_max = float(np.sum(np.log(self.model.action_space.nvec)))
        # this call's own timesteps budget, not the resumed total — model.
        # num_timesteps keeps accumulating across --resume calls, but
        # total_timesteps below is this invocation's --timesteps.
        self._start_timesteps = self.model.num_timesteps

    def _on_step(self) -> bool:
        return True

    def _on_rollout_end(self) -> None:
        if not self._is_multi_discrete:
            return
        import torch as th

        batch = next(self.model.rollout_buffer.get(batch_size=self.model.batch_size))
        with th.no_grad():
            distribution = self.model.policy.get_distribution(batch.observations)
        per_head = getattr(distribution, "distribution", None)
        if per_head is None:
            return

        entropies = [dist.entropy().mean().item() for dist in per_head]
        for name, entropy in zip(ACTION_HEAD_NAMES, entropies):
            self.logger.record(f"train/entropy_head_{name}", entropy)

        mean_entropy = sum(entropies)
        progress = min((self.model.num_timesteps - self._start_timesteps) / self.total_timesteps, 1.0)
        target_frac = self.target_start_frac + (self.target_end_frac - self.target_start_frac) * progress
        target = target_frac * self._h_max
        if mean_entropy < target:
            self.model.ent_coef = min(self.model.ent_coef * self.adjust_rate, self.ent_coef_bounds[1])
        else:
            self.model.ent_coef = max(self.model.ent_coef / self.adjust_rate, self.ent_coef_bounds[0])
        self.logger.record("train/ent_coef_adaptive", self.model.ent_coef)


class AbortIfCallback(BaseCallback):
    """Optional kill criterion (see curriculum.py's per-stage `abort_if`):
    ends model.learn() early once `metric` (a rollout/* key logged by
    FlightTelemetryCallback — must run earlier in the callback list so the
    value exists by the time this checks it) is below `below` at or past
    `at_steps`. Stops via SB3's own "_on_step returning False halts
    training" contract rather than an exception, so the enclosing
    try/finally in main() still runs and saves/exports/commits whatever
    checkpoint exists — an aborted stage still leaves a usable, logged
    artifact instead of either running a doomed stage to completion
    unattended or leaving one stranded and uncommitted.
    """

    def __init__(self, metric: str, below: float, at_steps: int):
        super().__init__()
        self.metric = metric
        self.below = below
        self.at_steps = at_steps
        self._checked = False
        self._should_stop = False

    def _on_step(self) -> bool:
        return not self._should_stop

    def _on_rollout_end(self) -> None:
        if self._checked or self.model.num_timesteps < self.at_steps:
            return
        self._checked = True
        value = self.logger.name_to_value.get(self.metric)
        if value is not None and value < self.below:
            print(
                f"AbortIfCallback: {self.metric}={value:.4f} < {self.below} "
                f"at {self.model.num_timesteps} steps — stopping early"
            )
            self._should_stop = True


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--godot_bin",
        default=os.environ.get("GODOT_BIN", DEFAULT_GODOT_MACOS),
        help="Path to the Godot binary (or set GODOT_BIN)",
    )
    parser.add_argument(
        "--exported-binary",
        default=None,
        help="Path to a pre-built game executable (see export_linux.sh) instead of running the "
        "project from source — skips per-instance script/resource import for faster parallel "
        "startup. Overrides --godot_bin when set.",
    )
    parser.add_argument("--experiment", default="default", help="Run name for logs/checkpoints")
    parser.add_argument("--timesteps", type=int, default=200_000)
    parser.add_argument("--n-parallel", type=int, default=2, help="Parallel Godot instances (2 agents each)")
    parser.add_argument("--speedup", type=int, default=8, help="Physics speedup factor inside Godot")
    parser.add_argument("--port", type=int, default=11008, help="Base TCP port (one per instance)")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--resume", default=None, help="Checkpoint .zip to resume from")
    parser.add_argument(
        "--ent-coef",
        type=float,
        default=0.01,
        help="Entropy bonus coefficient (applied on resume too). Raised from 0.0001 for curriculum "
        "generation 4: that value was tuned for a continuous Gaussian's differential entropy "
        "(unbounded, can go negative); MultiDiscrete entropy is bounded (~10 nats for this action "
        "space) and needs an order of magnitude more coefficient to matter. See --entropy-floor.",
    )
    parser.add_argument("--n-steps", type=int, default=256, help="Rollout length per env between updates (applied on resume too)")
    parser.add_argument("--batch-size", type=int, default=256, help="PPO minibatch size (applied on resume too)")
    parser.add_argument(
        "--reset-logits",
        type=float,
        default=None,
        help="On resume, multiply the policy's action_net weights/bias by this scale (e.g. 0.1), "
        "pulling every head's softmax back toward uniform without discarding learned features — "
        "the MultiDiscrete analogue of the old continuous --reset-std. Combine with "
        "--reset-logits-heads to reset only specific heads.",
    )
    parser.add_argument(
        "--reset-logits-heads",
        default=None,
        help=f"Comma-separated subset of {ACTION_HEAD_NAMES} to apply --reset-logits to (default: all heads)",
    )
    parser.add_argument(
        "--entropy-floor",
        action="store_true",
        help="Enable EntropyFloorCallback: a persistent per-rollout controller nudging ent_coef to "
        "hold policy entropy near a decaying target, replacing the one-shot --reset-std/"
        "--reset-logits shock as the primary exploration mechanism (that flag remains for "
        "resume-time recovery after a diagnosed collapse; this runs continuously).",
    )
    parser.add_argument(
        "--checkpoint-every", type=int, default=10_000_000,
        help="Timesteps between checkpoints. Raised from 100_000 for curriculum generation 4: at the "
        "old value a single 240M-step stage wrote ~2400 intermediate checkpoint files (only final.zip "
        "is ever committed, see .gitignore/run_training.sh, but they still accumulate in the working "
        "tree during the run).",
    )
    parser.add_argument("--viz", action="store_true", help="Show game windows (debugging; slow)")
    parser.add_argument("--wandb", action="store_true", help="Also log to Weights & Biases")
    parser.add_argument(
        "--abort-metric", default=None,
        help="Optional kill criterion (see curriculum.py's per-stage abort_if): a rollout/* metric name to watch",
    )
    parser.add_argument("--abort-below", type=float, default=None, help="Stop early if --abort-metric drops below this")
    parser.add_argument(
        "--abort-at-steps", type=int, default=None,
        help="Don't check --abort-metric until at least this many timesteps have elapsed",
    )

    curriculum = parser.add_argument_group(
        "curriculum", "Stage the training run — see TRAINING.md's Curriculum training section"
    )
    curriculum.add_argument(
        "--opponent-mode",
        choices=["self_play", "inert", "frozen", "league"],
        default=None,
        help="self_play (default): both ships are live trainees. inert: team 1 is a "
        "do-nothing placeholder (isolated scoring practice). frozen: team 1 runs a "
        "fixed exported policy (--opponent-model); league: sample a fixed policy per episode "
        "from --opponent-pool",
    )
    curriculum.add_argument("--opponent-model", default=None, help="Exported policy .json for --opponent-mode=frozen")
    curriculum.add_argument(
        "--opponent-pool", default=None,
        help="Comma-separated exported policy paths for --opponent-mode=league; one is sampled per episode",
    )
    curriculum.add_argument(
        "--draw-penalty", type=float, default=None, help="One-time penalty when an episode times out with no goal"
    )
    curriculum.add_argument(
        "--attack-goal-bias",
        type=float,
        default=None,
        help="0.5 = uniform between both goals (default); 1.0 = near-goal resets always target the goal team 0 attacks",
    )
    curriculum.add_argument("--kickoff-chance", type=float, default=None, help="Overrides kickoff_state_chance")
    curriculum.add_argument("--near-goal-chance", type=float, default=None, help="Overrides ball_near_goal_chance")
    curriculum.add_argument(
        "--air-drill-chance", type=float, default=None,
        help="Overrides air_drill_chance: ball spawned high, both ships spawned low and lateral — "
        "unsolvable without climbing (curriculum generation 4's state-setter aerial curriculum)",
    )
    curriculum.add_argument(
        "--air-intercept-chance", type=float, default=None,
        help="Moving high-ball interception starts aimed at a real goal (generation-5 aerial stage)",
    )
    curriculum.add_argument(
        "--team-size", type=int, choices=range(1, 6), default=None,
        help="Ships per team (1-5); generation-5 automated stages remain 1v1 until 2v2 evaluation exists",
    )
    curriculum.add_argument(
        "--tilt-penalty", type=float, default=None,
        help="Overrides ShipAIController.tilt_penalty (dense per-tick cost scaled by non-upright tilt)",
    )
    curriculum.add_argument(
        "--velocity-to-ball-weight", type=float, default=None,
        help="Overrides ShipAIController.velocity_to_ball_weight (dense reward for closing speed toward the ball)",
    )
    curriculum.add_argument(
        "--forward-velocity-to-ball-weight", type=float, default=None,
        help="Low-altitude dense reward for nose-led planar approach toward the ball",
    )
    curriculum.add_argument(
        "--ball-distance-penalty", type=float, default=None,
        help="Overrides ShipAIController.ball_distance_penalty (dense per-tick cost scaled by distance to the ball)",
    )
    curriculum.add_argument(
        "--ball-touch-reward", type=float, default=None,
        help="Overrides ShipAIController.ball_touch_reward (event reward on ball contact, cooldown-gated)",
    )
    curriculum.add_argument(
        "--airborne-penalty", type=float, default=None,
        help="Overrides ShipAIController.airborne_penalty (dense per-tick cost scaled by height above the floor)",
    )
    curriculum.add_argument(
        "--ground-tilt-penalty", type=float, default=None,
        help="Low-altitude-only tilt cost that fades to zero by the handling-height threshold",
    )
    curriculum.add_argument(
        "--non-forward-penalty", type=float, default=None,
        help="Overrides ShipAIController.non_forward_penalty (low-altitude dense cost on sideways/reverse "
        "planar velocity, independent of the ball)",
    )
    curriculum.add_argument(
        "--speed-reward-weight", type=float, default=None,
        help="Overrides the orientation-agnostic own-speed reward (generation 5 handling sets it to zero)",
    )
    curriculum.add_argument(
        "--ball-velocity-to-goal-weight", type=float, default=None,
        help="Overrides ShipAIController.ball_velocity_to_goal_weight (dense reward for the ball's velocity toward the attack goal)",
    )
    curriculum.add_argument(
        "--goal-reward", type=float, default=None,
        help="Overrides TrainingMode.goal_reward (terminal reward for actually scoring)",
    )

    return parser.parse_args()


def _curriculum_kwargs(args) -> dict:
    """Maps train.py's curriculum flags to the --key=value args training_mode.gd's
    _parse_curriculum_args() reads, omitting anything not explicitly passed so
    unset flags leave Godot's own @export defaults in place."""
    mapping = {
        "opponent_mode": args.opponent_mode,
        "opponent_model": args.opponent_model,
        "opponent_model_pool": args.opponent_pool,
        "draw_penalty": args.draw_penalty,
        "attack_goal_bias": args.attack_goal_bias,
        "kickoff_state_chance": args.kickoff_chance,
        "ball_near_goal_chance": args.near_goal_chance,
        "air_drill_chance": args.air_drill_chance,
        "air_intercept_chance": args.air_intercept_chance,
        "team_size": args.team_size,
        "ai_tilt_penalty": args.tilt_penalty,
        "ai_ground_tilt_penalty": args.ground_tilt_penalty,
        "ai_non_forward_penalty": args.non_forward_penalty,
        "ai_velocity_to_ball_weight": args.velocity_to_ball_weight,
        "ai_forward_velocity_to_ball_weight": args.forward_velocity_to_ball_weight,
        "ai_ball_distance_penalty": args.ball_distance_penalty,
        "ai_ball_touch_reward": args.ball_touch_reward,
        "ai_airborne_penalty": args.airborne_penalty,
        "ai_speed_reward_weight": args.speed_reward_weight,
        "ai_ball_velocity_to_goal_weight": args.ball_velocity_to_goal_weight,
        "goal_reward": args.goal_reward,
    }
    return {key: value for key, value in mapping.items() if value is not None}


def main():
    args = parse_args()
    log_dir = TRAINING_DIR / "logs"
    checkpoint_dir = TRAINING_DIR / "checkpoints" / args.experiment
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    if args.wandb:
        import wandb

        wandb.init(project="cosmic-clash-rl", name=args.experiment, sync_tensorboard=True)

    exported_binary = args.exported_binary or None
    env = CosmicClashVecEnv(
        godot_bin=exported_binary or args.godot_bin,
        exported=exported_binary is not None,
        n_parallel=args.n_parallel,
        seed=args.seed,
        port=args.port,
        show_window=args.viz,
        speedup=args.speedup,
        **_curriculum_kwargs(args),
    )
    env = VecMonitor(
        env,
        info_keywords=(
            "goal_scored",
            "airborne_fraction",
            "mean_altitude",
            "air_touch_fraction",
            "vertical_thrust_mean",
            "productive_air_touch_fraction",
            "upright_fraction",
            "forward_motion_fraction",
        ),
    )

    if args.resume:
        model = PPO.load(
            args.resume,
            env=env,
            tensorboard_log=str(log_dir),
            ent_coef=args.ent_coef,
            n_steps=args.n_steps,
            batch_size=args.batch_size,
        )
        print(
            f"Resumed from {args.resume} at {model.num_timesteps} timesteps "
            f"(ent_coef={args.ent_coef}, n_steps={args.n_steps}, batch_size={args.batch_size})"
        )
        if args.reset_logits is not None:
            import torch

            heads = args.reset_logits_heads.split(",") if args.reset_logits_heads else ACTION_HEAD_NAMES
            nvec = list(model.action_space.nvec)
            offset = 0
            offsets = {}
            for name, size in zip(ACTION_HEAD_NAMES, nvec):
                offsets[name] = (offset, offset + size)
                offset += size
            with torch.no_grad():
                for name in heads:
                    start, end = offsets[name]
                    model.policy.action_net.weight[start:end].mul_(args.reset_logits)
                    model.policy.action_net.bias[start:end].mul_(args.reset_logits)
            print(f"Reset action_net logits for heads {heads} by scale {args.reset_logits}")
    else:
        model = PPO(
            "MultiInputPolicy",
            env,
            verbose=1,
            ent_coef=args.ent_coef,
            n_steps=args.n_steps,
            batch_size=args.batch_size,
            learning_rate=3e-4,
            tensorboard_log=str(log_dir),
        )

    checkpoint_callback = CheckpointCallback(
        save_freq=max(args.checkpoint_every // env.num_envs, 1),
        save_path=str(checkpoint_dir),
        name_prefix="ppo",
    )
    # Order matters for AbortIfCallback (must run after FlightTelemetryCallback
    # so the rollout/* metric it watches has already been logged this round).
    callbacks = [checkpoint_callback, GoalRateCallback(), FlightTelemetryCallback()]
    if args.entropy_floor:
        callbacks.append(EntropyFloorCallback(total_timesteps=args.timesteps))
    if args.abort_metric is not None and args.abort_below is not None and args.abort_at_steps is not None:
        callbacks.append(AbortIfCallback(args.abort_metric, args.abort_below, args.abort_at_steps))

    try:
        model.learn(
            args.timesteps,
            callback=callbacks,
            tb_log_name=args.experiment,
            reset_num_timesteps=not args.resume,
        )
    finally:
        final_path = checkpoint_dir / "final.zip"
        model.save(str(final_path))
        print(f"Saved {final_path}")
        env.close()


if __name__ == "__main__":
    main()