class_name PoseFeatures
extends RefCounted

## 关键点 → DTW 特征向量 —— 契约 contracts/skeleton-topology.md §5（冻结）
##
## 每帧特征 = [33 × (x, y, z), 9 个关节角余弦] = 108 维。
## 低置信度点（conf < 0.5）对应的维度权重置 0，不参与代价计算。
##
## 关节角三元组顺序冻结为：
##   [shoulder_L, shoulder_R, elbow_L, elbow_R, hip_L, hip_R, knee_L, knee_R, spine]

const KeypointFrame = preload("res://src/net/keypoint_frame.gd")

const KEYPOINT_COUNT := 33
const VIRTUAL_NECK := 33
const VIRTUAL_PELVIS := 34
const ANGLE_COUNT := 9
const FEATURE_DIM := KEYPOINT_COUNT * 3 + ANGLE_COUNT  # 108

## 关节角 (a, b, c) 三元组索引
const ANGLE_TRIPLETS := [
	[KeypointFrame.LEFT_ELBOW, KeypointFrame.LEFT_SHOULDER, KeypointFrame.LEFT_HIP],
	[KeypointFrame.RIGHT_ELBOW, KeypointFrame.RIGHT_SHOULDER, KeypointFrame.RIGHT_HIP],
	[KeypointFrame.LEFT_SHOULDER, KeypointFrame.LEFT_ELBOW, KeypointFrame.LEFT_WRIST],
	[KeypointFrame.RIGHT_SHOULDER, KeypointFrame.RIGHT_ELBOW, KeypointFrame.RIGHT_WRIST],
	[KeypointFrame.LEFT_SHOULDER, KeypointFrame.LEFT_HIP, KeypointFrame.LEFT_KNEE],
	[KeypointFrame.RIGHT_SHOULDER, KeypointFrame.RIGHT_HIP, KeypointFrame.RIGHT_KNEE],
	[KeypointFrame.LEFT_HIP, KeypointFrame.LEFT_KNEE, KeypointFrame.LEFT_ANKLE],
	[KeypointFrame.RIGHT_HIP, KeypointFrame.RIGHT_KNEE, KeypointFrame.RIGHT_ANKLE],
	[KeypointFrame.NOSE, VIRTUAL_NECK, VIRTUAL_PELVIS],
]

const ANGLE_NAMES := [
	"shoulder_L", "shoulder_R", "elbow_L", "elbow_R",
	"hip_L", "hip_R", "knee_L", "knee_R", "spine",
]


## 一帧特征：108 维值 + 108 维权重
class FeatureVector extends RefCounted:
	var values: PackedFloat32Array = PackedFloat32Array()
	var weights: PackedFloat32Array = PackedFloat32Array()
	var time_ms: float = 0.0

	func _init() -> void:
		values.resize(PoseFeatures.FEATURE_DIM)
		weights.resize(PoseFeatures.FEATURE_DIM)
		values.fill(0.0)
		weights.fill(1.0)

	func dimension() -> int:
		return values.size()


