class_name TrainingMode
extends GameMode
const SimConstants = preload("res://scripts/sim_constants.gd")

# Headless self-play training mode: two RL-driven ships, no HUD, no camera.
# The scene also contains the godot_rl_agents Sync node, which speaks TCP to
# the Python trainer; this mode owns the environment rules — episodes, goal
# rewards, and randomized episode-start states (the RLGym "state setter"
# lesson: varied starts massively speed up learning versus kickoff-only).
#
# Run: godot --headless --path Game res://scenes/training.tscn
# (started automatically by training/train.py; boots into idle ships with a
# warning if no trainer is listening).
#
# Eval mode (used by training/evaluate.py): pass --eval_model_a=<path> and
# --eval_model_b=<path> (+ optional --eval_episodes=N) and both ships are
# instead driven by those exported policies via AIShipController; each episode
# ends at the first goal (or a draw on timeout), and a final machine-readable
# "EVAL_RESULT {...}" line is printed before quitting.
#
# Curriculum mode (used by training/train.py's --opponent-mode/--draw-penalty/
# etc. flags, see TRAINING.md): --opponent_mode=inert|frozen swaps team 1's
# live self-play agent for a do-nothing placeholder or a fixed exported
# policy; --ai_<name>=<value> and the TrainingMode-level overrides below let
# a run retune reward shaping / episode-start mix without touching script
# defaults. See _parse_curriculum_args.

@export var episode_length_seconds := 30.0
# Run01-vs-run02 eval (see training/eval_history.json) came back 87.5% draws:
# with a 30s episode, the dense per-tick terms on ShipAIController can sum to
# several times this value before it was raised, so scoring and forfeiting
# the rest of the episode's farmable reward was worse than never finishing.
# Raised well above that ceiling so a real scoring chance always beats
# continuing to farm dense reward for however long is left in the episode.
@export var goal_reward := 40.0
# One-time penalty applied to every agent when an episode times out with no
# goal scored (see _physics_process's truncation branch) — distinct from
# ShipAIController's per-tick time_penalty, which accrues regardless of
# outcome and doesn't specifically mark "this episode ended undecided."
# Default 0 (off) so ordinary runs are unaffected; curriculum stage 3 turns
# this on via --draw_penalty to teach that a draw is still a failure.
@export var draw_penalty := 0.0

# Episode-start state mix; remaining probability = fully random state.
@export_range(0.0, 1.0) var kickoff_state_chance := 0.2
# Raised from 0.2: fixing the reward incentive to score (see goal_reward,
# ball_touch_reward, time_penalty on ShipAIController) only helps if the
# policy also gets enough reps at actually finishing. At 0.2 that scenario
# was 1 in 5 episode starts; most training time was spent in generic
# midfield play where a finish never comes up.
@export_range(0.0, 1.0) var ball_near_goal_chance := 0.35
# Which goal _place_ball_near_goal() favors: 0.5 = uniform between both goals
# (default, matches historical behaviour). 1.0 = always the goal team 0
# attacks — used by curriculum stage 1 (--attack_goal_bias=1.0) so a lone
# trainee's near-goal resets are always finishing chances, not a coin flip
# between attacking and defending an empty net.
@export_range(0.0, 1.0) var attack_goal_bias := 0.5
# Fourth episode-start branch (after kickoff/near-goal, before the fully-
# random fallback): ball spawned high, both ships spawned low and lateral —
# unsolvable without climbing. Default 0 (off) so ordinary runs are
# unaffected. Added for curriculum generation 4: the existing random branch
# already samples ship/ball Y across the full arena height, but that only
# randomizes the *initial* state — under gravity+drag a floor-pinned policy
# sinks back to the floor in ~1.5s, so the *stationary* state distribution
# stayed floor-pinned even though the initial one wasn't. See
# _place_air_drill.
@export_range(0.0, 1.0) var air_drill_chance := 0.0
# Moving-ball aerial interception branch used by generation 5. Unlike the
# stationary/random air drill, the ball follows a reachable trajectory toward
# a real goal and ships start low behind/lateral to it, so a useful touch is
# naturally reinforced by the existing goal-directed ball rewards.
@export_range(0.0, 1.0) var air_intercept_chance := 0.0
# Wall-play and rebound starts are separate: wall-play begins beside a wall
# with the ball travelling inward, while rebound begins just before an
# outward wall impact. Both default off to preserve existing distributions.
@export_range(0.0, 1.0) var wall_play_chance := 0.0
@export_range(0.0, 1.0) var rebound_chance := 0.0
# Ground-start branch for the generation-5 handling stage: ships spawn level
# and resting on the floor with a low, floor-level ball. Every other branch
# samples ship Y uniformly across the full 18m volume (see _random_position),
# so ~65% of episodes previously began at a mean altitude near 8.7m — the
# measured airborne_fraction ~0.44 was largely that spawn distribution rather
# than a policy preference, and the ground-handling reward terms (which all
# fade out above 3m) barely ever applied. A stage that means to teach driving
# has to actually start the ship on the ground.
@export_range(0.0, 1.0) var ground_start_chance := 0.0

