"""T9 SPIKE · 合成人形图生成器（纯几何绘制，不使用任何真实人脸/人体数据——合规红线）。

rtmpose_benchmark / birefnet_matting 的输入图、body_measure_prototype 的体型变体
全部由本模块确定性生成（同一参数永远产出同一张图），并把「生成时即已知」的真值
（三线高度 / 三线宽度 / 期望体型）写入 ground_truth.json——这就是体型分类一致性
验证的最小真值来源。
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np

ASSETS_DIR = Path(__file__).resolve().parent / "assets"

# ── 画布与几何（像素；画布 384x512，人形居中站立） ────────────────────────────
CANVAS_W, CANVAS_H = 384, 512
CX = CANVAS_W // 2
HEAD_R = 34
HEAD_TOP = 26
HEAD_CY = HEAD_TOP + HEAD_R
NECK_HW = 8
Y_SHOULDER, Y_WAIST, Y_HIP = 112, 200, 280
LEG_BOTTOM, LEG_GAP = 486, 8
ARM_T = 10  # 手臂沿躯干边缘绘制，每侧超出边缘 ARM_T/2 → 三线宽度 ≈ 2×半宽 + ARM_T

BODY_COLOR: tuple[int, int, int] = (150, 168, 198)  # BGR，米肤色调
BG_FLAT: tuple[int, int, int] = (242, 242, 242)  # 纯色背景
GRAD_TOP = np.array([250, 250, 250], dtype=np.float32)  # 渐变背景（上亮）
GRAD_BOTTOM = np.array([158, 166, 182], dtype=np.float32)  # 渐变背景（下暗）


@dataclass(frozen=True, slots=True)
class BodyShape:
    """体型变体：三处剖面半宽（像素）+ 期望分类（生成时即已知，即一致性断言真值）。"""

    key: str
    expect: str
    shoulder_hw: int
    waist_hw: int
    hip_hw: int


# STD 是 §3.2/§3.3 的标准输入图（中性比例，分类为 H）；其余五种覆盖 H/A/V/O/X 全类别
VARIANTS: tuple[BodyShape, ...] = (
    BodyShape("STD", "H", 60, 50, 60),
    BodyShape("H", "H", 58, 55, 57),
    BodyShape("A", "A", 44, 46, 68),
    BodyShape("V", "V", 72, 48, 54),
    BodyShape("X", "X", 60, 34, 62),
    BodyShape("O", "O", 52, 66, 56),
)


def _canvas(background: str) -> np.ndarray:
    if background == "gradient":
        t = np.linspace(0.0, 1.0, CANVAS_H, dtype=np.float32)[:, None]
        rows = GRAD_TOP[None, :] * (1.0 - t) + GRAD_BOTTOM[None, :] * t  # (H,3) 逐行插值
        img = np.repeat(rows[:, None, :], CANVAS_W, axis=1).astype(np.uint8)
        return np.ascontiguousarray(img)
    return np.full((CANVAS_H, CANVAS_W, 3), BG_FLAT, dtype=np.uint8)


def draw_figure(shape: BodyShape, background: str = "flat") -> np.ndarray:
    """按体型参数绘制人形剪影：头圆 + 颈 + 躯干三段梯形 + 贴边手臂 + 双腿。"""
    img = _canvas(background)
    left_edge = [
        (CX - shape.shoulder_hw, Y_SHOULDER),
        (CX - shape.waist_hw, Y_WAIST),
        (CX - shape.hip_hw, Y_HIP),
    ]
    right_edge = [
        (CX + shape.shoulder_hw, Y_SHOULDER),
        (CX + shape.waist_hw, Y_WAIST),
        (CX + shape.hip_hw, Y_HIP),
    ]
    poly = np.array([*left_edge, right_edge[2], right_edge[1], right_edge[0]], dtype=np.int32)
    cv2.fillPoly(img, [poly], BODY_COLOR)
    # 手臂：沿躯干左右轮廓的粗折线，圆角过渡，对三条剖面线近似均匀加宽 ARM_T/2×2
    for edge in (left_edge, right_edge):
        cv2.polylines(
            img,
            [np.array(edge, dtype=np.int32)],
            isClosed=False,
            color=BODY_COLOR,
            thickness=ARM_T,
            lineType=cv2.LINE_8,
        )
    # 颈 + 头
    cv2.rectangle(img, (CX - NECK_HW, HEAD_CY), (CX + NECK_HW, Y_SHOULDER + 6), BODY_COLOR, -1)
    cv2.circle(img, (CX, HEAD_CY), HEAD_R, BODY_COLOR, -1)
    # 双腿：臀线以下的两个矩形，总宽 = 臀宽，中间留 LEG_GAP
    leg_w = shape.hip_hw - LEG_GAP // 2
    for x0 in (CX - shape.hip_hw, CX + LEG_GAP // 2):
        cv2.rectangle(img, (x0, Y_HIP - 2), (x0 + leg_w, LEG_BOTTOM), BODY_COLOR, -1)
    return img


def _ground_truth(shape: BodyShape) -> dict[str, Any]:
    """生成时即已知的真值：三线高度、三线宽度（含手臂加宽）、包围盒、期望分类。"""
    half = max(shape.shoulder_hw, shape.hip_hw)
    return {
        "key": shape.key,
        "expect": shape.expect,
        "half_widths": {
            "shoulder": shape.shoulder_hw,
            "waist": shape.waist_hw,
            "hip": shape.hip_hw,
        },
        "gt_widths": {
            "shoulder": 2 * shape.shoulder_hw + ARM_T,
            "waist": 2 * shape.waist_hw + ARM_T,
            "hip": 2 * shape.hip_hw + ARM_T,
        },
        "line_y": {"shoulder": Y_SHOULDER, "waist": Y_WAIST, "hip": Y_HIP},
        "bbox": [CX - half - ARM_T, HEAD_TOP, CX + half + ARM_T, LEG_BOTTOM],
    }


def ensure_assets(assets_dir: Path = ASSETS_DIR) -> dict[str, Any]:
    """（重新）生成全部合成图与真值 JSON。确定性输出，可反复执行（幂等）。"""
    variants_dir = assets_dir / "variants"
    variants_dir.mkdir(parents=True, exist_ok=True)
    gt: dict[str, Any] = {"canvas": [CANVAS_W, CANVAS_H], "arm_thickness": ARM_T, "variants": {}}
    for shape in VARIANTS:
        cv2.imwrite(str(variants_dir / f"{shape.key}.png"), draw_figure(shape, "flat"))
        gt["variants"][shape.key] = _ground_truth(shape)
    # §3.2/§3.3 的标准输入：纯色背景 + 渐变背景变体（各一张）
    cv2.imwrite(str(assets_dir / "synthetic_person.png"), draw_figure(VARIANTS[0], "flat"))
    cv2.imwrite(str(assets_dir / "synthetic_person_grad.png"), draw_figure(VARIANTS[0], "gradient"))
    (assets_dir / "ground_truth.json").write_text(
        json.dumps(gt, indent=2, ensure_ascii=False), encoding="utf-8"
    )
    return gt