class_name ShipObservations
extends RefCounted

# Canonical, team-relative observation builder. Shared by training
# (ShipAIController) and in-game inference (AIShipController) so a trained
# policy sees byte-identical inputs in both contexts — do not fork this logic.
#
# Self-play trick: observations for team 1 are rotated 180° about Y
# (x → -x, z → -z), so every ship perceives itself attacking toward -Z
# regardless of which side it spawned on. One policy can then play both teams.
# The same rotation must be inverted when interpreting actions (see canon —
# it is its own inverse).

# Normalization scales. Standard arena volume (see ArenaBoundary): x ±18,
# z ±27, height 18, goals at z ±27 (flush with the end walls); positions are
# soft-normalized to roughly [-1, 1]. Do not retune without retraining every
# model in Game/bots/.
const POSITION_SCALE := Vector3(30.0, 15.0, 30.0)
const BALL_SPEED_SCALE := 30.0
const GOAL_DISTANCE_SCALE := 60.0

# Contact normals with y above this are floor contact; below it they read as
# wall (sideways) or ceiling (downward) — mirrors
# ShipAIController.FLOOR_NORMAL_MIN_Y (kept here too since ShipAIController's
# wall_contact_penalty and this observation feature must agree on what
# counts as "in contact" for the reward/observation to stay consistent).
const FLOOR_NORMAL_MIN_Y := 0.7

# Fixed roster caps for the padded teammate/opponent slots below — the
# largest supported match size is 5v5. Slots beyond the real teammate/
# opponent count are zero-filled, mirroring the old single-opponent's
# null-zero-fill (see build()). Callers must pass teammates/opponents already
# sorted by Ship.spawn_index, so a given ship occupies the same slot in every
# tick's observation for the whole match, in both training and in-game
# inference (see AIShipController._discover_scene_refs /
# TrainingMode._start).
const MAX_TEAMMATES := 4
const MAX_OPPONENTS := 5

# Number of floats build() returns; the policy input size.
# 15 (own) + 6 (ball) + 6*MAX_TEAMMATES + 6*MAX_OPPONENTS + 4 (goal) + 4 (contact)
# Game/bots/promoted/*.json were exported against the old single-opponent,
# SIZE=35 layout and are not migrated — this is a from-scratch retrain, so
# those checkpoints are expected to go stale rather than keep decoding.
const SIZE := 15 + 6 + 6 * MAX_TEAMMATES + 6 * MAX_OPPONENTS + 4 + 4


# 180° rotation about Y for team 1; identity for team 0. A proper rotation
# (preserves handedness), and its own inverse — used for both observations
# and mapping canonical-frame actions back to world intent.
static func canon(v: Vector3, team: int) -> Vector3:
	return v if team == 0 else Vector3(-v.x, v.y, -v.z)