# Ships per team. Default 1 preserves every existing curriculum script's 1v1
# behaviour unchanged; up to 5 matches ShipObservations.MAX_TEAMMATES/
# MAX_OPPONENTS. Team-credit reward and paired 2v2 evaluation are opt-in;
# no teamplay training stage is enabled by default.
@export_range(1, 5) var team_size: int = 1

# Placement bounds for randomized episode starts, derived from the standard
# enclosure (ArenaBoundary). The inset keeps a randomly oriented ship
# (1.6x0.6x4 box, worst-case half-extent ~2.18) from spawning intersecting
# the walls, ceiling, or goal sensors.
const SPAWN_INSET := 2.5
const FIELD_HALF_X := ArenaBoundary.INNER_HALF_X - SPAWN_INSET
const FIELD_HALF_Z := ArenaBoundary.GOAL_LINE_Z - SPAWN_INSET
const FIELD_MIN_Y := 1.5
# Resting heights for the ground-start branch: half the ship hull's 0.6 height
# and the ball's 0.5 radius (Godot's SphereShape3D default, see ball.tscn),
# each plus a little clearance so bodies settle onto the floor instead of
# spawning interpenetrated with it.
const GROUND_START_Y := 0.35
const GROUND_START_BALL_Y := 0.55
const WALL_PLAY_BALL_CLEARANCE := 1.0
const REBOUND_BALL_CLEARANCE := 0.75
const WALL_PLAY_SPEED := Vector2(4.0, 9.0)
const FIELD_MAX_Y := ArenaBoundary.INNER_HEIGHT - SPAWN_INSET
# The corner curves reach at most their chord plane |x| + |z| = INNER_HALF_X
# + INNER_HALF_Z - CORNER_RADIUS; spawns keep the same SPAWN_INSET clearance
# from that plane as from the walls (perpendicular distance, hence the
# sqrt(2) when expressed in |x| + |z| terms). The true curve bulges outward
# from the chord, so this is conservative.
const CORNER_LIMIT := ArenaBoundary.INNER_HALF_X + ArenaBoundary.INNER_HALF_Z \
	- ArenaBoundary.CORNER_RADIUS - SPAWN_INSET * sqrt(2.0)
# Below this height a tilted ship could reach down into the wall-base
# fillets, so low spawns stay an extra BASE_RADIUS off the walls.
const FILLET_CLEAR_Y := ArenaBoundary.BASE_RADIUS + FIELD_MIN_Y
const MAX_RANDOM_BALL_SPEED := 12.0
const MAX_RANDOM_SHIP_SPEED := 8.0
# Minimum centre-to-centre separation enforced between ships placed in the
# same randomized reset (team_size > 1) — without it, _place_ships_random/
# _place_air_drill sample each ship independently and can spawn them
# interpenetrating. Twice the ~2.05m worst-case rotated half-extent noted
# above clears any relative orientation; matches arena_base.tscn's spawn
# marker spacing, which uses the same margin for the same reason.
const MIN_SHIP_SEPARATION := 4.5

# Sim runs at SimConstants.TICK_HZ physics ticks per sim-second regardless
# of speedup.
const TICKS_PER_SIM_SECOND := float(SimConstants.TICK_HZ)

var _agents: Array[ShipAIController] = []

# Eval mode state (see header comment)
var _eval := false
var _eval_models: Array[String] = ["", ""]
# Per-model locomotion mask — must match how each model was actually trained
# (see AIShipController's identical exports), so a stage 1/2 (grounded)
# candidate isn't unfairly penalized by untrained aerial noise during eval
# that its training environment never had.
var _eval_allow_vertical: Array[bool] = [true, true]
var _eval_allow_pitch_roll: Array[bool] = [true, true]
var _eval_episodes := 20
var _eval_goals := {0: 0, 1: 0}
var _eval_draws := 0
var _eval_episodes_done := 0
var _episode_ticks := 0
var _eval_team_size := 1

# Curriculum mode state (see _parse_curriculum_args). "self_play" (default)
# is today's only historical behaviour: both ships are live trainees sharing
# the policy. "inert" gives team 1 a do-nothing placeholder ship (no bot
# configured, same as MatchMode's fallback) so a lone trainee can drill
# scoring against an empty net. "frozen" gives team 1 a fixed exported
# policy via AIShipController — the same wiring the eval branch above
# already uses, just for one side of a live training episode.
var _opponent_mode := "self_play"
var _opponent_model_path := ""
var _opponent_model_pool: Array[String] = []
var _frozen_opponent_bots: Array[AIShipController] = []
# ShipAIController @export overrides collected from --ai_<name>=<value> args,
# applied to every ShipAIController this run creates (see _attach_agent).
var _ai_overrides := {}
# Ships excluded from _attach_agent (the "inert" opponent) skip randomized
# per-episode placement in _place_ships_random so they stay parked at their
# arena spawn instead of drifting into the play area as a stray obstacle.
var _inert_ships: Array[Ship] = []

