#!/usr/bin/env python3
"""
NEXT Athlete OS · RC38 pose atlas cleanup + runtime limb binding.

- Removes disconnected atlas bleed fragments from non-special sprite images.
- Crops cleaned sprites back to the main athlete/object component with safe padding.
- Adds full-sprite-percent runtimeLandmarks and limbBinding metadata to pose-atlas-packs.json.

This is intentionally deterministic and offline-safe: it gives the runtime stable feet/hands
identity now, while leaving room for true MediaPipe landmarks later under metrics.skeleton.mediaPipe.
"""
from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Dict, List, Tuple, Any

import cv2
import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
PACKS_PATH = ROOT / "assets/data/image-admin/pose-atlas-packs.json"
REPORT_PATH = ROOT / "assets/data/image-admin/rc38_pose_binding_cleanup_report.json"
PAD = 20
ALPHA_THRESHOLD = 16
MIN_FRAGMENT_AREA = 30

SPECIAL_KEEP_MULTICOMPONENT = ("sheet_12_bar_box_partner_special",)

Point = Dict[str, float]
Box = Tuple[int, int, int, int]  # inclusive x0,y0,x1,y1 for cv stats/internal


def pct(v: float, total: float) -> float:
    if total <= 0:
        return 50.0
    return round(max(0.0, min(100.0, (float(v) / float(total)) * 100.0)), 2)


def load_rgba(path: Path) -> Image.Image:
    return Image.open(path).convert("RGBA")


def alpha_mask(img: Image.Image) -> np.ndarray:
    arr = np.array(img)
    return (arr[:, :, 3] >= ALPHA_THRESHOLD).astype("uint8")


def components(mask: np.ndarray) -> List[Dict[str, Any]]:
    n, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, 8)
    out = []
    for i in range(1, n):
        x, y, w, h, area = stats[i]
        if int(area) <= 0:
            continue
        out.append({
            "label": int(i),
            "area": int(area),
            "box": (int(x), int(y), int(x + w - 1), int(y + h - 1)),
            "centroid": (float(centroids[i][0]), float(centroids[i][1])),
        })
    out.sort(key=lambda c: c["area"], reverse=True)
    return out


def bbox_from_mask(mask: np.ndarray) -> Box | None:
    ys, xs = np.where(mask > 0)
    if xs.size == 0 or ys.size == 0:
        return None
    return int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())


def crop_to_box(img: Image.Image, box: Box, pad: int = PAD) -> Image.Image:
    w, h = img.size
    x0, y0, x1, y1 = box
    cx0 = max(0, x0 - pad)
    cy0 = max(0, y0 - pad)
    cx1 = min(w - 1, x1 + pad)
    cy1 = min(h - 1, y1 + pad)
    return img.crop((cx0, cy0, cx1 + 1, cy1 + 1))


def should_component_clean(pose_key: str) -> bool:
    if any(s in pose_key for s in SPECIAL_KEEP_MULTICOMPONENT):
        return False
    # In RC37 the observed bleed is from packed atlas cells in footwork, burpee and
    # medball sheets. Other sheets are intentionally left untouched unless later
    # promoted here after visual QA.
    return pose_key.startswith((
        "sheet_06_footwork_shuffle",
        "sheet_09_floor_to_stand_burpee",
        "sheet_10_medball_throw_slam",
    ))


def remove_non_main_and_crop(path: Path, pose_key: str) -> Dict[str, Any] | None:
    if not should_component_clean(pose_key):
        return None
    img = load_rgba(path)
    mask = alpha_mask(img)
    n, labels_mask, stats, centroids = cv2.connectedComponentsWithStats(mask, 8)
    comps = []
    for i in range(1, n):
        x, y, ww, hh, area = stats[i]
        if int(area) > 0:
            comps.append({
                "label": int(i),
                "area": int(area),
                "box": (int(x), int(y), int(x + ww - 1), int(y + hh - 1)),
            })
    comps.sort(key=lambda c: c["area"], reverse=True)
    if len(comps) <= 1:
        return None
    secondary = [c for c in comps[1:] if c["area"] >= MIN_FRAGMENT_AREA]
    if not secondary:
        return None
    main_label = comps[0]["label"]
    keep = (labels_mask == main_label)
    arr = np.array(img)
    arr[:, :, 3] = np.where(keep, arr[:, :, 3], 0).astype(np.uint8)
    cleaned = Image.fromarray(arr, "RGBA")
    main_box = bbox_from_mask(keep.astype("uint8"))
    if main_box is None:
        return None
    cropped = crop_to_box(cleaned, main_box, PAD)
    before_size = img.size
    cropped.save(path)
    return {
        "poseKey": pose_key,
        "src": str(path.relative_to(ROOT)),
        "beforeSize": list(before_size),
        "afterSize": list(cropped.size),
        "keptArea": int(comps[0]["area"]),
        "removedFragments": [
            {"area": int(c["area"]), "box": [int(v) for v in c["box"]]} for c in secondary
        ],
    }


