utils/pose_analysis.py

"""
Lógica pura de análisis postural usando la MediaPipe Tasks API (>= 0.10.14).
La API clásica mp.solutions fue eliminada en versiones recientes de MediaPipe.
"""

import cv2
import numpy as np
import urllib.request
from pathlib import Path

import mediapipe as mp
from mediapipe.tasks import python as mp_python
from mediapipe.tasks.python import vision as mp_vision

# ──────────────────────────────────────────────
# Modelo de pose — se descarga la primera vez
# ──────────────────────────────────────────────
_MODEL_PATH = Path(__file__).parent / "pose_landmarker_lite.task"
_MODEL_URL = (
    "https://storage.googleapis.com/mediapipe-models/"
    "pose_landmarker/pose_landmarker_lite/float16/latest/pose_landmarker_lite.task"
)


def _ensure_model() -> str:
    if not _MODEL_PATH.exists():
        print(f"Descargando modelo MediaPipe Pose (~5 MB)... → {_MODEL_PATH}")
        urllib.request.urlretrieve(_MODEL_URL, _MODEL_PATH)
        print("Modelo descargado.")
    return str(_MODEL_PATH)


# ──────────────────────────────────────────────
# Conexiones del esqueleto (subset de los 33 landmarks)
# ──────────────────────────────────────────────
# Índices estándar de MediaPipe Pose Landmarker
NOSE           = 0
LEFT_EYE       = 2
RIGHT_EYE      = 5
LEFT_SHOULDER  = 11
RIGHT_SHOULDER = 12
LEFT_HIP       = 23
RIGHT_HIP      = 24
LEFT_KNEE      = 25
RIGHT_KNEE     = 26
LEFT_ANKLE     = 27
RIGHT_ANKLE    = 28

POSE_CONNECTIONS = [
    # Cabeza
    (NOSE, LEFT_EYE), (NOSE, RIGHT_EYE),
    # Tronco
    (LEFT_SHOULDER, RIGHT_SHOULDER),
    (LEFT_SHOULDER, LEFT_HIP),
    (RIGHT_SHOULDER, RIGHT_HIP),
    (LEFT_HIP, RIGHT_HIP),
    # Piernas
    (LEFT_HIP, LEFT_KNEE),   (LEFT_KNEE, LEFT_ANKLE),
    (RIGHT_HIP, RIGHT_KNEE), (RIGHT_KNEE, RIGHT_ANKLE),
]

# ──────────────────────────────────────────────
# Rangos saludables para cada ángulo (en grados)
# ──────────────────────────────────────────────
ANGLE_THRESHOLDS = {
    "shoulder_tilt": {
        "min": -5.0,
        "max": 5.0,
        "label": "Inclinación de hombros",
        "correction": "Nivelá los hombros: mantenerlos a la misma altura.",
    },
    "lateral_alignment": {
        "min": 160.0,
        "max": 180.0,
        "label": "Alineación lateral (hombro–cadera–tobillo)",
        "correction": "Corregí la alineación lateral: alineá hombro, cadera y tobillo.",
    },
    "head_tilt": {
        "min": -10.0,
        "max": 10.0,
        "label": "Inclinación de cabeza",
        "correction": "Corregí la inclinación de la cabeza: mantenela centrada sobre los hombros.",
    },
}


# ──────────────────────────────────────────────
# Funciones públicas
# ──────────────────────────────────────────────

def detect_pose(image_rgb: np.ndarray):
    """
    Ejecuta MediaPipe Pose sobre una imagen RGB (ndarray HxWx3).
    Retorna la lista de landmarks del primer cuerpo detectado, o None si no hay persona.
    """
    model_path = _ensure_model()

    base_options = mp_python.BaseOptions(model_asset_path=model_path)
    options = mp_vision.PoseLandmarkerOptions(
        base_options=base_options,
        running_mode=mp_vision.RunningMode.IMAGE,   # imagen estática, no video
        num_poses=1,
        min_pose_detection_confidence=0.5,
        min_pose_presence_confidence=0.5,
        min_tracking_confidence=0.5,
    )

    with mp_vision.PoseLandmarker.create_from_options(options) as landmarker:
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image_rgb)
        result = landmarker.detect(mp_image)

    if not result.pose_landmarks:
        return None

    # Retornamos los landmarks del primer cuerpo detectado
    return result.pose_landmarks[0]


def calculate_angle(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> float:
    """
    Ángulo en grados formado en el punto B, entre los vectores BA y BC.
    a, b, c: arrays [x, y] en cualquier sistema de coordenadas.
    """
    ba = a - b
    bc = c - b
    cos_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc) + 1e-6)
    return float(np.degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0))))


def _lm_xy(landmark, img_w: int, img_h: int) -> np.ndarray:
    """Convierte un landmark normalizado a coordenadas de píxel [x, y]."""
    return np.array([landmark.x * img_w, landmark.y * img_h])