# Whether this run's arena is an ELEVATED-goal variant (see
# ArenaBoundary.GoalMode) — read once _start() runs so _place_ball_near_goal
# can sample a height range matching the goal's real position. Read here
# rather than _ready(): TrainingMode has no _ready() override, and
# GameMode._ready() is what discovers `arena` before calling _start().
var _elevated := false


func _start() -> void:
	_elevated = (arena.get_node("Boundary") as ArenaBoundary).goal_mode == ArenaBoundary.GoalMode.ELEVATED
	_parse_eval_args()
	_parse_curriculum_args()
	spawn_ball()
	if _eval:
		for team in [0, 1]:
			for spawn_index in _eval_team_size:
				var bot := AIShipController.new()
				bot.model_path = _eval_models[team]
				bot.allow_vertical = _eval_allow_vertical[team]
				bot.allow_pitch_roll = _eval_allow_pitch_roll[team]
				spawn_ship(team, spawn_index, bot)
		return

	var team0_ships: Array[Ship] = []
	for i in team_size:
		team0_ships.append(spawn_ship(0, i, RLShipController.new()))

	# The opponent_mode branch applies uniformly to every ship on team 1: an
	# "inert"/"frozen"/"league" run means the whole opposing team gets that treatment,
	# not just one ship.
	var team1_ships: Array[Ship] = []
	for i in team_size:
		var ship1: Ship
		match _opponent_mode:
			"inert":
				ship1 = spawn_ship(1, i, ShipController.new())
				_inert_ships.append(ship1)
			"frozen":
				var bot := AIShipController.new()
				bot.model_path = _opponent_model_path
				ship1 = spawn_ship(1, i, bot)
				_frozen_opponent_bots.append(bot)
			"league":
				var bot := AIShipController.new()
				bot.model_path = _opponent_model_pool[0] if not _opponent_model_pool.is_empty() else ""
				ship1 = spawn_ship(1, i, bot)
				_frozen_opponent_bots.append(bot)
			_:
				ship1 = spawn_ship(1, i, RLShipController.new())
		team1_ships.append(ship1)

	# All ships spawn before any agent attaches, so every agent's
	# teammates/opponents lists see the other side's full roster.
	for ship in team0_ships:
		_attach_agent(ship, _other_ships(team0_ships, ship), team1_ships)
	if _opponent_mode == "self_play":
		for ship in team1_ships:
			_attach_agent(ship, _other_ships(team1_ships, ship), team0_ships)


# `roster` minus `ship`, preserving order — rosters are built by spawn_index
# already, so this stays spawn_index-sorted (see ShipObservations' slot-
# stability requirement).
func _other_ships(roster: Array[Ship], ship: Ship) -> Array[Ship]:
	var others: Array[Ship] = []
	for s in roster:
		if s != ship:
			others.append(s)
	return others


# Shared "--key=value" cmdline scan used by both eval and curriculum parsing.
func _cmdline_kv_args() -> Dictionary:
	var args := {}
	for argument in OS.get_cmdline_args():
		if argument.begins_with("--") and argument.find("=") > -1:
			var key_value := argument.lstrip("--").split("=", true, 1)
			args[key_value[0]] = key_value[1]
	return args


func _parse_eval_args() -> void:
	var args := _cmdline_kv_args()
	if args.has("eval_model_a") and args.has("eval_model_b"):
		_eval = true
		_eval_models[0] = args["eval_model_a"]
		_eval_models[1] = args["eval_model_b"]
		_eval_episodes = int(args.get("eval_episodes", str(_eval_episodes)))
		_eval_team_size = clampi(int(args.get("eval_team_size", str(_eval_team_size))), 1, 2)
		_eval_allow_vertical[0] = _typed_like(args.get("eval_allow_vertical_a", "true"), true)
		_eval_allow_vertical[1] = _typed_like(args.get("eval_allow_vertical_b", "true"), true)
		_eval_allow_pitch_roll[0] = _typed_like(args.get("eval_allow_pitch_roll_a", "true"), true)
		_eval_allow_pitch_roll[1] = _typed_like(args.get("eval_allow_pitch_roll_b", "true"), true)


# TrainingMode @export names a curriculum run may override from the cmdline.
# Explicit allow-list (not reflection) so a typo'd flag fails loudly instead
# of silently matching an unrelated inherited export.
const TRAINING_MODE_OVERRIDES := [
	"goal_reward", "draw_penalty", "kickoff_state_chance",
	"ball_near_goal_chance", "attack_goal_bias", "air_drill_chance",
	"air_intercept_chance", "ground_start_chance", "team_size",
	"wall_play_chance", "rebound_chance",
]
# ShipAIController @export names a curriculum run may override, read as
# --ai_<name>=<value> to avoid colliding with the names above.
const SHIP_AI_OVERRIDES := [
	"ball_touch_reward", "ball_touch_cooldown_ticks", "ball_touch_direction_floor", "team_touch_credit_weight",
	"velocity_to_ball_weight", "ball_velocity_to_goal_weight", "ball_distance_penalty",
	"forward_velocity_to_ball_weight", "air_approach_weight", "air_touch_bonus_weight", "wall_contact_penalty", "tilt_penalty",
	"ground_tilt_penalty", "non_forward_penalty", "grounded_upright_reward",
	"speed_reward_weight", "time_penalty", "airborne_penalty",
]


