class_name TrainingMode
extends GameMode

# 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

# Placement bounds for randomized episode starts, derived from the standard
# enclosure (ArenaBoundary). The inset keeps a randomly oriented ship (1x1x4
# box, worst-case half-extent ~2.05) 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
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

# Sim runs at 60 physics ticks per sim-second regardless of speedup.
const TICKS_PER_SIM_SECOND := 60.0

# The arena is physically enclosed, so nothing should ever get this far out.
# If a body escapes anyway (physics regression, boundary edit), it is warned
# about and respawned with no reward change and no episode end — a multi-hour
# training run must survive it, and the escape must not shape rewards.
const ESCAPE_MARGIN := 15.0

var _agents: Array[ShipAIController] = []

# Eval mode state (see header comment)
var _eval := false
var _eval_models: Array[String] = ["", ""]
var _eval_episodes := 20
var _eval_goals := {0: 0, 1: 0}
var _eval_draws := 0
var _eval_episodes_done := 0
var _episode_ticks := 0

# 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 := ""
# 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] = []


func _start() -> void:
	_parse_eval_args()
	_parse_curriculum_args()
	spawn_ball()
	if _eval:
		for team in [0, 1]:
			var bot := AIShipController.new()
			bot.model_path = _eval_models[team]
			spawn_ship(team, 0, bot)
		return

	var ship_team0 := spawn_ship(0, 0, RLShipController.new())
	var ship_team1: Ship
	match _opponent_mode:
		"inert":
			ship_team1 = spawn_ship(1, 0, ShipController.new())
			_inert_ships.append(ship_team1)
		"frozen":
			var bot := AIShipController.new()
			bot.model_path = _opponent_model_path
			ship_team1 = spawn_ship(1, 0, bot)
		_:
			ship_team1 = spawn_ship(1, 0, RLShipController.new())

	_attach_agent(ship_team0, ship_team1)
	if _opponent_mode == "self_play":
		_attach_agent(ship_team1, ship_team0)


# 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)))


# 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",
]
# 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",
	"velocity_to_ball_weight", "ball_velocity_to_goal_weight", "ball_distance_penalty",
	"wall_contact_penalty", "tilt_penalty", "speed_reward_weight", "time_penalty",
	"allow_vertical", "allow_pitch_roll",
]


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)

	for name in TRAINING_MODE_OVERRIDES:
		if args.has(name):
			set(name, _typed_like(args[name], get(name)))

	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
		"velocity_to_ball_weight": return 0.02
		"ball_velocity_to_goal_weight": return 0.004
		"ball_distance_penalty": return 0.002
		"wall_contact_penalty": return 0.0025
		"tilt_penalty": return 0.002
		"speed_reward_weight": return 0.004
		"time_penalty": return 0.001
		"allow_vertical", "allow_pitch_roll": return true
		_: return null


func _attach_agent(ship: Ship, opponent: 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, opponent, _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
		_reset_episode()
		return


# Escape failsafe: see ESCAPE_MARGIN.
func _respawn_escaped_bodies() -> void:
	for ship in ships:
		if is_instance_valid(ship) and _escaped(ship.global_position):
			push_warning("TrainingMode: ship escaped the enclosed arena — check boundary colliders")
			_place_body(ship, _ship_spawn_transforms[ship], Vector3.ZERO, Vector3.ZERO)
	if is_instance_valid(ball) and _escaped(ball.global_position):
		push_warning("TrainingMode: ball escaped the enclosed arena — check boundary colliders")
		_place_body(ball, arena.get_ball_spawn(), Vector3.ZERO, Vector3.ZERO)


func _escaped(position: Vector3) -> bool:
	return absf(position.x) > ArenaBoundary.INNER_HALF_X + ESCAPE_MARGIN \
		or absf(position.z) > ArenaBoundary.INNER_HALF_Z + ESCAPE_MARGIN \
		or position.y < -ESCAPE_MARGIN \
		or position.y > ArenaBoundary.INNER_HEIGHT + ESCAPE_MARGIN


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
	_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()

	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()
	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)


# 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)
	var position := Vector3(
		randf_range(-4.0, 4.0),
		randf_range(FIELD_MIN_Y, 4.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:
	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)
		_place_body(ship, Transform3D(orientation, _random_position()), velocity, Vector3.ZERO)


func _random_position() -> Vector3:
	# Resample anything too close to a corner curve or wall-base fillet (see
	# CORNER_LIMIT / FILLET_CLEAR_Y); the violating region is a few percent
	# of the volume, so 20 attempts effectively never fall through.
	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):
			break
	return position


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)