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 ±12,
# z ±18, height 12, goals at z ±18 (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(20.0, 10.0, 20.0)
const BALL_SPEED_SCALE := 30.0
const GOAL_DISTANCE_SCALE := 40.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

# Number of floats build() returns; the policy input size. APPEND-ONLY: new
# features go on the end and existing indices never move, so an old exported
# model (whose network was trained against a shorter SIZE) still decodes its
# first N inputs identically when SIZE grows — see PolicyNetwork.forward's
# input_size slice/guard. Do not retune an *existing* index without
# retraining every model in Game/bots/.
const SIZE := 35


# 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).
static func build(ship: Ship, opponent: 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)

	# Opponent, relative to self (zeros if absent, e.g. a 1-ship drill)
	if is_instance_valid(opponent):
		var opp_rel := opponent.global_position - ship.global_position
		_append(obs, canon(opp_rel, team) / POSITION_SCALE)
		_append(obs, canon(opponent.linear_velocity, team) / ship.max_speed)
	else:
		_append(obs, Vector3.ZERO)
		_append(obs, Vector3.ZERO)

	# 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 (appended — see SIZE's append-only invariant).
	# 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


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