# -*- coding: utf-8 -*-
"""关节角组合可达空间（工具尖点）计算：FK 采样 + J1 旋转对称包络 + 体积估计"""
import json
import math
import os
import sys

import numpy as np

from arm import kinematics as HL
from arm.paths import DOCS, CACHE

TRAVEL = {t["axis"]: (t["lo"], t["hi"]) for t in
          json.load(open(DOCS + "/travel.json", encoding="utf-8"))}

TIP = (HL.X_ARM, HL.Y_AXIS, HL.Z_TIP)          # P09 法兰底面中心（工具安装面）
J1_AX = (HL.Y_AXIS, HL.Z_RL)                   # J1 转轴: 过 Y=128, Z=610, 平行 X


def rot_about(p, k, ang):
    """绕轴 (p 点, k 方向) 旋转 ang(弧度, 可为数组M,)；返回 Rm(M,3,3), t(M,3)。
    p' = Rm @ (p_pt - p) + p，与 AB.R 的 gp_Trsf.SetRotation 右手定则一致。"""
    k = np.asarray(k, dtype=float)
    k = k / np.linalg.norm(k)
    ang = np.atleast_1d(np.asarray(ang, dtype=float))
    px, py, pz = p
    K = np.array([[0.0, -k[2], k[1]], [k[2], 0.0, -k[0]], [-k[1], k[0], 0.0]])
    K2 = K @ K
    Rm = np.eye(3)[None, :, :] + np.sin(ang)[:, None, None] * K[None, :, :] \
        + (1.0 - np.cos(ang))[:, None, None] * K2[None, :, :]
    pp = np.array([px, py, pz])
    t = pp[None, :] - np.einsum("mij,j->mi", Rm, pp)
    return Rm, t


def fk_batch(angles, tip=TIP):
    """angles: dict idx->弧度数组(M,)；返回 (M,3) 世界坐标。链式 acc 与 sweep_v2 一致：
    父关节变换在最外层 p' = R1∘R2∘…∘R5 (p)，各 R_i 绕其零位世界轴（human_layout 唯一事实源）。"""
    pts = np.tile(np.asarray(tip, dtype=float), (angles[1].shape[0], 1))
    for i in range(5, 0, -1):                  # J6 跳过：尖点在 J6 轴上，其转角不改变尖点位置
        rear, d = HL.JOINTS[i - 1][2], HL.JOINTS[i - 1][3]
        Rm, t = rot_about(rear, d, angles[i])
        pts = np.einsum("mij,mj->mi", Rm, pts) + t
    return pts


def sweep(step=2.0):
    j2 = np.radians(np.arange(TRAVEL["J2"][0], TRAVEL["J2"][1] + 1e-9, step))
    j4 = np.radians(np.arange(TRAVEL["J4"][0], TRAVEL["J4"][1] + 1e-9, step))
    j5 = np.radians(np.arange(TRAVEL["J5"][0], TRAVEL["J5"][1] + 1e-9, step))
    g2, g4, g5 = np.meshgrid(j2, j4, j5, indexing="ij")
    M = g2.size
    angles = {1: np.zeros(M), 2: g2.ravel(), 3: np.zeros(M),
              4: g4.ravel(), 5: g5.ravel()}
    pts = fk_batch(angles)
    # 自证1：零位尖点必须回到 (18.5,128,93)
    z0 = fk_batch({i: np.zeros(1) for i in range(1, 6)})
    assert np.allclose(z0[0], TIP, atol=1e-9), z0[0]
    # 自证2：J1/J6 旋转不改变 (x, r) 与极值位置——J1 取中值扫一遍，r 分布应与 J1=0 相同
    for phi in (0.5, 1.2, 2.4):
        v1 = np.full(M, phi)
        p1 = fk_batch({1: v1, 2: g2.ravel(), 3: np.zeros(M),
                       4: g4.ravel(), 5: g5.ravel()})
        r1 = np.hypot(p1[:, 1] - J1_AX[0], p1[:, 2] - J1_AX[1])
        r0 = np.hypot(pts[:, 1] - J1_AX[0], pts[:, 2] - J1_AX[1])
        assert np.allclose(np.sort(r0), np.sort(r1), atol=1e-9), "J1 对称性被破坏"
    # 自证3：J3 旋转（大臂自转）不改变尖点位置——关节轴过尖点所在纵线
    print("J3 对称性: 尖点 x=%.1f 与 J3 轴 x=%.1f 一致=%s, y=%.1f 与轴 y=%.1f 一致=%s"
          % (TIP[0], HL.X_ARM, abs(TIP[0] - HL.X_ARM) < 1e-9,
             TIP[1], HL.Y_AXIS, abs(TIP[1] - HL.Y_AXIS) < 1e-9))
    print("J6 含于轴: 尖点 (%.1f,%.1f) vs 轴 (%.1f,%.1f)"
          % (TIP[0], TIP[1], HL.JOINTS[5][2][0], HL.JOINTS[5][2][1]))
    return pts


