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"},
#     ...
#   ]
# }

var input_size: int = 0
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))
	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:
	var x := PackedFloat64Array(observation)
	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)