class_name PolicyNetwork
extends RefCounted

# Minimal MLP forward pass for running trained policies in pure GDScript —
# no .NET build or ONNX runtime needed. Weights come from a JSON file written
# by training/export_policy.py (see TRAINING.md). The policy net is tiny
# (31 → 64 → 64 → 7 by default), and the bot only thinks every few physics
# ticks, so GDScript is plenty fast.
#
# JSON shape:
# {
#   "input_size": 31,
#   "layers": [
#     {"weights": [[out x in floats]], "biases": [out floats], "activation": "tanh" | "linear"},
#     ...
#   ]
#   "action_space": {"type": "multi_discrete", "heads": [{"name","bins"}, ...]}  // optional
# }
#
# "action_space" is absent from every model exported before curriculum
# generation 4 (e.g. Game/bots/promoted/easy.json) — absence means
# {"type": "continuous"}, decoded via ShipActionCodec.from_continuous, the
# same flattened-Box(7)-mean-output path this class has always produced.
# This class itself never changes behaviour based on it; only the caller
# (AIShipController._decide) branches on action_space["type"].

var input_size: int = 0
var action_space: Dictionary = {"type": "continuous"}
var _layers: Array = []


static func load_from_file(path: String) -> PolicyNetwork:
	if not FileAccess.file_exists(path):
		push_error("PolicyNetwork: model file not found: %s" % path)
		return null
	var text := FileAccess.get_file_as_string(path)
	var data: Variant = JSON.parse_string(text)
	if data == null or not (data is Dictionary) or not data.has("layers"):
		push_error("PolicyNetwork: invalid model file: %s" % path)
		return null

	var net := PolicyNetwork.new()
	net.input_size = int(data.get("input_size", 0))
	net.action_space = data.get("action_space", {"type": "continuous"})
	for layer in data["layers"]:
		# Flatten each layer's weights into a PackedFloat64Array for speed
		var out_size: int = layer["biases"].size()
		var in_size: int = layer["weights"][0].size()
		var flat := PackedFloat64Array()
		flat.resize(out_size * in_size)
		var i := 0
		for row in layer["weights"]:
			for value in row:
				flat[i] = value
				i += 1
		var biases := PackedFloat64Array(layer["biases"])
		net._layers.append({
			"weights": flat,
			"biases": biases,
			"in_size": in_size,
			"out_size": out_size,
			"tanh": layer.get("activation", "linear") == "tanh",
		})
	return net


func forward(observation: Array) -> Array:
	if observation.size() < input_size:
		push_error("PolicyNetwork: observation has %d values, model expects %d" % [observation.size(), input_size])
	# Slice rather than trust the caller: ShipObservations.SIZE only ever
	# grows (append-only), so an older/smaller model must still decode
	# correctly against a newer, longer observation vector — the extra
	# trailing values it never trained on are simply dropped here rather
	# than corrupting the first layer's dot product by accident.
	var x := PackedFloat64Array(observation.slice(0, input_size))
	for layer in _layers:
		var in_size: int = layer["in_size"]
		var out_size: int = layer["out_size"]
		var weights: PackedFloat64Array = layer["weights"]
		var biases: PackedFloat64Array = layer["biases"]
		var y := PackedFloat64Array()
		y.resize(out_size)
		for row in out_size:
			var sum := biases[row]
			var offset := row * in_size
			for col in in_size:
				sum += weights[offset + col] * x[col]
			y[row] = tanh(sum) if layer["tanh"] else sum
		x = y
	return Array(x)