src/envs/subenv2/node4_clip_extractor.py

"""
Node 4: Clip Signal Extractor — Sub-env 2.

Extracts pre-computed CV signals from a raw video clip using OpenCV and
MediaPipe Tasks FaceLandmarker. The resulting ``ClipSignalObservation`` is
consumed by the Clip Signal Extractor agent (Node 4) which does diagnostic
reasoning, not perception.

**No model inference is performed inline.** Phoneme sequences are accepted from
an optional pre-run forced-aligner output (e.g. Montreal Forced Aligner)
passed as an argument. Identity drift signals are computed from normalized
landmark vectors, avoiding heavyweight ArcFace runtime dependencies.

Blur score normalization
------------------------
``blur_score = clip(mean_laplacian_variance / pixel_count / CEILING, 0.0, 1.0)``

``_BLUR_CALIBRATION_CEILING`` is a calibration constant derived from the test
set. It maps per-pixel Laplacian variance of a sharp reference frame to 1.0.
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any, Optional
import urllib.error
import urllib.request

import cv2
import mediapipe as mp
import numpy as np
from mediapipe.tasks.python import BaseOptions
from mediapipe.tasks.python.vision import (
    FaceLandmarker,
    FaceLandmarkerOptions,
    RunningMode,
)
from numpy.typing import NDArray

from src.schemas.subenv2 import ClipSignalObservation

log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_MIN_FRAMES: int = 24
# Per-pixel Laplacian variance calibration ceiling.
# Empirically derived from sharp talking-head face ROIs at 480p–1080p:
# a sharp 300×300 face crop has lap_var ≈ 150–600, giving per-pixel ≈ 0.0017–0.0067.
# Setting the ceiling to 0.005 maps a sharp face to ≈ 0.33–1.0 and
# a blurry face (lap_var ≈ 10–30) to ≈ 0.002–0.02.
_BLUR_CALIBRATION_CEILING: float = 0.005
_EAR_BLINK_THRESHOLD: float = 0.20

# 468-landmark topology indices (Tasks API keeps FaceMesh indexing).
_LEFT_EYE_IDX: tuple[int, ...] = (362, 385, 387, 263, 373, 380)
_RIGHT_EYE_IDX: tuple[int, ...] = (33, 160, 158, 133, 153, 144)
_UPPER_LIP_IDX: int = 13
_LOWER_LIP_IDX: int = 14

_PROJECT_ROOT = Path(__file__).resolve().parents[3]
_FACE_LANDMARKER_URL = (
    "https://storage.googleapis.com/mediapipe-models/face_landmarker/"
    "face_landmarker/float16/latest/face_landmarker.task"
)
_DEFAULT_MODEL_CANDIDATES: tuple[Path, ...] = (
    _PROJECT_ROOT / "data" / "models" / "face_landmarker.task",
    Path.home() / ".cache" / "talkingheadbench" / "models" / "face_landmarker.task",
)


# ---------------------------------------------------------------------------
# Private helpers — model setup
# ---------------------------------------------------------------------------


def _env_truthy(name: str, *, default: bool) -> bool:
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def _candidate_landmarker_model_paths() -> list[Path]:
    env_path = os.getenv("THB_FACE_LANDMARKER_MODEL", "").strip()
    candidates: list[Path] = []
    if env_path:
        candidates.append(Path(env_path).expanduser())
    candidates.extend(_DEFAULT_MODEL_CANDIDATES)

    deduped: list[Path] = []
    seen: set[str] = set()
    for path in candidates:
        key = str(path)
        if key in seen:
            continue
        seen.add(key)
        deduped.append(path)
    return deduped


def _download_landmarker_model(dest: Path) -> Path:
    dest.parent.mkdir(parents=True, exist_ok=True)
    urllib.request.urlretrieve(_FACE_LANDMARKER_URL, dest)
    return dest


def _resolve_landmarker_model_path() -> Path | None:
    for candidate in _candidate_landmarker_model_paths():
        if candidate.exists() and candidate.is_file():
            return candidate

    if not _env_truthy("THB_AUTO_DOWNLOAD_FACE_LANDMARKER", default=True):
        return None

    cache_target = _DEFAULT_MODEL_CANDIDATES[-1]
    try:
        downloaded = _download_landmarker_model(cache_target)
    except (OSError, urllib.error.URLError, ValueError) as exc:
        log.warning(
            "Unable to auto-download FaceLandmarker model to %s: %s",
            cache_target,
            exc,
        )
        return None

    log.info("Downloaded MediaPipe FaceLandmarker model to %s", downloaded)
    return downloaded


def _create_face_landmarker() -> Any | None:
    model_path = _resolve_landmarker_model_path()
    if model_path is None:
        log.warning(
            "FaceLandmarker model file not found. Checked: %s",
            ", ".join(str(p) for p in _candidate_landmarker_model_paths()),
        )
        return None

    try:
        options = FaceLandmarkerOptions(
            base_options=BaseOptions(model_asset_path=str(model_path)),
            running_mode=RunningMode.IMAGE,
            num_faces=1,
            min_face_detection_confidence=0.5,
            min_face_presence_confidence=0.5,
            output_face_blendshapes=False,
            output_facial_transformation_matrixes=False,
        )
        return FaceLandmarker.create_from_options(options)
    except Exception as exc:  # noqa: BLE001
        log.warning(
            "Failed to initialize FaceLandmarker from %s: %s",
            model_path,
            exc,
        )
        return None


# ---------------------------------------------------------------------------
# Private helpers — signal computation
# ---------------------------------------------------------------------------


def _landmark_embedding(landmarks: list[Any]) -> NDArray[np.float32]:
    coords = np.array([(lm.x, lm.y, lm.z) for lm in landmarks], dtype=np.float32)
    centered = coords - coords.mean(axis=0, keepdims=True)
    scale = float(np.std(centered) + 1e-6)
    return (centered / scale).flatten().astype(np.float32)


def _face_bbox_from_landmarks(
    landmarks: list[Any],
    width: int,
    height: int,
    *,
    padding_ratio: float = 0.2,
) -> tuple[int, int, int, int]:
    xs = np.array([lm.x * width for lm in landmarks], dtype=np.float32)
    ys = np.array([lm.y * height for lm in landmarks], dtype=np.float32)

    x0 = int(np.clip(np.floor(xs.min()), 0, width - 1))
    x1 = int(np.clip(np.ceil(xs.max()), 1, width))
    y0 = int(np.clip(np.floor(ys.min()), 0, height - 1))
    y1 = int(np.clip(np.ceil(ys.max()), 1, height))

    pad_x = int((x1 - x0) * padding_ratio)
    pad_y = int((y1 - y0) * padding_ratio)

    x0 = max(0, x0 - pad_x)
    y0 = max(0, y0 - pad_y)
    x1 = min(width, x1 + pad_x)
    y1 = min(height, y1 + pad_y)

    if x1 <= x0:
        x1 = min(width, x0 + 1)
    if y1 <= y0:
        y1 = min(height, y0 + 1)

    return x0, y0, x1, y1


def _eye_aspect_ratio(landmarks: list[Any], indices: tuple[int, ...]) -> float:
    pts = np.array([(landmarks[i].x, landmarks[i].y) for i in indices], dtype=np.float32)
    v1 = np.linalg.norm(pts[1] - pts[5])
    v2 = np.linalg.norm(pts[2] - pts[4])
    h = np.linalg.norm(pts[0] - pts[3])
    return (v1 + v2) / (2.0 * h + 1e-6)


def _cosine_distance(a: NDArray[np.float32], b: NDArray[np.float32]) -> float:
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a < 1e-8 or norm_b < 1e-8:
        return 1.0
    return float(1.0 - np.dot(a, b) / (norm_a * norm_b))


def _laplacian_blur_score(gray: NDArray[np.uint8]) -> float:
    pixel_count = gray.shape[0] * gray.shape[1]
    lap_var = float(cv2.Laplacian(gray, cv2.CV_64F).var())
    raw = lap_var / pixel_count
    return float(np.clip(raw / _BLUR_CALIBRATION_CEILING, 0.0, 1.0))


def _exposure_score(gray: NDArray[np.uint8]) -> float:
    hist = cv2.calcHist([gray], [0], None, [256], [0, 256]).flatten()
    total = gray.size
    clipping = float((hist[0] + hist[255]) / total)
    mean_norm = float(gray.mean() / 255.0)
    mean_score = 1.0 - abs(mean_norm - 0.5) * 2.0
    return float(np.clip(mean_score * (1.0 - clipping), 0.0, 1.0))


def _parse_aligner_phonemes(aligner_output: dict) -> list[str]:
    if "phonemes" in aligner_output:
        return [str(p) for p in aligner_output["phonemes"]]

    try:
        entries = aligner_output["tiers"]["phones"]["entries"]
        return [str(entry[2]) for entry in entries]
    except (KeyError, IndexError, TypeError) as exc:
        raise ValueError(
            "aligner_output does not match expected MFA formats. "
            "Provide either {'phonemes': [...]} or the MFA TextGrid JSON export."
        ) from exc


def _phoneme_coverage_new(
    phoneme_sequence: list[str],
    current_phoneme_coverage: dict,
) -> float:
    unique_in_clip = set(phoneme_sequence)
    if not unique_in_clip:
        return 0.0
    new_count = sum(1 for p in unique_in_clip if current_phoneme_coverage.get(p, 0) == 0)
    return new_count / len(unique_in_clip)


def _lip_sync_confidence_proxy(lip_openings: list[float]) -> float:
    """Map mouth-opening variance to a lip-sync confidence score in [0, 1].

    Lip openings are normalized landmark Y-distances (range ~ 0.00–0.08).
    A talking sequence has std ≈ 0.003–0.010; silence is near 0.
    Divisor 0.008 maps:
      - active talking  (std ≈ 0.006–0.010) → 0.75–1.00
      - mild movement   (std ≈ 0.003–0.006) → 0.38–0.75
      - near-silence    (std < 0.003)        → < 0.38
    """
    if not lip_openings:
        return 0.0
    arr = np.array(lip_openings, dtype=np.float32)
    std = float(arr.std())
    return float(np.clip(std / 0.008, 0.0, 1.0))


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def extract_clip_signals(
    clip_path: Path,
    dataset_context: dict,
    aligner_output: Optional[dict] = None,
) -> ClipSignalObservation:
    """Extract CV signals from a raw video clip for Node 4."""
    clip_path = Path(clip_path)
    if not clip_path.exists():
        raise FileNotFoundError(f"Clip not found: {clip_path}")

    clip_id = clip_path.stem

    cap = cv2.VideoCapture(str(clip_path))
    if not cap.isOpened():
        raise ValueError(f"OpenCV could not open video file: {clip_path}")

    try:
        frames_bgr: list[NDArray[np.uint8]] = []
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            frames_bgr.append(frame)
    finally:
        cap.release()

    if len(frames_bgr) < _MIN_FRAMES:
        raise ValueError(
            f"Clip '{clip_id}' has only {len(frames_bgr)} frames; at least {_MIN_FRAMES} are required."
        )

    n_frames = len(frames_bgr)
    h, w = frames_bgr[0].shape[:2]

    face_landmarker = _create_face_landmarker()
    if face_landmarker is None:
        raise ValueError(
            "FaceLandmarker model file not found or failed to initialize. "
            "Set THB_FACE_LANDMARKER_MODEL or place model at data/models/face_landmarker.task."
        )

    landmark_sets: list[Optional[list[Any]]] = []
    landmark_embeddings: list[NDArray[np.float32]] = []
    lip_openings: list[float] = []
    blur_scores: list[float] = []
    exposure_scores: list[float] = []
    ear_values: list[float] = []
    occlusion_frame_count: int = 0

    try:
        for frame_bgr in frames_bgr:
            gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)

            rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb.copy())
            result = face_landmarker.detect(mp_image)

            if result.face_landmarks:
                lm = result.face_landmarks[0]
                landmark_sets.append(lm)
                landmark_embeddings.append(_landmark_embedding(lm))

                x0, y0, x1, y1 = _face_bbox_from_landmarks(lm, w, h)
                face_gray = gray[y0:y1, x0:x1]
                if face_gray.size == 0:
                    face_gray = gray
                blur_scores.append(_laplacian_blur_score(face_gray))
                exposure_scores.append(_exposure_score(face_gray))

                ear = 0.5 * (_eye_aspect_ratio(lm, _LEFT_EYE_IDX) + _eye_aspect_ratio(lm, _RIGHT_EYE_IDX))
                ear_values.append(ear)

                lip_open = abs(lm[_LOWER_LIP_IDX].y - lm[_UPPER_LIP_IDX].y)
                lip_openings.append(lip_open)
            else:
                landmark_sets.append(None)
                blur_scores.append(_laplacian_blur_score(gray))
                exposure_scores.append(_exposure_score(gray))
                ear_values.append(1.0)
                lip_openings.append(0.0)
                occlusion_frame_count += 1
    finally:
        if hasattr(face_landmarker, "close"):
            face_landmarker.close()

    if len(landmark_embeddings) >= 2:
        emb_matrix = np.stack(landmark_embeddings, axis=0)
        face_embedding_variance = float(np.var(emb_matrix, axis=0).mean())
        identity_cosine_drift = _cosine_distance(emb_matrix[0], emb_matrix[-1])
    elif len(landmark_embeddings) == 1:
        face_embedding_variance = 0.0
        identity_cosine_drift = 0.0
    else:
        face_embedding_variance = 1.0
        identity_cosine_drift = 1.0

    detected_lm = [(i, lm) for i, lm in enumerate(landmark_sets) if lm is not None]
    if len(detected_lm) >= 2:
        jitter_values: list[float] = []
        for (_, lm_a), (_, lm_b) in zip(detected_lm, detected_lm[1:]):
            pts_a = np.array([(p.x, p.y) for p in lm_a], dtype=np.float32)
            pts_b = np.array([(p.x, p.y) for p in lm_b], dtype=np.float32)
            jitter_values.append(float(np.mean(np.linalg.norm(pts_a - pts_b, axis=1))))
        landmark_stability_score = float(np.mean(jitter_values))
    else:
        landmark_stability_score = 1.0

    blink_count = 0
    in_blink = False
    for ear in ear_values:
        if ear < _EAR_BLINK_THRESHOLD:
            if not in_blink:
                blink_count += 1
                in_blink = True
        else:
            in_blink = False

    if n_frames >= 2:
        diffs: list[float] = []
        for fa_fr, fb_fr in zip(frames_bgr, frames_bgr[1:]):
            diffs.append(float(np.mean(np.abs(fa_fr.astype(np.float32) - fb_fr.astype(np.float32)))))
        frame_difference_mean = float(np.mean(diffs))
    else:
        frame_difference_mean = 0.0

    if n_frames >= 2:
        face_flows: list[float] = []
        bg_flows: list[float] = []

        for i in range(min(n_frames - 1, 30)):
            g1 = cv2.cvtColor(frames_bgr[i], cv2.COLOR_BGR2GRAY)
            g2 = cv2.cvtColor(frames_bgr[i + 1], cv2.COLOR_BGR2GRAY)
            flow = cv2.calcOpticalFlowFarneback(g1, g2, None, 0.5, 3, 15, 3, 5, 1.2, 0)
            mag = np.sqrt(flow[..., 0] ** 2 + flow[..., 1] ** 2)

            lm_a = landmark_sets[i]
            if lm_a is not None:
                xs = [int(p.x * w) for p in lm_a]
                ys = [int(p.y * h) for p in lm_a]
                x1, x2 = max(min(xs), 0), min(max(xs), w - 1)
                y1, y2 = max(min(ys), 0), min(max(ys), h - 1)
                face_mask = np.zeros((h, w), dtype=bool)
                face_mask[y1:y2, x1:x2] = True
            else:
                cx, cy = w // 2, h // 2
                face_mask = np.zeros((h, w), dtype=bool)
                face_mask[cy - h // 5 : cy + h // 5, cx - w // 5 : cx + w // 5] = True

            face_mean = float(mag[face_mask].mean()) if face_mask.any() else 0.0
            face_flows.append(face_mean)
            bg_flows.append(float(mag[~face_mask].mean() + 1e-6))

        optical_flow_magnitude = float(np.mean(face_flows)) / float(np.mean(bg_flows))
    else:
        optical_flow_magnitude = 1.0

    blur_score = float(np.mean(blur_scores))
    exposure_score_val = float(np.mean(exposure_scores))
    lip_sync_confidence = _lip_sync_confidence_proxy(lip_openings)

    if aligner_output is not None:
        phoneme_sequence = _parse_aligner_phonemes(aligner_output)
    else:
        phoneme_sequence = []

    current_phoneme_coverage: dict = dataset_context.get("current_phoneme_coverage", {})
    phone_cov_new = _phoneme_coverage_new(phoneme_sequence, current_phoneme_coverage)

    return ClipSignalObservation(
        clip_id=clip_id,
        face_embedding_variance=face_embedding_variance,
        landmark_stability_score=landmark_stability_score,
        identity_cosine_drift=identity_cosine_drift,
        frame_difference_mean=frame_difference_mean,
        optical_flow_magnitude=optical_flow_magnitude,
        blink_count=blink_count,
        lip_sync_confidence=lip_sync_confidence,
        phoneme_sequence=phoneme_sequence,
        phoneme_coverage_new=phone_cov_new,
        blur_score=blur_score,
        exposure_score=exposure_score_val,
        occlusion_frames=occlusion_frame_count,
        clips_audited_so_far=int(dataset_context.get("clips_audited_so_far", 0)),
        current_phoneme_coverage=current_phoneme_coverage,
        current_pose_distribution=dataset_context.get("current_pose_distribution", {}),
        similar_clips_accepted=int(dataset_context.get("similar_clips_accepted", 0)),
    )