def point(x: float, y: float, w: int, h: int, confidence: float = 0.55, source: str = "alpha") -> Point:
    return {"x": pct(x, w), "y": pct(y, h), "confidence": round(float(confidence), 2), "source": source}


def visible_metrics(path: Path) -> Dict[str, Any]:
    img = load_rgba(path)
    w, h = img.size
    mask = alpha_mask(img)
    box = bbox_from_mask(mask)
    if box is None:
        # Safe fallback for an empty sprite.
        return {
            "spriteSize": [w, h],
            "visibleBox": [0, 0, w, h],
            "bottomContactX": 50.0,
            "bottomContactY": 100.0,
            "bottomContactMinX": 42.0,
            "bottomContactMaxX": 58.0,
            "lowerContactMinX": 42.0,
            "lowerContactMaxX": 58.0,
            "frontContactX": 42.0,
            "rearFootX": 58.0,
            "alphaCoverage": 0.0,
            "scaleHeightPx": float(h),
            "runtimeLandmarks": {},
            "limbBinding": {},
        }

    x0, y0, x1, y1 = box
    vb_w = max(1, x1 - x0 + 1)
    vb_h = max(1, y1 - y0 + 1)
    ys, xs = np.where(mask > 0)

    # Bottom-contact pixels: very bottom of the visible alpha, expanded a little so
    # light shoe shadows / antialiasing don't choose a single random toe pixel.
    bottom_y = int(ys.max())
    bottom_band_y = max(y0, bottom_y - max(3, int(vb_h * 0.035)))
    bottom_sel = mask.copy()
    bottom_sel[:bottom_band_y, :] = 0
    bys, bxs = np.where(bottom_sel > 0)
    if bxs.size == 0:
        bys, bxs = ys, xs
    bx_min, bx_max = int(bxs.min()), int(bxs.max())
    by_max = int(bys.max())
    bcx = float((bx_min + bx_max) / 2.0)

    # Lower support area: captures separated feet better than one-pixel bottom contact.
    lower_y = max(y0, int(y0 + vb_h * 0.62))
    lower_sel = mask.copy()
    lower_sel[:lower_y, :] = 0
    lys, lxs = np.where(lower_sel > 0)
    if lxs.size == 0:
        lys, lxs = bys, bxs
    lx_min, lx_max = int(lxs.min()), int(lxs.max())
    lower_cx = float((lx_min + lx_max) / 2.0)
    lower_cy = float(np.percentile(lys, 82)) if lys.size else float(by_max)

    # Upper/mid extremities for hand binding. These are deliberately "runtime binding"
    # hints, not anatomical truth. True anatomical labels can be added by MediaPipe later.
    mid_top = int(y0 + vb_h * 0.20)
    mid_bot = int(y0 + vb_h * 0.68)
    mid_sel = np.zeros_like(mask)
    mid_sel[mid_top:mid_bot + 1, :] = mask[mid_top:mid_bot + 1, :]
    mys, mxs = np.where(mid_sel > 0)
    if mxs.size:
        mh_left_x, mh_right_x = int(mxs.min()), int(mxs.max())
        # Median y near each horizontal extreme is more stable than using max/min y.
        left_band = mxs <= (mh_left_x + max(6, int(vb_w * 0.08)))
        right_band = mxs >= (mh_right_x - max(6, int(vb_w * 0.08)))
        left_y = float(np.median(mys[left_band])) if left_band.any() else float(y0 + vb_h * 0.42)
        right_y = float(np.median(mys[right_band])) if right_band.any() else float(y0 + vb_h * 0.42)
    else:
        mh_left_x = int(x0 + vb_w * 0.34)
        mh_right_x = int(x0 + vb_w * 0.66)
        left_y = right_y = float(y0 + vb_h * 0.42)

    head_top_y = y0
    head_center_y = y0 + vb_h * 0.12
    shoulder_y = y0 + vb_h * 0.30
    hip_y = y0 + vb_h * 0.60
    knee_y = y0 + vb_h * 0.79
    center_x = x0 + vb_w * 0.50

    runtime_landmarks = {
        "headTop": point(center_x, head_top_y, w, h, 0.48),
        "headCenter": point(center_x, head_center_y, w, h, 0.46),
        "shoulderCenter": point(center_x, shoulder_y, w, h, 0.44),
        "hipCenter": point(center_x, hip_y, w, h, 0.44),
        "kneeCenter": point(center_x, knee_y, w, h, 0.42),
        "leftHand": point(mh_left_x, left_y, w, h, 0.40),
        "rightHand": point(mh_right_x, right_y, w, h, 0.40),
        "handCenter": point((mh_left_x + mh_right_x) / 2.0, (left_y + right_y) / 2.0, w, h, 0.38),
        "leftFoot": point(lx_min, by_max, w, h, 0.62),
        "rightFoot": point(lx_max, by_max, w, h, 0.62),
        "footCenter": point(lower_cx, by_max, w, h, 0.66),
        "primaryContactFoot": point(bcx, by_max, w, h, 0.66),
        "rearFoot": point(lx_max, by_max, w, h, 0.55),
        "frontFoot": point(lx_min, by_max, w, h, 0.55),
        "supportCenter": point((lower_cx + (mh_left_x + mh_right_x) / 2.0) / 2.0, (lower_cy + (left_y + right_y) / 2.0) / 2.0, w, h, 0.38),
        "lowerCenter": point(lower_cx, lower_cy, w, h, 0.50),
    }

    return {
        "spriteSize": [w, h],
        "visibleBox": [int(x0), int(y0), int(x1 + 1), int(y1 + 1)],
        "bottomContactX": pct(bcx, w),
        "bottomContactY": pct(by_max, h),
        "bottomContactMinX": pct(bx_min, w),
        "bottomContactMaxX": pct(bx_max, w),
        "lowerContactMinX": pct(lx_min, w),
        "lowerContactMaxX": pct(lx_max, w),
        "frontContactX": pct(lx_min, w),
        "rearFootX": pct(lx_max, w),
        "alphaCoverage": round(float(mask.sum()) / float(max(1, w * h)), 4),
        "scaleHeightPx": float(vb_h),
        "runtimeLandmarks": runtime_landmarks,
        "limbBinding": {
            "schema": "next-limb-binding-v1",
            "coordinateSpace": "sprite-percent-unflipped",
            "leftRightPolicy": "visual-left-right; image and skeleton are mirrored together by runtime flip",
            "anchorPolicy": "prefer explicit anchorLandmark, otherwise contact-aware bound landmark",
            "feet": {
                "visualLeft": "leftFoot",
                "visualRight": "rightFoot",
                "center": "footCenter",
                "primaryContact": "primaryContactFoot",
                "rear": "rearFoot",
                "front": "frontFoot",
            },
            "hands": {
                "visualLeft": "leftHand",
                "visualRight": "rightHand",
                "center": "handCenter",
            },
            "support": {
                "center": "supportCenter",
                "lower": "lowerCenter",
            },
        },
    }