## 由姿态帧构造特征向量
static func build(frame: KeypointFrame.PoseFrame, use_unmirrored: bool = true) -> FeatureVector:
	var out := FeatureVector.new()
	if frame == null:
		return out

	var points: PackedVector3Array = frame.unmirrored_points() if use_unmirrored else frame.points
	var confidence: PackedFloat32Array = frame.confidence
	out.time_ms = float(frame.received_at_ms)

	# ---- 33 × 3 坐标 ----
	for i in KEYPOINT_COUNT:
		var p := points[i] if i < points.size() else Vector3.ZERO
		var c := confidence[i] if i < confidence.size() else 0.0
		var w := 1.0 if c >= KeypointFrame.CONFIDENCE_THRESHOLD else 0.0
		out.values[i * 3] = p.x
		out.values[i * 3 + 1] = p.y
		out.values[i * 3 + 2] = p.z
		out.weights[i * 3] = w
		out.weights[i * 3 + 1] = w
		out.weights[i * 3 + 2] = w

	# ---- 9 个关节角 ----
	var neck := _midpoint(points, KeypointFrame.LEFT_SHOULDER, KeypointFrame.RIGHT_SHOULDER)
	var pelvis := _midpoint(points, KeypointFrame.LEFT_HIP, KeypointFrame.RIGHT_HIP)
	var positions := _extended_positions(points, neck, pelvis)

	for a in ANGLE_COUNT:
		var triple: Array = ANGLE_TRIPLETS[a]
		var idx := KEYPOINT_COUNT * 3 + a
		out.values[idx] = _angle_cos(positions[triple[0]], positions[triple[1]], positions[triple[2]])
		# 三个端点均可信时该角度才参与代价计算
		var trustworthy := (
			_weight_for_point(confidence, triple[0]) > 0.0
			and _weight_for_point(confidence, triple[1]) > 0.0
			and _weight_for_point(confidence, triple[2]) > 0.0
		)
		out.weights[idx] = 1.0 if trustworthy else 0.0

	return out


## 33 点 + neck(33) + pelvis(34)
static func _extended_positions(points: PackedVector3Array, neck: Vector3, pelvis: Vector3) -> Array:
	var out: Array = []
	out.resize(VIRTUAL_PELVIS + 1)
	for i in KEYPOINT_COUNT:
		out[i] = points[i] if i < points.size() else Vector3.ZERO
	out[VIRTUAL_NECK] = neck
	out[VIRTUAL_PELVIS] = pelvis
	return out


static func _midpoint(points: PackedVector3Array, a: int, b: int) -> Vector3:
	var pa: Vector3 = points[a] if a < points.size() else Vector3.ZERO
	var pb: Vector3 = points[b] if b < points.size() else Vector3.ZERO
	return (pa + pb) * 0.5


static func _weight_for_point(confidence: PackedFloat32Array, index: int) -> float:
	if index == VIRTUAL_NECK or index == VIRTUAL_PELVIS or index >= KeypointFrame.KEYPOINT_COUNT:
		return 1.0
	if index >= confidence.size():
		return 0.0
	return 1.0 if confidence[index] >= KeypointFrame.CONFIDENCE_THRESHOLD else 0.0


## 向量 ba 与 bc 夹角的余弦，落在 [-1, 1]
static func _angle_cos(a: Vector3, b: Vector3, c: Vector3) -> float:
	var ba := a - b
	var bc := c - b
	var denom := ba.length() * bc.length()
	if denom < 1e-9:
		return 0.0
	return clampf(ba.dot(bc) / denom, -1.0, 1.0)


## 加权欧氏距离（逐维归一化，使不同维数可比）
static func weighted_distance(a: FeatureVector, b: FeatureVector) -> float:
	var dim := mini(a.values.size(), b.values.size())
	if dim == 0:
		return 0.0
	var sum := 0.0
	var weight_total := 0.0
	for i in dim:
		var w: float = minf(a.weights[i], b.weights[i])
		if w <= 0.0:
			continue
		var d := a.values[i] - b.values[i]
		sum += w * d * d
		weight_total += w
	if weight_total <= 0.0:
		return 0.0
	return sqrt(sum / weight_total)


## 逐维绝对差（用于调试与可视化）
static func dimension_deltas(a: FeatureVector, b: FeatureVector) -> PackedFloat32Array:
	var out := PackedFloat32Array()
	out.resize(mini(a.values.size(), b.values.size()))
	for i in out.size():
		out[i] = absf(a.values[i] - b.values[i])
	return out


## 关节角名称索引 → 中文标签（UI 展示用）
static func angle_label(index: int) -> String:
	match index:
		0: return "左肩"
		1: return "右肩"
		2: return "左肘"
		3: return "右肘"
		4: return "左髋"
		5: return "右髋"
		6: return "左膝"
		7: return "右膝"
		8: return "躯干"
	return "角度%d" % index