func _parse_curriculum_args() -> void:
	var args := _cmdline_kv_args()
	if args.has("opponent_mode"):
		_opponent_mode = args["opponent_mode"]
	_opponent_model_path = args.get("opponent_model", _opponent_model_path)
	if args.has("opponent_model_pool"):
		for path in String(args["opponent_model_pool"]).split(",", false):
			if not path.is_empty():
				_opponent_model_pool.append(path)
	if _opponent_mode == "league" and _opponent_model_pool.is_empty():
		push_error("TrainingMode: opponent_mode=league requires --opponent_model_pool=path,path")

	for name in TRAINING_MODE_OVERRIDES:
		if args.has(name):
			set(name, _typed_like(args[name], get(name)))
	var start_probability := kickoff_state_chance + ball_near_goal_chance \
		+ air_drill_chance + air_intercept_chance + ground_start_chance \
		+ wall_play_chance + rebound_chance
	if start_probability > 1.0:
		push_error("TrainingMode: episode-start probabilities sum to %.3f (> 1.0)" % start_probability)

	for name in SHIP_AI_OVERRIDES:
		var key := "ai_%s" % name
		if args.has(key):
			_ai_overrides[name] = _typed_like(args[key], _ai_default(name))


# Parses a cmdline string into the same Variant type as `sample` (bool/int/
# float pass through Godot's str()-based conversions; anything else stays a
# String), so callers can `set()` it straight onto a typed @export var.
func _typed_like(value: String, sample) -> Variant:
	match typeof(sample):
		TYPE_BOOL:
			return value.to_lower() in ["1", "true", "yes"]
		TYPE_INT:
			return value.to_int()
		TYPE_FLOAT:
			return value.to_float()
		_:
			return value


# ShipAIController isn't in the scene tree until _attach_agent instantiates
# one, so overrides need a default to type-match against up front; this
# mirrors ship_ai_controller.gd's own @export defaults.
func _ai_default(name: String) -> Variant:
	match name:
		"ball_touch_reward": return 0.4
		"ball_touch_cooldown_ticks": return 60
		"ball_touch_direction_floor": return 0.3
		"team_touch_credit_weight": return 0.0
		"velocity_to_ball_weight": return 0.02
		"forward_velocity_to_ball_weight": return 0.0
		"air_approach_weight": return 0.0
		"air_touch_bonus_weight": return 0.0
		"ball_velocity_to_goal_weight": return 0.004
		"ball_distance_penalty": return 0.002
		"wall_contact_penalty": return 0.0025
		"tilt_penalty": return 0.0005
		"ground_tilt_penalty": return 0.0
		"non_forward_penalty": return 0.0
		"grounded_upright_reward": return 0.0
		"speed_reward_weight": return 0.004
		"time_penalty": return 0.001
		"airborne_penalty": return 0.0
		_: return null


func _attach_agent(ship: Ship, teammates: Array[Ship], opponents: Array[Ship]) -> void:
	var agent := ShipAIController.new()
	agent.name = "ShipAIController"
	agent.reset_after = int(episode_length_seconds * TICKS_PER_SIM_SECOND)
	for key in _ai_overrides:
		agent.set(key, _ai_overrides[key])
	ship.add_child(agent)
	agent.setup(ship, ship.controller as RLShipController, ball, teammates, opponents, _attack_goal_position(ship.team))
	_agents.append(agent)


# The goal a team scores into: the one the opponent defends/concedes.
func _goal_for_team(team: int) -> Goal:
	for goal in arena.get_goals():
		if goal.team == 1 - team:
			return goal
	push_error("TrainingMode: no goal found for team %d to attack" % team)
	return null


func _attack_goal_position(team: int) -> Vector3:
	var goal := _goal_for_team(team)
	return goal.global_position if goal else Vector3.ZERO


