"""
Alpeccaai Pose — Hugging Face Space (ZeroGPU)
Runs DWPose / RTMPose whole-body 2D pose via rtmlib and returns keypoint JSON
in the exact format RIGFORGE's "AI keypoints -> Import pose keypoints" reads.

Endpoint: /detect  (the RIGFORGE "Detect via HF Space" button calls this).
"""
import os, json, tempfile
import numpy as np
import cv2
import gradio as gr

# ZeroGPU: CUDA is only visible INSIDE a @spaces.GPU function, so the model is
# built lazily on first call (see get_model). Off ZeroGPU this is a no-op.
try:
    import spaces
    gpu = spaces.GPU(duration=60)
except Exception:
    def gpu(fn):
        return fn

# COCO-17 body order == first 17 COCO-WholeBody points == RIGFORGE's order:
#  0 nose 1 L_eye 2 R_eye 3 L_ear 4 R_ear 5 L_sho 6 R_sho 7 L_elb 8 R_elb
#  9 L_wri 10 R_wri 11 L_hip 12 R_hip 13 L_knee 14 R_knee 15 L_ank 16 R_ank
SKELETON = [(5,7),(7,9),(6,8),(8,10),(5,6),(11,12),(5,11),(6,12),
            (11,13),(13,15),(12,14),(14,16),(0,5),(0,6)]

_model = None
def get_model():
    global _model
    if _model is None:
        import onnxruntime as ort
        from rtmlib import Wholebody
        dev = "cuda" if "CUDAExecutionProvider" in ort.get_available_providers() else "cpu"
        _model = Wholebody(mode="balanced", backend="onnxruntime", device=dev)
    return _model

def _openpose18(b):
    """COCO-17 -> OpenPose-COCO18 flat list (neck = shoulder midpoint)."""
    neck = [(b[5][0]+b[6][0])/2, (b[5][1]+b[6][1])/2, min(b[5][2], b[6][2])]
    order = [b[0], neck, b[6], b[8], b[10], b[5], b[7], b[9],
             b[12], b[14], b[16], b[11], b[13], b[15], b[2], b[1], b[4], b[3]]
    flat = []
    for x, y, s in order:
        flat += [float(x), float(y), float(s)]
    return flat

@gpu
def detect(image, conf=0.3):
    if image is None:
        raise gr.Error("Upload a single, isolated character figure (not a sheet).")
    rgb = np.array(image.convert("RGB"))
    bgr = rgb[:, :, ::-1].copy()
    h, w = bgr.shape[:2]

    kpts, scores = get_model()(bgr)
    if kpts is None or len(kpts) == 0:
        raise gr.Error("No person detected. Use a clean front T/A-pose crop.")

    p = int(np.argmax(scores.mean(axis=1)))
    kp, sc = kpts[p], scores[p]
    body = [[float(kp[i][0]), float(kp[i][1]), float(sc[i])] for i in range(17)]

    out = {
        "format": "coco17",
        "canvas_width": int(w), "canvas_height": int(h),
        "keypoints": body,
        "people": [{"pose_keypoints_2d": _openpose18(body)}],
        "model": "rtmlib.Wholebody(balanced)",
    }

    prev = rgb.copy()
    r = max(3, w // 200); lw = max(2, w // 300)
    for x, y, s in body:
        if s >= conf:
            cv2.circle(prev, (int(x), int(y)), r, (61, 240, 224), -1)
    for a, b in SKELETON:
        if body[a][2] >= conf and body[b][2] >= conf:
            cv2.line(prev, (int(body[a][0]), int(body[a][1])),
                     (int(body[b][0]), int(body[b][1])), (255, 61, 127), lw)

    path = os.path.join(tempfile.gettempdir(), "rigpose.json")
    with open(path, "w") as f:
        json.dump(out, f, indent=2)
    return prev, path

with gr.Blocks(title="Alpeccaai Pose") as demo:
    gr.Markdown(
        "# Alpeccaai Pose — DWPose keypoints\n"
        "Upload **one isolated character figure** (front T/A-pose is best — *not* a "
        "multi-pose reference sheet). Returns whole-body 2D keypoints as JSON for "
        "RIGFORGE → **AI keypoints → Import pose keypoints**, or call `/detect` from "
        "the in-tool **Detect via HF Space** button."
    )
    with gr.Row():
        inp = gr.Image(type="pil", label="Single character figure")
        with gr.Column():
            conf = gr.Slider(0.0, 1.0, value=0.3, label="Keypoint confidence threshold")
            btn = gr.Button("Detect", variant="primary")
    with gr.Row():
        out_img = gr.Image(label="Detected skeleton")
        out_file = gr.File(label="rigpose.json")
    btn.click(detect, inputs=[inp, conf], outputs=[out_img, out_file], api_name="detect")

if __name__ == "__main__":
    demo.launch()