"""策略输入特征（83 维）—— 与 webapp/js/ik-policy.js features() 逐项一致。改动须两边同步并升 FEAT_VER。"""
import math
import torch

FEAT_VER = 1
VMAX = [60., 60., 90., 90., 120., 180.]
DT = 0.02
STEP = [v * DT for v in VMAX]
ACC = 5.0
H, NJ, PREFIX = 16, 6, 3
GP = [(0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (2, 4), (2, 5), (3, 5)]
GAP_CAP = 40.0
F_DIM = 83
D2R = math.pi / 180.0
ROT6 = [0, 3, 6, 1, 4, 7]   # 行主序 3×3 的前两列


def wrap180(x):
    return torch.remainder(x + 180.0, 360.0) - 180.0


def features(kin, q, q_prev, T, R, prefix, prefix_len, sub, sub_flag, gm=None, xf=None):
    """全部为 batch 张量：
    q, q_prev [B,6] 度；T [B,3] mm；R [B,9]（pos 模式全 0）；pose 标志由 R 是否全 0 判定
    prefix [B,3,6] 增量（度），prefix_len [B]；sub [B,6]，sub_flag [B]（0/1）
    gm：kin.gaps(q, 40)['gm'] 可传入复用。"""
    B = q.shape[0]
    dt = q.dtype
    step = torch.tensor(STEP, dtype=dt, device=q.device)
    xf = xf if xf is not None else kin.cum_xf(q)
    out = []
    rad = q * D2R
    out.append(torch.stack([torch.sin(rad), torch.cos(rad)], -1).reshape(B, 12))
    lo, hi = kin.lo.to(dt), kin.hi.to(dt)
    idx = [1, 3, 4]
    out.append(2 * (q[:, idx] - lo[idx]) / (hi[idx] - lo[idx]) - 1)
    d = torch.where(kin.cont, wrap180(q - q_prev), q - q_prev)
    out.append(d / step)
    p = kin.fk(q, xf); Rq = kin.orient(q, xf).reshape(B, 9)
    out.append(p / 600)
    out.append(Rq[:, ROT6])
    out.append(T / 600)
    pose = (R.abs().sum(-1) > 0).to(dt)
    out.append(R[:, ROT6] * pose[:, None])
    out.append((T - p) / 600)
    out.append(((T - p) / 50).clamp(-1, 1))
    out.append(torch.stack([1 - pose, pose], -1))
    m = (torch.arange(PREFIX, device=q.device)[None, :] < prefix_len[:, None]).to(dt)   # [B,3]
    out.append((prefix / step * m[:, :, None]).reshape(B, 18))
    out.append(m)
    sd = torch.where(kin.cont, wrap180(sub - q), sub - q) / 180 * sub_flag[:, None].to(dt)
    out.append(sd)
    out.append(sub_flag[:, None].to(dt))
    if gm is None:
        gm = kin.gaps(q, GAP_CAP, xf)['gm']
    out.append(torch.stack([gm[:, a * 7 + b].clamp(max=GAP_CAP) / GAP_CAP for a, b in GP], -1))
    f = torch.cat(out, -1)
    assert f.shape[1] == F_DIM, f.shape
    return f