func _physics_process(_delta):
	_respawn_escaped_bodies()
	if _eval:
		_episode_ticks += 1
		if _episode_ticks > int(episode_length_seconds * TICKS_PER_SIM_SECOND):
			_eval_draws += 1
			_end_eval_episode()
		return

	# Both trainer-requested resets and truncation (reset_after ticks elapsed)
	# surface as needs_reset. Only truncation is an episode end the trainer
	# must be told about via done — a trainer-requested reset already knows.
	var needs_reset := false
	var truncated := false
	for agent in _agents:
		needs_reset = needs_reset or agent.needs_reset
		truncated = truncated or agent.n_steps > agent.reset_after
	if needs_reset:
		if truncated:
			for agent in _agents:
				agent.reward -= draw_penalty
				agent.done = true
				agent.goal_scored_this_episode = false
				# Snapshot BEFORE _reset_episode() below, which moves the ship/
				# ball and would otherwise make this the post-reset state, not
				# the terminal one PPO needs to bootstrap V(s) from (see
				# ShipAIController.get_info / cosmic_env.py's truncation remap).
				# A goal (_on_goal_scored) does NOT do this — a goal is a
				# genuine terminal, V(s)=0 is correct there.
				agent.truncated_this_episode = true
				agent.terminal_obs = ShipObservations.build(agent.ship, agent.teammates, agent.opponents, agent.ball, agent.attack_goal_position)
		_reset_episode()
		return


func _on_goal_scored(conceding_team: int) -> void:
	if _eval:
		_eval_goals[1 - conceding_team] += 1
		_end_eval_episode()
		return
	for agent in _agents:
		agent.reward += goal_reward if agent.ship.team != conceding_team else -goal_reward
		agent.done = true
		agent.goal_scored_this_episode = true
		agent.truncated_this_episode = false  # genuine terminal, not a timeout
	_reset_episode()


func _end_eval_episode() -> void:
	_eval_episodes_done += 1
	_episode_ticks = 0
	if _eval_episodes_done >= _eval_episodes:
		print("EVAL_RESULT " + JSON.stringify({
			"model_a": _eval_models[0],
			"model_b": _eval_models[1],
			"episodes": _eval_episodes_done,
			"goals_a": _eval_goals[0],
			"goals_b": _eval_goals[1],
			"draws": _eval_draws,
		}))
		get_tree().quit()
		return
	# Randomized states (not kickoff): deterministic policies would otherwise
	# replay the identical episode every time.
	_reset_episode()


func _reset_episode() -> void:
	for agent in _agents:
		agent.reset()
	_select_league_opponent()

	var roll := randf()
	if roll < kickoff_state_chance:
		reset_ball()
		reset_ships()
	elif roll < kickoff_state_chance + ball_near_goal_chance:
		_place_ships_random()
		_place_ball_near_goal()
	elif roll < kickoff_state_chance + ball_near_goal_chance + air_drill_chance:
		_place_air_drill()
	elif roll < kickoff_state_chance + ball_near_goal_chance + air_drill_chance + air_intercept_chance:
		_place_air_intercept()
	elif roll < kickoff_state_chance + ball_near_goal_chance + air_drill_chance \
			+ air_intercept_chance + ground_start_chance:
		_place_ground_start()
	elif roll < kickoff_state_chance + ball_near_goal_chance + air_drill_chance \
			+ air_intercept_chance + ground_start_chance + wall_play_chance:
		_place_wall_state(false)
	elif roll < kickoff_state_chance + ball_near_goal_chance + air_drill_chance \
			+ air_intercept_chance + ground_start_chance + wall_play_chance + rebound_chance:
		_place_wall_state(true)
	else:
		_place_ships_random()
		_place_ball_random()


func _place_ball_random() -> void:
	var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_BALL_SPEED)
	_place_body(ball, Transform3D(Basis.IDENTITY, _random_position()), velocity, Vector3.ZERO)


# Extra clearance for the air drill's ball placement specifically — well
# beyond SPAWN_INSET, and well beyond the ball's own radius. The ball (unlike
# _random_position) has no collision-avoidance resample, so this is the
# anti-exploit measure: the RLGym wall-bounce exploit ("hits the ball off a
# wall high up instead of doing a real aerial") needs a wall to bounce off,
# so simply not generating ball states anywhere near one removes the exploit
# from the training distribution entirely, rather than trying to price it
# out via reward shaping.
const AIR_DRILL_BALL_WALL_CLEARANCE := 5.0


func _select_league_opponent() -> void:
	if _opponent_mode != "league" or _opponent_model_pool.is_empty():
		return
	var path := _opponent_model_pool[randi() % _opponent_model_pool.size()]
	for bot in _frozen_opponent_bots:
		bot.load_policy(path)