def pose_class_from_box(size: Tuple[int, int], visible_box: List[int]) -> str:
    w, h = size
    x0, y0, x1, y1 = visible_box
    vw, vh = max(1, x1 - x0), max(1, y1 - y0)
    aspect = vw / vh
    if aspect >= 1.15 and vh < 560:
        return "floor"
    if vh >= 500 and aspect < 1.05:
        return "upright"
    if vh >= 560 and aspect < 1.25:
        return "upright"
    return "compact"


def enrich_pose_entry(pose: Dict[str, Any], athlete_id: str, athlete_data: Dict[str, Any]) -> Dict[str, Any]:
    src = athlete_data.get("src")
    if not src:
        return {"poseKey": pose.get("poseKey"), "athleteId": athlete_id, "status": "missing-src"}
    path = ROOT / src
    if not path.exists():
        return {"poseKey": pose.get("poseKey"), "athleteId": athlete_id, "status": "missing-file", "src": src}

    cleanup = remove_non_main_and_crop(path, pose.get("poseKey", ""))
    vm = visible_metrics(path)

    metrics = athlete_data.setdefault("metrics", {})
    skeleton = metrics.setdefault("skeleton", {})

    # Keep legacy estimated landmarks untouched, but provide new runtimeLandmarks that are
    # in the same coordinate space as the displayed sprite. The JS runtime now prefers these.
    metrics["visibleBox"] = vm["visibleBox"]
    metrics["spriteSize"] = vm["spriteSize"]
    metrics["bottomContactX"] = vm["bottomContactX"]
    metrics["bottomContactY"] = vm["bottomContactY"]
    metrics["bottomContactMinX"] = vm["bottomContactMinX"]
    metrics["bottomContactMaxX"] = vm["bottomContactMaxX"]
    metrics["lowerContactMinX"] = vm["lowerContactMinX"]
    metrics["lowerContactMaxX"] = vm["lowerContactMaxX"]
    metrics["frontContactX"] = vm["frontContactX"]
    metrics["rearFootX"] = vm["rearFootX"]
    metrics["alphaCoverage"] = vm["alphaCoverage"]
    metrics["scaleHeightPx"] = vm["scaleHeightPx"]

    athlete_data["spriteSize"] = vm["spriteSize"]
    skeleton["schema"] = "rc38-bound-runtime-skeleton-v1"
    skeleton["quality"] = "estimated-alpha-bound"
    skeleton["coordinateSpace"] = "sprite-percent-unflipped"
    skeleton["runtimeLandmarks"] = vm["runtimeLandmarks"]
    skeleton["limbBinding"] = vm["limbBinding"]
    skeleton["poseClass"] = skeleton.get("poseClass") or pose_class_from_box(tuple(vm["spriteSize"]), vm["visibleBox"])
    skeleton["humanBox"] = vm["visibleBox"]
    skeleton["mainComponentBox"] = vm["visibleBox"]
    skeleton["scaleHeightPx"] = vm["scaleHeightPx"]
    skeleton["headToFootPx"] = vm["scaleHeightPx"]

    return {
        "poseKey": pose.get("poseKey"),
        "athleteId": athlete_id,
        "status": "cleaned+bound" if cleanup else "bound",
        "cleanup": cleanup,
        "spriteSize": vm["spriteSize"],
        "visibleBox": vm["visibleBox"],
        "anchors": {
            "leftFoot": vm["runtimeLandmarks"].get("leftFoot"),
            "rightFoot": vm["runtimeLandmarks"].get("rightFoot"),
            "footCenter": vm["runtimeLandmarks"].get("footCenter"),
            "leftHand": vm["runtimeLandmarks"].get("leftHand"),
            "rightHand": vm["runtimeLandmarks"].get("rightHand"),
        },
    }