def analyze_posture(landmarks, img_w: int, img_h: int) -> dict:
    """
    Calcula los 3 ángulos posturales y evalúa si están en rango saludable.
    """
    ls   = _lm_xy(landmarks[LEFT_SHOULDER],  img_w, img_h)
    rs   = _lm_xy(landmarks[RIGHT_SHOULDER], img_w, img_h)
    lh   = _lm_xy(landmarks[LEFT_HIP],       img_w, img_h)
    la   = _lm_xy(landmarks[LEFT_ANKLE],     img_w, img_h)
    nose = _lm_xy(landmarks[NOSE],           img_w, img_h)

    # 1. Inclinación de hombros: ángulo del segmento RS→LS respecto a la horizontal
    delta_y = rs[1] - ls[1]
    delta_x = rs[0] - ls[0]
    shoulder_tilt_deg = float(np.degrees(np.arctan2(delta_y, delta_x)))

    # 2. Alineación lateral hombro–cadera–tobillo
    lateral_angle = calculate_angle(ls, lh, la)

    # 3. Inclinación de cabeza: nariz respecto al eje vertical del punto medio de hombros
    shoulder_mid  = (ls + rs) / 2
    vertical_ref  = shoulder_mid - np.array([0, 50])  # punto 50px arriba (eje vertical)
    head_tilt_deg = calculate_angle(nose, shoulder_mid, vertical_ref)
    if nose[0] > shoulder_mid[0]:                      # añadimos signo según lado
        head_tilt_deg = -head_tilt_deg

    angles = {
        "shoulder_tilt":     shoulder_tilt_deg,
        "lateral_alignment": lateral_angle,
        "head_tilt":         head_tilt_deg,
    }

    results = {}
    all_ok = True
    for key, value in angles.items():
        thresh = ANGLE_THRESHOLDS[key]
        in_range = thresh["min"] <= value <= thresh["max"]
        if not in_range:
            all_ok = False
        results[key] = {
            "angle":      round(value, 1),
            "ok":         in_range,
            "label":      thresh["label"],
            "correction": thresh["correction"],
        }

    results["all_ok"] = all_ok
    return results


def draw_skeleton(image_rgb: np.ndarray, landmarks, img_w: int, img_h: int) -> np.ndarray:
    """
    Dibuja el skeleton (conexiones + puntos) sobre la imagen con OpenCV.
    En la Tasks API no hay drawing_utils, lo hacemos a mano.
    """
    out = image_rgb.copy()

    # Puntos como píxeles
    pts = {
        i: (int(lm.x * img_w), int(lm.y * img_h))
        for i, lm in enumerate(landmarks)
    }

    # Conexiones en celeste
    for (a, b) in POSE_CONNECTIONS:
        if a in pts and b in pts:
            cv2.line(out, pts[a], pts[b], (0, 200, 255), 2, cv2.LINE_AA)

    # Puntos en blanco
    for idx, pt in pts.items():
        cv2.circle(out, pt, 4, (255, 255, 255), -1, cv2.LINE_AA)
        cv2.circle(out, pt, 4, (0, 150, 200),   1, cv2.LINE_AA)

    return out


def annotate_angles(image_rgb: np.ndarray, analysis: dict, img_w: int, img_h: int) -> np.ndarray:
    """
    Agrega texto con los valores de ángulo y un banner de resultado global.
    """
    out = image_rgb.copy()

    GREEN = (34,  197,  94)
    RED   = (239,  68,  68)

    y_start, line_h = 32, 30

    for i, (key, data) in enumerate(analysis.items()):
        if key == "all_ok":
            continue
        color  = GREEN if data["ok"] else RED
        status = "OK" if data["ok"] else "Revisar"
        text   = f"{data['label']}: {data['angle']}  [{status}]"
        cv2.putText(out, text, (10, y_start + i * line_h),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2, cv2.LINE_AA)

    # Banner inferior
    banner_color = GREEN if analysis["all_ok"] else RED
    banner_text  = "POSTURA CORRECTA" if analysis["all_ok"] else "POSTURA REQUIERE CORRECCION"
    cv2.rectangle(out, (0, img_h - 50), (img_w, img_h), banner_color, -1)
    cv2.putText(out, banner_text, (10, img_h - 15),
                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2, cv2.LINE_AA)

    return out


def generate_feedback(analysis: dict) -> str:
    """Genera el texto de feedback Markdown para mostrar en Gradio."""
    lines = ["## Analisis Postural\n"]

    for key, data in analysis.items():
        if key == "all_ok":
            continue
        icon = "OK" if data["ok"] else "Revisar"
        lines.append(f"**{data['label']}**: {data['angle']} [{icon}]")
        if not data["ok"]:
            lines.append(f"   → {data['correction']}")

    lines.append("")
    if analysis["all_ok"]:
        lines.append("### Postura correcta")
        lines.append("Todos los angulos posturales estan dentro del rango saludable.")
    else:
        lines.append("### Postura requiere correccion")
        lines.append("Revisa las indicaciones marcadas con [Revisar].")

    return "\n".join(lines)