# Air drill state (see air_drill_chance): ball spawned high, both ships
# spawned low and lateral, so the state is unsolvable without climbing.
func _place_air_drill() -> void:
	var ball_half_x := ArenaBoundary.INNER_HALF_X - AIR_DRILL_BALL_WALL_CLEARANCE
	var ball_half_z := ArenaBoundary.GOAL_LINE_Z - AIR_DRILL_BALL_WALL_CLEARANCE
	var ball_position := Vector3(
		randf_range(-ball_half_x, ball_half_x),
		randf_range(ArenaBoundary.INNER_HEIGHT * 0.45, FIELD_MAX_Y),
		randf_range(-ball_half_z, ball_half_z)
	)
	var ball_velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_BALL_SPEED * 0.5)
	_place_body(ball, Transform3D(Basis.IDENTITY, ball_position), ball_velocity, Vector3.ZERO)

	# Each ship's lateral offset is sampled independently, so with more than
	# one ship per team (team_size > 1) two could otherwise land within their
	# own hulls of each other — resample against every ship already placed
	# this reset (see MIN_SHIP_SEPARATION).
	var placed: Array[Vector3] = []
	for ship in ships:
		if ship in _inert_ships:
			continue
		var ship_position := Vector3.ZERO
		for _attempt in 20:
			var lateral_offset := Vector3(randf_range(-1, 1), 0.0, randf_range(-1, 1))
			lateral_offset = lateral_offset.normalized() if lateral_offset.length_squared() > 0.001 else Vector3.FORWARD
			lateral_offset *= randf_range(6.0, 14.0)
			ship_position = Vector3(
				clampf(ball_position.x + lateral_offset.x, -FIELD_HALF_X, FIELD_HALF_X),
				randf_range(FIELD_MIN_Y, 4.0),
				clampf(ball_position.z + lateral_offset.z, -FIELD_HALF_Z, FIELD_HALF_Z)
			)
			if _far_enough_from(ship_position, placed):
				break
		placed.append(ship_position)
		var orientation := Basis.from_euler(Vector3(
			randf_range(-0.4, 0.4), randf_range(-PI, PI), randf_range(-0.4, 0.4)
		))
		_place_body(ship, Transform3D(orientation, ship_position), Vector3.ZERO, Vector3.ZERO)


# Ground start (see ground_start_chance): ships resting level on the floor,
# yaw-only so they begin belly-down rather than needing to recover attitude
# first, and a floor-level ball rolling slowly. This is the state the
# handling stage's rewards are actually written for — every ground term
# (ground_tilt_penalty, non_forward_penalty, the nose-led approach bonus)
# fades out by GROUND_HANDLING_HEIGHT, so they only bite in states like this
# one. The ball gets a modest planar-only velocity so it stays reachable
# without a climb.
func _place_ground_start() -> void:
	var ball_velocity := _random_direction()
	ball_velocity.y = 0.0
	ball_velocity = ball_velocity.normalized() * randf_range(0.0, MAX_RANDOM_BALL_SPEED * 0.5)
	var ball_position := Vector3(
		randf_range(-FIELD_HALF_X, FIELD_HALF_X),
		GROUND_START_BALL_Y,
		randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
	)
	_place_body(ball, Transform3D(Basis.IDENTITY, ball_position), ball_velocity, Vector3.ZERO)

	var placed: Array[Vector3] = []
	for ship in ships:
		if ship in _inert_ships:
			continue
		var ship_position := Vector3.ZERO
		for _attempt in 20:
			ship_position = Vector3(
				randf_range(-FIELD_HALF_X, FIELD_HALF_X),
				GROUND_START_Y,
				randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
			)
			if _spawn_position_clear(ship_position) and _far_enough_from(ship_position, placed):
				break
		placed.append(ship_position)
		var yaw := randf_range(-PI, PI)
		_place_body(
			ship, Transform3D(Basis.from_euler(Vector3(0.0, yaw, 0.0)), ship_position),
			Vector3.ZERO, Vector3.ZERO
		)


# Wall-play/rebound states (see wall_play_chance/rebound_chance). The ball is
# placed against a side wall, never in a corner or goal sensor. A wall-play
# state starts after the bounce and sends the ball inward; a rebound state
# starts before contact and sends it outward so the physics engine supplies
# the reflected trajectory. Ships use the ordinary randomized placement, so
# the policy has to read the wall/rebound context instead of memorising a
# fixed attacker spawn.
func _place_wall_state(rebound: bool) -> void:
	_place_ships_random()
	var side := -1.0 if randf() < 0.5 else 1.0
	var clearance := REBOUND_BALL_CLEARANCE if rebound else WALL_PLAY_BALL_CLEARANCE
	var ball_position := Vector3(
		side * (ArenaBoundary.INNER_HALF_X - clearance),
		randf_range(1.0, minf(FIELD_MAX_Y, 7.0)),
		randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
	)
	var velocity := wall_state_velocity(
		rebound, side, randf_range(WALL_PLAY_SPEED.x, WALL_PLAY_SPEED.y),
		randf_range(-0.15, 0.15), randf_range(-0.15, 0.15)
	)
	_place_body(ball, Transform3D(Basis.IDENTITY, ball_position), velocity, Vector3.ZERO)


# Pure geometry seam for adversarial tests. `side` identifies the selected
# wall (+1 or -1); a wall-play vector points into the field and a rebound
# vector points into that wall. Normalize the perturbed normal before applying
# speed so random tangential components cannot accidentally change the speed
# distribution between the two state types.
static func wall_state_velocity(rebound: bool, side: float, speed: float, vertical: float, lateral: float) -> Vector3:
	if speed < 0.0:
		return Vector3.ZERO
	var wall_side := -1.0 if side < 0.0 else 1.0
	var toward_field := Vector3(-wall_side, vertical, lateral).normalized()
	return (-toward_field if rebound else toward_field) * speed