# attack_goal_position: centre of the goal this ship is trying to score in
# (the goal whose `team` == the opponent's team). teammates/opponents must
# already be sorted by Ship.spawn_index (ascending) by the caller — see
# MAX_TEAMMATES/MAX_OPPONENTS's comment for why slot stability matters; this
# function only pads/truncates to the fixed cap, it doesn't sort.
static func build(
	ship: Ship, teammates: Array[Ship], opponents: Array[Ship],
	ball: RigidBody3D, attack_goal_position: Vector3
) -> Array:
	var team := ship.team
	var obs := []

	# Own kinematics
	_append(obs, canon(ship.global_position, team) / POSITION_SCALE)
	_append(obs, canon(-ship.global_transform.basis.z, team))  # forward
	_append(obs, canon(ship.global_transform.basis.y, team))   # up
	_append(obs, canon(ship.linear_velocity, team) / ship.max_speed)
	_append(obs, canon(ship.angular_velocity, team) / ship.max_angular_speed)

	# Ball, relative to self
	var ball_rel := ball.global_position - ship.global_position
	_append(obs, canon(ball_rel, team) / POSITION_SCALE)
	_append(obs, canon(ball.linear_velocity, team) / BALL_SPEED_SCALE)

	# Teammates and opponents, relative to self, each padded/truncated to a
	# fixed slot count (zeros past the real roster size, e.g. a 1v1 match or
	# a solo drill) so the vector shape never depends on match size.
	_append_ship_slots(obs, ship, team, teammates, MAX_TEAMMATES)
	_append_ship_slots(obs, ship, team, opponents, MAX_OPPONENTS)

	# Goal we are attacking, relative to self
	var goal_rel := attack_goal_position - ship.global_position
	_append(obs, canon(goal_rel, team) / POSITION_SCALE)
	obs.append(goal_rel.length() / GOAL_DISTANCE_SCALE)

	# Own contact state. Added for generation 4: ShipAIController's
	# wall_contact_penalty used to fire on a condition the observation vector
	# couldn't see coming, leaving the value function to predict a reward
	# with no supporting signal. Also gives the policy a direct "am I resting
	# on a surface" signal it can use to push off (a real aerial mechanic),
	# distinct from inferring it indirectly from position/up-vector.
	var normal := contact_normal(ship)
	_append(obs, canon(normal, team))
	obs.append(1.0 if normal != Vector3.ZERO else 0.0)

	return obs


# Appends up to `slot_count` other ships' (relative position, relative
# velocity) — 6 floats each — zero-filling any slots beyond the real roster
# size, or beyond slot_count if the roster somehow has more (sorted-by-
# spawn_index order means truncation drops the highest spawn_index ships,
# not the nearest ones — acceptable since slot_count already covers the
# largest supported match size, 5v5).
static func _append_ship_slots(
	obs: Array, ship: Ship, team: int, others: Array[Ship], slot_count: int
) -> void:
	for i in slot_count:
		if i < others.size() and is_instance_valid(others[i]):
			var other := others[i]
			var rel := other.global_position - ship.global_position
			_append(obs, canon(rel, team) / POSITION_SCALE)
			_append(obs, canon(other.linear_velocity, team) / ship.max_speed)
		else:
			_append(obs, Vector3.ZERO)
			_append(obs, Vector3.ZERO)


static func _append(obs: Array, v: Vector3) -> void:
	obs.append(v.x)
	obs.append(v.y)
	obs.append(v.z)


# Aggregate wall/ceiling contact normal (zero if none, or if only floor
# contact — floor contact is excluded so "in_contact" means "touching
# something other than the ground it's expected to rest on", matching
# ShipAIController.wall_contact_penalty's own floor exemption). Requires
# ship.contact_monitor (see ship.gd's _ready — on unconditionally for every
# ship so training and in-game inference see identical observations).
static func contact_normal(ship: Ship) -> Vector3:
	var state := PhysicsServer3D.body_get_direct_state(ship.get_rid())
	if state == null:
		return Vector3.ZERO
	for i in state.get_contact_count():
		if not state.get_contact_collider_object(i) is ArenaBoundary:
			continue
		var normal := state.get_contact_local_normal(i)
		if normal.y < FLOOR_NORMAL_MIN_Y:
			return normal
	return Vector3.ZERO


# True belly-on-floor contact — the complement of contact_normal, which
# deliberately excludes floor contact (see its comment). Used by
# ShipAIController.grounded_upright_reward to reward genuinely resting on
# the floor rather than just being below the GROUND_HANDLING_HEIGHT proxy
# altitude, so a ship can't collect ground-handling reward by hovering just
# under the threshold without ever touching down.
static func is_floor_contact(ship: Ship) -> bool:
	var state := PhysicsServer3D.body_get_direct_state(ship.get_rid())
	if state == null:
		return false
	for i in state.get_contact_count():
		if not state.get_contact_collider_object(i) is ArenaBoundary:
			continue
		if state.get_contact_local_normal(i).y >= FLOOR_NORMAL_MIN_Y:
			return true
	return false