def envelope(pts, grid=1.0):
    """(x, r) 半平面栅格化占用 + 体积（绕 J1 轴回转）。"""
    x = pts[:, 0]
    r = np.hypot(pts[:, 1] - J1_AX[0], pts[:, 2] - J1_AX[1])
    x0, x1 = math.floor(x.min()), math.ceil(x.max())
    r0, r1 = 0.0, math.ceil(r.max())
    nx = int(round((x1 - x0) / grid)) + 1
    nr = int(round((r1 - r0) / grid)) + 1
    xi = np.clip(((x - x0) / grid).round().astype(int), 0, nx - 1)
    ri = np.clip((r / grid).round().astype(int), 0, nr - 1)
    occ = np.zeros((nx, nr), dtype=bool)
    occ[xi, ri] = True
    # 每个 x 列只保留 r <= 该列可达最大 r（回转体沿 r 连续填充）
    for i in range(nx):
        j = np.nonzero(occ[i])[0]
        if j.size:
            occ[i, : j.max() + 1] = True
    rr = r0 + np.arange(nr) * grid
    vol = float((occ * (2.0 * math.pi * rr)[None, :] * grid * grid).sum())
    return dict(occ=occ, x0=x0, grid=grid, rm_axes=rr, vol_mm3=vol,
                nx=nx, nr=nr)


def main():
    pts = sweep()
    x, y, z = pts[:, 0], pts[:, 1], pts[:, 2]
    r = np.hypot(y - J1_AX[0], z - J1_AX[1])
    env = envelope(pts)
    # 每个 x 的 r 包络（侧视/俯视剖面用）
    xs = env["x0"] + np.arange(env["nx"]) * env["grid"]
    rmax_x, rmin_x = [], []
    for i in range(env["nx"]):
        j = np.nonzero(env["occ"][i])[0]
        if j.size:
            rmax_x.append(env["rm_axes"][j.max()])
            rmin_x.append(env["rm_axes"][j.min()])
        else:
            rmax_x.append(np.nan)
            rmin_x.append(np.nan)
    os.makedirs(CACHE, exist_ok=True)
    np.savez_compressed(CACHE + "/ws_cloud.npz",
                        pts=pts, xs=np.array(xs),
                        rmax=np.array(rmax_x), rmin=np.array(rmin_x))
    met = dict(
        z_max=float(z.max()), z_min=float(z.min()),
        x_fwd=float(x.max()), x_back=float(x.min()),
        y_lo=float(y.min()), y_hi=float(y.max()),
        r_max=float(r.max()), r_min=float(r.min()),
        reach_from_j1_axis=float(r.max()),
        spread=float(z.max() - z.min()),
        vol_dm3=env["vol_mm3"] / 1e6,
        tip=TIP, j1_axis=list(J1_AX),
        limits={k: list(v) for k, v in TRAVEL.items()},
    )
    # 自证4：极值点姿态可解释（打印取极值的角度组合索引）
    imax = int(np.argmax(z)); imin = int(np.argmin(z))
    ifwd = int(np.argmax(x)); iback = int(np.argmin(x))
    met["pose_extreme"] = dict(
        zmax_idx=imax, zmin_idx=imin, xfwd_idx=ifwd, xback_idx=iback)
    print(json.dumps({k: v for k, v in met.items()
                      if k not in ("limits", "pose_extreme")},
                     ensure_ascii=False, indent=1))
    json.dump(met, open(CACHE + "/ws_metrics.json", "w", encoding="utf-8"),
              ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()