# Air-intercept drill geometry. These six ranges are not free tuning knobs —
# together they decide whether the drill is solvable at all, and the original
# values made it arithmetically impossible (see the Round 9 note in
# training/generation5.py). The constraint: the ball is only above
# ShipAIController.AIR_TOUCH_HEIGHT (5m) for a fixed window after the spawn,
# and the ship has to cross the gap within it. The ship's own numbers cap what
# it can do — vertical_thrust 120 / mass 5 = 24 m/s^2 up, less 9.8 gravity, and
# drag_coefficient 0.98/tick caps climb at roughly 12 m/s — so the window has
# to be sized against those, not chosen for how the drill looks. The values
# below were picked by simulating the spawn distribution against that flight
# envelope: an ideal interceptor now reaches the ball in ~98% of episodes and
# can do so above 5m in ~37%, versus 0% before.
# Retuned 8-14m -> 6-10m on 2026-08-24, together with AIR_TOUCH_HEIGHT going
# 5.0 -> 3.0 (see ship_ai_controller.gd for the measurements behind that). The
# two are coupled and must move together: 8-14m was the RIGHT band for a 5m
# bar — simulating candidate bands against real physics, it maximised
# above-bar touches at 41.2% while 5-8m collapsed them to 4.3%, because a ball
# spawned near the bar drops under it almost immediately. Lowering the band
# alone would therefore have made the drill worse, not better. Against a 3m
# bar the ordering changes and 6-10m becomes the best row: 67.8% reach (was
# 53.2%) and 57.3% above-bar touches (was 41.2%), needing 5.2m of climb
# instead of 8.2m. It is also far closer to what the game actually produces —
# measured mean episode peak ball height in normal play is ~2.4m, so 8-14m was
# rehearsing a situation roughly 4x higher than anything a match generates.
const AIR_INTERCEPT_BALL_Y := Vector2(6.0, 10.0)      # matched to AIR_TOUCH_HEIGHT 3.0
const AIR_INTERCEPT_BALL_SPEED := Vector2(4.0, 8.0)   # slower: the ball outran the ship
const AIR_INTERCEPT_BEHIND := Vector2(4.0, 9.0)       # closer: less gap to close
const AIR_INTERCEPT_LATERAL := 5.0
const AIR_INTERCEPT_SHIP_Y := 4.0                     # upper bound; FIELD_MIN_Y is the lower
# A ship in real play is already moving; spawning at a dead stop spent most of
# the drill window just building speed, which was the single largest cause of
# the old geometry being unreachable. Planar only, aimed at the ball, so the
# climb itself is still the ship's own problem to solve.
const AIR_INTERCEPT_SHIP_SPEED := Vector2(6.0, 14.0)


# Goal-relevant aerial intercept: a high ball is already travelling toward a
# randomly selected goal, while ships begin low and behind/lateral to its
# path, already carrying planar speed toward it. The generous wall clearance
# prevents rebound farming and an upright yaw-only spawn avoids wasting the
# short drill window on random recovery.
func _place_air_intercept() -> void:
	var goal := _goal_for_team(randi() % 2)
	var ball_position := Vector3(
		randf_range(-8.0, 8.0),
		randf_range(AIR_INTERCEPT_BALL_Y.x, minf(AIR_INTERCEPT_BALL_Y.y, FIELD_MAX_Y)),
		randf_range(-10.0, 10.0)
	)
	var to_goal := (goal.global_position - ball_position).normalized()
	var ball_velocity := (to_goal + Vector3(randf_range(-0.15, 0.15), randf_range(0.0, 0.15), 0.0)).normalized() \
		* randf_range(AIR_INTERCEPT_BALL_SPEED.x, AIR_INTERCEPT_BALL_SPEED.y)
	_place_body(ball, Transform3D(Basis.IDENTITY, ball_position), ball_velocity, Vector3.ZERO)

	var placed: Array[Vector3] = []
	var behind := -Vector3(ball_velocity.x, 0.0, ball_velocity.z).normalized()
	for ship in ships:
		if ship in _inert_ships:
			continue
		var ship_position := Vector3.ZERO
		for _attempt in 20:
			var lateral := Vector3(-behind.z, 0.0, behind.x) \
				* randf_range(-AIR_INTERCEPT_LATERAL, AIR_INTERCEPT_LATERAL)
			ship_position = ball_position \
				+ behind * randf_range(AIR_INTERCEPT_BEHIND.x, AIR_INTERCEPT_BEHIND.y) + lateral
			ship_position.x = clampf(ship_position.x, -FIELD_HALF_X, FIELD_HALF_X)
			ship_position.y = randf_range(FIELD_MIN_Y, AIR_INTERCEPT_SHIP_Y)
			ship_position.z = clampf(ship_position.z, -FIELD_HALF_Z, FIELD_HALF_Z)
			if _spawn_position_clear(ship_position) and _far_enough_from(ship_position, placed):
				break
		placed.append(ship_position)
		var face_ball := ball_position - ship_position
		var yaw := atan2(-face_ball.x, -face_ball.z)
		var run_up := Vector3(face_ball.x, 0.0, face_ball.z)
		run_up = run_up.normalized() * randf_range(
			AIR_INTERCEPT_SHIP_SPEED.x, AIR_INTERCEPT_SHIP_SPEED.y
		) if run_up.length_squared() > 0.0001 else Vector3.ZERO
		_place_body(ship, Transform3D(Basis.from_euler(Vector3(0.0, yaw, 0.0)), ship_position), run_up, Vector3.ZERO)


