class_name ShipActionCodec
extends RefCounted

# Single source of truth for the RL action layout — shared by training
# (ShipAIController.get_action_space/set_action) and in-game inference
# (AIShipController._decide via PolicyNetwork) so a trained policy's action
# output is decoded identically in both contexts. Mirrors ShipObservations'
# "do not fork this logic" role for observations; the train/inference seam
# broke once before over exactly this kind of divergence (commit 8c15c46).
#
# Ship thrust is body-local, so its axes must not be mirrored for team 1.
# Ship rotation, however, is applied directly as world-space torque in
# Ship.apply_rotation_forces(). Team 1 observes a canonical frame rotated
# 180 degrees about world Y, so its canonical pitch/roll outputs must be
# rotated back to world space before they reach the ship. apply_team_frame()
# is the shared training/inference seam for that conversion.
#
# Curriculum generation 4 replaces the old continuous Gaussian action space
# (Box(7), see the "continuous" path below) with a per-axis MultiDiscrete
# space: PPO's Gaussian std reliably collapsed to ~0.13-0.15 within the first
# ~10% of every training run across 3 generations and never recovered, which
# made a *sustained* set-point (e.g. hovering, thrust.y ~= 0.408 given this
# ship's mass/thrust — see TRAINING.md) essentially unreachable: the
# collapsed distribution can brush the hover value but never hold it long
# enough to accumulate the reward signal that would move the mean. A
# discrete bin is a single, atomic, repeatable choice with non-zero
# probability under any softmax, which does not have that failure mode.
#
# HEADS order is deliberately gymnasium's *sorted* key order (verified:
# "rot_x" < "rot_y" < "rot_z" < "thrust_x" < "thrust_y" < "thrust_z" <
# "turbo") — godot_rl's ActionSpaceProcessor builds the Tuple action space
# from a gymnasium Dict, which sorts keys regardless of insertion order, so
# this order is what SB3/PPO actually samples/trains against and what
# set_action() receives keyed by. Do not reorder without re-verifying that
# sort order.
const HEADS := [
	{"name": "rot_x", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
	{"name": "rot_y", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
	{"name": "rot_z", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
	{"name": "thrust_x", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
	# Deliberately asymmetric: hovering this ship (mass 5.0, vertical_thrust
	# 120, default gravity 9.8 m/s^2 — see ship.gd/ship.tscn) requires a
	# sustained thrust.y ~= 0.408. Uniform-random selection over these 5 bins
	# averages 0.34 — just below neutral buoyancy, so a fresh policy drifts
	# gently through the volume instead of pinning to the floor (symmetric
	# bins) or sticking to the ceiling (ceiling_pull_strength 11.5 > gravity
	# 9.8, so the ceiling is easy to over-shoot into). This is the direct
	# analogue of the RLGym/RLBot community fix for the same failure mode
	# ("add more jump actions to the discrete action parser").
	{"name": "thrust_y", "bins": [-0.5, 0.0, 0.45, 0.75, 1.0]},
	{"name": "thrust_z", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
	{"name": "turbo", "bins": [0.0, 1.0]},
]


static func action_space_dict() -> Dictionary:
	var space := {}
	for head in HEADS:
		space[head["name"]] = {"size": head["bins"].size(), "action_type": "discrete"}
	return space


# Training side: `action` is the Dictionary godot_rl's Sync node hands
# set_action() — one entry per HEADS key, each an int (or int-valued float)
# bin index in [0, bins.size()).
static func from_indices(action: Dictionary) -> ShipAction:
	var result := ShipAction.new()
	var values := {}
	for head in HEADS:
		var index: int = clampi(int(round(float(action[head["name"]]))), 0, head["bins"].size() - 1)
		values[head["name"]] = head["bins"][index]
	result.rotation = Vector3(values["rot_x"], values["rot_y"], values["rot_z"])
	result.thrust = Vector3(values["thrust_x"], values["thrust_y"], values["thrust_z"])
	result.turbo = values["turbo"] > 0.0
	return result


# In-game inference for a MultiDiscrete-trained export: `logits` is the raw
# policy_network.gd output — 32 floats (5+5+5+5+5+5+2), one contiguous slice
# per head in HEADS order (matches export_policy.py's action_net layer,
# which concatenates SB3's per-head categorical logits in that same order).
# argmax within each slice picks that head's bin, same as SB3's
# MultiCategoricalDistribution.mode() under deterministic inference.
static func from_logits(logits: Array, noise: float) -> ShipAction:
	var result := ShipAction.new()
	var values := {}
	var offset := 0
	for head in HEADS:
		var bins: Array = head["bins"]
		var index := 0
		if noise > 0.0 and randf() < noise:
			# eps-random-bin: the discrete analogue of continuous action_noise
			# (see ai_ship_controller.gd) — degrades gracefully and keeps the
			# same 0..1 monotonic difficulty semantics as the continuous path.
			index = randi() % bins.size()
		else:
			var best_value: float = logits[offset]
			for i in range(1, bins.size()):
				if logits[offset + i] > best_value:
					best_value = logits[offset + i]
					index = i
		values[head["name"]] = bins[index]
		offset += bins.size()
	result.rotation = Vector3(values["rot_x"], values["rot_y"], values["rot_z"])
	result.thrust = Vector3(values["thrust_x"], values["thrust_y"], values["thrust_z"])
	result.turbo = values["turbo"] > 0.0
	return result


# Map a policy's canonical-frame rotation intent back into the physical
# team's world frame. A 180-degree Y rotation negates X and Z and leaves Y
# unchanged. Translation remains untouched because Ship applies it through
# the ship's local basis rather than as a world-space vector.
static func apply_team_frame(action: ShipAction, team: int) -> ShipAction:
	if team == 1:
		action.rotation.x = -action.rotation.x
		action.rotation.z = -action.rotation.z
	return action


# Legacy continuous decode — moved verbatim from ai_ship_controller.gd so
# every model exported before generation 4 (no "action_space" block in its
# JSON, e.g. Game/bots/promoted/easy.json) keeps behaving byte-identically.
# `out` is the trainer's flattened Box(7) output, gymnasium-sorted: rotation
# xyz, thrust xyz, turbo (> 0 means on) — NOT ShipAction's thrust-first
# declaration order.
static func from_continuous(out: Array, noise: float) -> ShipAction:
	var result := ShipAction.new()
	result.rotation = Vector3(
		_continuous_axis(out[0], noise),
		_continuous_axis(out[1], noise),
		_continuous_axis(out[2], noise)
	)
	result.thrust = Vector3(
		_continuous_axis(out[3], noise),
		_continuous_axis(out[4], noise),
		_continuous_axis(out[5], noise)
	)
	result.turbo = out[6] > 0.0
	return result


static func _continuous_axis(value: float, noise: float) -> float:
	if noise > 0.0:
		value += randf_range(-noise, noise)
	return clampf(value, -1.0, 1.0)