def main() -> None:
    data = json.loads(PACKS_PATH.read_text(encoding="utf-8"))
    report_items = []
    cleaned = []
    bound_count = 0
    for pose in data.get("poses", []):
        for athlete_id, athlete_data in (pose.get("athletes") or {}).items():
            item = enrich_pose_entry(pose, athlete_id, athlete_data)
            report_items.append(item)
            if item.get("status") in ("bound", "cleaned+bound"):
                bound_count += 1
            if item.get("cleanup"):
                cleaned.append(item["cleanup"])

    data["version"] = "v97-rc38-pose-binding-cleanup"
    data["runtimeLandmarkSchema"] = "rc38-bound-runtime-skeleton-v1"
    PACKS_PATH.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    report = {
        "version": "rc38-pose-binding-cleanup",
        "intent": "Clean atlas bleed fragments and add stable sprite-linked runtime landmarks for feet/hands.",
        "boundSpriteCount": bound_count,
        "cleanedSpriteCount": len(cleaned),
        "cleanedSprites": cleaned,
        "notes": [
            "runtimeLandmarks are full-sprite percentage coordinates and are preferred by pose-motion-service.js.",
            "Legacy estimated landmarks are retained for traceability.",
            "Sheet 12 multi-component partner/bar/box special poses are intentionally not auto-cleaned.",
            "True MediaPipe anatomical landmarks can be layered into metrics.skeleton.mediaPipe later without changing runtime contracts."
        ],
        "items": report_items,
    }
    REPORT_PATH.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"RC38 pose binding cleanup complete: bound={bound_count}, cleaned={len(cleaned)}")
    for c in cleaned:
        print(f" cleaned {c['src']} {c['beforeSize']} -> {c['afterSize']} removed={len(c['removedFragments'])}")


if __name__ == "__main__":
    main()