# Attacking/defending drill states: ball close to a goal, moving toward it.
# Which goal is picked is biased by attack_goal_bias (0.5 = uniform between
# both, matching historical behaviour; 1.0 = always the goal team 0 attacks).
func _place_ball_near_goal() -> void:
	var goal := _goal_for_team(0) if randf() < attack_goal_bias else _goal_for_team(1)
	var toward_centre := -signf(goal.global_position.z)
	# Upper Y bound is a no-op on FLOOR arenas (goal.global_position.y ~0.79,
	# so maxf(4.0, ...) stays 4.0); on ELEVATED arenas it widens to sample
	# near the goal's real height instead of always landing near the floor.
	var position := Vector3(
		randf_range(-4.0, 4.0),
		randf_range(FIELD_MIN_Y, maxf(4.0, goal.global_position.y + 2.0)),
		goal.global_position.z + toward_centre * randf_range(3.0, 6.0)
	)
	var to_goal := (goal.global_position - position).normalized()
	var velocity := (to_goal + _random_direction() * 0.3).normalized() * randf_range(2.0, MAX_RANDOM_BALL_SPEED)
	_place_body(ball, Transform3D(Basis.IDENTITY, position), velocity, Vector3.ZERO)


func _place_ships_random() -> void:
	# Placed one at a time, resampling each against every position already
	# placed this reset (see MIN_SHIP_SEPARATION) — otherwise a team_size > 1
	# roster is sampled independently per ship and can spawn interpenetrating.
	var placed: Array[Vector3] = []
	for ship in ships:
		# Inert opponents (opponent_mode=inert) stay parked at their arena
		# spawn instead of drifting into the play area as a stray obstacle —
		# see _inert_ships.
		if ship in _inert_ships:
			continue
		var orientation := Basis.from_euler(Vector3(
			randf_range(-0.4, 0.4),
			randf_range(-PI, PI),
			randf_range(-0.4, 0.4)
		))
		var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_SHIP_SPEED)
		var position := _random_position(placed)
		placed.append(position)
		_place_body(ship, Transform3D(orientation, position), velocity, Vector3.ZERO)


func _random_position(exclude: Array[Vector3] = []) -> Vector3:
	# Resample anything too close to a corner curve or wall-base fillet (see
	# CORNER_LIMIT / FILLET_CLEAR_Y), or too close to an already-placed ship
	# this same reset (see MIN_SHIP_SEPARATION); the violating region is a
	# few percent of the volume, so 20 attempts effectively never fall
	# through even placing a full 5v5 roster one at a time.
	var position := Vector3.ZERO
	for _attempt in 20:
		position = Vector3(
			randf_range(-FIELD_HALF_X, FIELD_HALF_X),
			randf_range(FIELD_MIN_Y, FIELD_MAX_Y),
			randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
		)
		if _spawn_position_clear(position) and _far_enough_from(position, exclude):
			break
	return position


func _far_enough_from(position: Vector3, others: Array[Vector3]) -> bool:
	for other in others:
		if position.distance_squared_to(other) < MIN_SHIP_SEPARATION * MIN_SHIP_SEPARATION:
			return false
	return true


func _spawn_position_clear(position: Vector3) -> bool:
	if absf(position.x) + absf(position.z) > CORNER_LIMIT:
		return false
	if position.y >= FILLET_CLEAR_Y:
		return true
	return absf(position.x) <= FIELD_HALF_X - ArenaBoundary.BASE_RADIUS \
		and absf(position.z) <= FIELD_HALF_Z - ArenaBoundary.BASE_RADIUS


func _random_direction() -> Vector3:
	var direction := Vector3(randf_range(-1, 1), randf_range(-1, 1), randf_range(-1, 1))
	return direction.normalized() if direction.length_squared() > 0.001 else Vector3.FORWARD


func _place_body(body: RigidBody3D, to: Transform3D, linear_velocity: Vector3, angular_velocity: Vector3) -> void:
	# Deferred: a RigidBody3D transform can't be set mid-physics-step
	body.set_deferred("global_transform", to)
	body.set_deferred("linear_velocity", linear_velocity)
	body.set_deferred("angular_velocity", angular_velocity)