"""
Geometric body-region segmentation using MediaPipe pose landmarks.

SAM2 segments by pixel similarity, which on a clothed photo yields the tank-top
color region rather than the underlying anatomy. For breast/buttocks regions we
build elliptical masks positioned by pose landmarks (shoulders, hips, knees) so
the result is anatomically located regardless of clothing.
"""
import sys
import numpy as np
from PIL import Image


def log(msg: str):
    print(msg, flush=True)


def _ellipse_mask(W: int, H: int, cx: float, cy: float, rx: float, ry: float, angle: float = 0.0) -> np.ndarray:
    """Build a (H, W) bool mask: True inside the ellipse at (cx,cy) with semi-axes (rx,ry), rotated by `angle` rad."""
    ys, xs = np.mgrid[0:H, 0:W]
    dx = xs - cx
    dy = ys - cy
    cos_a = np.cos(-angle)
    sin_a = np.sin(-angle)
    rx_d = (dx * cos_a - dy * sin_a) / max(rx, 1e-3)
    ry_d = (dx * sin_a + dy * cos_a) / max(ry, 1e-3)
    return (rx_d * rx_d + ry_d * ry_d) <= 1.0


def _bbox_of_mask(mask: np.ndarray) -> list[int]:
    rows = np.any(mask, axis=1)
    cols = np.any(mask, axis=0)
    if not rows.any() or not cols.any():
        return [0, 0, 0, 0]
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    return [int(cmin), int(rmin), int(cmax - cmin + 1), int(rmax - rmin + 1)]


def segment_anatomy(image: Image.Image, regions: list[str]) -> dict:
    """
    Build geometric masks for body regions using MediaPipe pose landmarks.

    Supported regions: breast_left, breast_right, buttocks.
    Returns {region: {"mask": list[list[bool]], "bbox": [x,y,w,h]}}.
    Skips regions for which the required landmarks aren't visible.
    """
    from pose import detect_landmarks  # reuse the same MediaPipe wrapper

    log(f"[Anatomy] Geometric segmentation for {regions} on image {image.size}")
    landmarks = detect_landmarks(image)

    W, H = image.size

    def get_pixel(idx: int) -> tuple[float, float] | None:
        lm = landmarks.get(idx)
        if not lm or lm.get("visibility", 0) < 0.3:
            return None
        return float(lm["px"]), float(lm["py"])

    sl = get_pixel(11)   # left shoulder (image-right side of subject)
    sr = get_pixel(12)   # right shoulder
    hl = get_pixel(23)   # left hip
    hr = get_pixel(24)   # right hip
    kl = get_pixel(25)   # left knee
    kr = get_pixel(26)   # right knee

    if not (sl and sr and hl and hr):
        log("[Anatomy] Missing shoulder/hip landmarks — cannot build geometric masks.")
        return {}

    # Sternum (midpoint between shoulders)
    sx = (sl[0] + sr[0]) / 2
    sy = (sl[1] + sr[1]) / 2
    # Pelvis (midpoint between hips)
    px = (hl[0] + hr[0]) / 2
    py = (hl[1] + hr[1]) / 2

    shoulder_width = float(np.hypot(sl[0] - sr[0], sl[1] - sr[1]))
    torso_height  = float(np.hypot(sx - px, sy - py))
    torso_angle   = float(np.arctan2(py - sy, px - sx) - np.pi / 2)  # 0 if perfectly upright

    log(f"[Anatomy] shoulders=({sl[0]:.0f},{sl[1]:.0f})↔({sr[0]:.0f},{sr[1]:.0f}) "
        f"hips=({hl[0]:.0f},{hl[1]:.0f})↔({hr[0]:.0f},{hr[1]:.0f}) "
        f"shoulder_w={shoulder_width:.0f} torso_h={torso_height:.0f}")

    results: dict = {}

    # ── Breasts: elliptical regions sitting on the upper torso, halfway between
    #    each shoulder and the sternum, dropped 25% of the torso height down. ──
    if "breast_left" in regions or "breast_right" in regions:
        # Vertical drop from shoulder line toward pelvis
        drop_x = (px - sx) * 0.30
        drop_y = (py - sy) * 0.30
        for region, shoulder in (("breast_left", sl), ("breast_right", sr)):
            if region not in regions:
                continue
            cx = (shoulder[0] + sx) / 2 + drop_x
            cy = (shoulder[1] + sy) / 2 + drop_y
            rx = abs(shoulder[0] - sx) * 0.55
            ry = torso_height * 0.18
            mask = _ellipse_mask(W, H, cx, cy, rx, ry, torso_angle)
            log(f"[Anatomy] {region} center=({cx:.0f},{cy:.0f}) rx={rx:.0f} ry={ry:.0f}")
            results[region] = {"mask": mask.tolist(), "bbox": _bbox_of_mask(mask)}

    # ── Buttocks: ellipse spanning hips, dropped halfway toward knees. ──
    if "buttocks" in regions:
        if kl and kr:
            kx = (kl[0] + kr[0]) / 2
            ky = (kl[1] + kr[1]) / 2
            cx = (px + kx) / 2 * 0.5 + px * 0.5     # weighted toward hips
            cy = py + (ky - py) * 0.25
        else:
            cx, cy = px, py + torso_height * 0.20
        hip_width = float(np.hypot(hl[0] - hr[0], hl[1] - hr[1]))
        rx = hip_width * 0.65
        ry = torso_height * 0.30
        mask = _ellipse_mask(W, H, cx, cy, rx, ry, torso_angle)
        log(f"[Anatomy] buttocks center=({cx:.0f},{cy:.0f}) rx={rx:.0f} ry={ry:.0f}")
        results["buttocks"] = {"mask": mask.tolist(), "bbox": _bbox_of_mask(mask)}

    log(f"[Anatomy] Built {len(results)} geometric masks.")
    return results