"""
Main-figure pose extractor — gold 4x1000 subset.

Tests upgrade (a) from the pose investigation: instead of averaging all
skeletons in an image (extract_pose.py), store ONLY the largest detected
figure's torso-normalised skeleton, plus quality indicators so garbage
detections can be filtered in analysis:

  filename
  n_persons        — YOLO detections in the image
  main_conf        — box confidence of the largest figure
  main_kpt_conf    — mean keypoint confidence of the largest figure
  main_torso_px    — torso height in pixels (tiny => unreliable normalisation)
  main_area        — bbox area fraction of the image
  main_skel_0..33  — 17 kpts x (dx,dy), torso-normalised, 0 where invisible

Output: data/features/pose.parquet   (checkpointed, resumable)
Usage:  python features/extract_pose_main_gold.py [--model yolov8m-pose.pt]
"""

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from tqdm import tqdm

ROOT = Path(__file__).resolve().parent.parent
IMAGES = ROOT / "data/images"
SELECTED = ROOT / "data/artwork_metadata.csv"
OUTPUT = ROOT / "data/features/pose.parquet"

KPT_THR = 0.3
L_SHOULDER, R_SHOULDER, L_HIP, R_HIP = 5, 6, 11, 12


def normalise_skeleton(kpts, conf):
    """Torso-normalised skeleton (matches extract_pose.py) + torso height px."""
    vis = conf >= KPT_THR
    anchors = [L_SHOULDER, R_SHOULDER, L_HIP, R_HIP]
    if sum(vis[i] for i in anchors) >= 2:
        mid_shoulder = (kpts[L_SHOULDER] + kpts[R_SHOULDER]) / 2.0
        mid_hip = (kpts[L_HIP] + kpts[R_HIP]) / 2.0
        centre = (mid_shoulder + mid_hip) / 2.0
        height = float(np.linalg.norm(mid_shoulder - mid_hip))
    else:
        visible = kpts[vis]
        if len(visible) == 0:
            return np.zeros(34, dtype=np.float32), 0.0
        centre = visible.mean(axis=0)
        height = float(np.linalg.norm(visible.max(axis=0) - visible.min(axis=0)))
    norm = (kpts - centre) / max(height, 1.0)
    norm[~vis] = 0.0
    return norm.ravel().astype(np.float32), height


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="yolov8m-pose.pt")
    ap.add_argument("--save-every", type=int, default=200)
    args = ap.parse_args()

    from ultralytics import YOLO
    model = YOLO(args.model)

    gold = (pd.read_csv(SELECTED, dtype=str)
            .drop_duplicates("filename")[["filename"]])
    existing = (pd.read_parquet(OUTPUT) if OUTPUT.exists()
                else pd.DataFrame(columns=["filename"]))
    have = set(existing["filename"])
    todo = gold[~gold["filename"].isin(have)]["filename"].tolist()
    print(f"gold={len(gold)}  done={len(have)}  todo={len(todo)}")

    rows, failed = [], 0
    for filename in tqdm(todo):
        try:
            res = model(str(IMAGES / filename), verbose=False)[0]
            row = {"filename": filename, "n_persons": 0, "main_conf": 0.0,
                   "main_kpt_conf": 0.0, "main_torso_px": 0.0, "main_area": 0.0,
                   **{f"main_skel_{i}": 0.0 for i in range(34)}}
            if res.boxes is not None and len(res.boxes) > 0:
                areas = ((res.boxes.xyxy[:, 2] - res.boxes.xyxy[:, 0])
                         * (res.boxes.xyxy[:, 3] - res.boxes.xyxy[:, 1]))
                j = int(areas.argmax())
                kpts = res.keypoints.xy[j].cpu().numpy()
                conf = (res.keypoints.conf[j].cpu().numpy()
                        if res.keypoints.conf is not None else np.ones(17))
                skel, torso = normalise_skeleton(kpts, conf)
                ih, iw = res.orig_shape
                row.update({"n_persons": len(res.boxes),
                            "main_conf": float(res.boxes.conf[j]),
                            "main_kpt_conf": float(conf.mean()),
                            "main_torso_px": torso,
                            "main_area": float(areas[j]) / (ih * iw)})
                row.update({f"main_skel_{i}": float(v) for i, v in enumerate(skel)})
            rows.append(row)
        except Exception as e:
            failed += 1
            sys.stderr.write(f"FAIL {filename}: {e}\n")
        if len(rows) >= args.save_every:
            existing = pd.concat([existing, pd.DataFrame(rows)], ignore_index=True)
            existing.to_parquet(OUTPUT, index=False)
            rows = []
    if rows:
        existing = pd.concat([existing, pd.DataFrame(rows)], ignore_index=True)
    existing.to_parquet(OUTPUT, index=False)
    print(f"Wrote {OUTPUT}: {len(existing)} rows. Failures: {failed}")


if __name__ == "__main__":
    main()