"""TEASER face processor: MediaPipe detection + crop/warp + TeaserEncoder.

Analogous to ``wilor_deps/wilor_processor.py`` — bundles detection,
pre-processing, and encoding into a single ``process_frames`` call.

Ported from TEASER ``demo_video.py``.  Uses MediaPipe FaceLandmarker for
face detection (matching the original TEASER implementation) instead of YOLO.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


@dataclass
class TeaserProcessorConfig:
    """Configuration for :class:`TeaserProcessorOptimized`."""

    checkpoint_path: str = ""
    input_image_size: int = 224
    crop_scale: float = 1.4
    detection_confidence: float = 0.5
    use_amp: bool = True


# ---------------------------------------------------------------------------
# Batched crop / warp helpers  (GPU, no loops)
# ---------------------------------------------------------------------------


def crop_face_batched(
    landmarks: torch.Tensor,
    scale: float = 1.0,
    image_size: int = 224,
) -> torch.Tensor:
    """Compute similarity-transform matrices from landmark bounding boxes.

    Args:
        landmarks: ``(B, N, 2)`` — face landmarks or 2-corner
            bounding boxes ``(B, 2, 2)`` reshaped.
        scale: Padding scale factor (1.4 is standard for TEASER).
        image_size: Output crop size.

    Returns:
        ``(B, 3, 3)`` forward transformation matrices (src → dst).
    """
    B = landmarks.shape[0]
    device = landmarks.device

    min_xy, _ = torch.min(landmarks, dim=1)  # (B, 2)
    max_xy, _ = torch.max(landmarks, dim=1)  # (B, 2)

    left = min_xy[:, 0]
    right = max_xy[:, 0]
    top = min_xy[:, 1]
    bottom = max_xy[:, 1]

    old_size = (right - left + bottom - top) / 2.0
    center_x = right - (right - left) / 2.0
    center_y = bottom - (bottom - top) / 2.0
    size = (old_size * scale).long()

    dst_width = float(image_size - 1)
    s = dst_width / size.float()

    src_origin_x = center_x - (size / 2.0)
    src_origin_y = center_y - (size / 2.0)

    tx = -s * src_origin_x
    ty = -s * src_origin_y

    tform = torch.eye(3, device=device).repeat(B, 1, 1)
    tform[:, 0, 0] = s
    tform[:, 1, 1] = s
    tform[:, 0, 2] = tx
    tform[:, 1, 2] = ty

    return tform


def warp_batched(
    images: torch.Tensor,
    tforms: torch.Tensor,
    output_size: tuple[int, int] = (224, 224),
) -> torch.Tensor:
    """Batched affine warp via ``F.affine_grid`` / ``F.grid_sample``.

    Equivalent to ``skimage.transform.warp(img, tform.inverse, ...)``.

    Args:
        images: ``(B, 3, H, W)`` float tensor.
        tforms: ``(B, 3, 3)`` forward transform matrices.
        output_size: ``(H_out, W_out)``.

    Returns:
        ``(B, 3, H_out, W_out)`` warped images.
    """
    H_out, W_out = output_size
    device = images.device
    B, C, H_in, W_in = images.shape

    M_inv = torch.inverse(tforms)

    # Normalisation matrices: pixel coords ↔ [-1, 1]
    N_dst = torch.eye(3, device=device).repeat(B, 1, 1)
    N_dst[:, 0, 0] = (W_out - 1) / 2.0
    N_dst[:, 0, 2] = (W_out - 1) / 2.0
    N_dst[:, 1, 1] = (H_out - 1) / 2.0
    N_dst[:, 1, 2] = (H_out - 1) / 2.0

    N_src = torch.eye(3, device=device).repeat(B, 1, 1)
    N_src[:, 0, 0] = 2.0 / (W_in - 1)
    N_src[:, 0, 2] = -1.0
    N_src[:, 1, 1] = 2.0 / (H_in - 1)
    N_src[:, 1, 2] = -1.0

    theta_mat = torch.bmm(N_src, torch.bmm(M_inv, N_dst))
    theta = theta_mat[:, :2, :]  # (B, 2, 3)

    grid = F.affine_grid(theta, torch.Size((B, C, H_out, W_out)), align_corners=True)
    return F.grid_sample(images, grid, align_corners=True, padding_mode="zeros")


# ---------------------------------------------------------------------------
# IoU helper
# ---------------------------------------------------------------------------


def _box_iou(box_a: torch.Tensor, box_b: torch.Tensor) -> float:
    """Compute IoU between two ``(4,)`` xyxy boxes on CPU."""
    x1 = max(float(box_a[0]), float(box_b[0]))
    y1 = max(float(box_a[1]), float(box_b[1]))
    x2 = min(float(box_a[2]), float(box_b[2]))
    y2 = min(float(box_a[3]), float(box_b[3]))
    inter = max(0.0, x2 - x1) * max(0.0, y2 - y1)
    area_a = float(box_a[2] - box_a[0]) * float(box_a[3] - box_a[1])
    area_b = float(box_b[2] - box_b[0]) * float(box_b[3] - box_b[1])
    union = area_a + area_b - inter
    return inter / union if union > 0 else 0.0


# ---------------------------------------------------------------------------
# MediaPipe face detector
# ---------------------------------------------------------------------------


def _init_mediapipe_detector(confidence: float = 0.5):
    """Create a MediaPipe FaceLandmarker (downloads model automatically)."""
    import mediapipe as mp
    from mediapipe.tasks.python import vision

    base_options = mp.tasks.python.BaseOptions(
        model_asset_path=_get_mediapipe_model_path(),
    )
    options = vision.FaceLandmarkerOptions(
        base_options=base_options,
        output_face_blendshapes=False,
        output_facial_transformation_matrixes=False,
        num_faces=1,
        min_face_detection_confidence=confidence,
        min_face_presence_confidence=confidence,
    )
    return vision.FaceLandmarker.create_from_options(options)


def _get_mediapipe_model_path() -> str:
    """Return path to face_landmarker.task, downloading if needed."""
    import os
    from pathlib import Path

    # Check common locations
    cache_dir = Path(
        os.environ.get("SLTK_WEIGHTS_DIR", Path.home() / ".cache" / "sltk" / "weights")
    )
    cached = cache_dir / "face_landmarker.task"
    if cached.exists():
        return str(cached)

    # Download from Google Storage (official MediaPipe model)
    url = (
        "https://storage.googleapis.com/mediapipe-models/"
        "face_landmarker/face_landmarker/float16/latest/face_landmarker.task"
    )
    logger.info("Downloading MediaPipe face_landmarker.task...")
    cached.parent.mkdir(parents=True, exist_ok=True)

    import urllib.request

    urllib.request.urlretrieve(url, str(cached))
    logger.info("Saved to %s", cached)
    return str(cached)


def _detect_face_mediapipe(
    detector,
    frame_rgb: np.ndarray,
) -> np.ndarray | None:
    """Run MediaPipe face detection on a single RGB frame.

    Args:
        detector: MediaPipe FaceLandmarker instance.
        frame_rgb: ``(H, W, 3)`` uint8 RGB.

    Returns:
        ``(478, 2)`` pixel-coordinate landmarks, or None if no face found.
    """
    import mediapipe as mp

    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb)
    result = detector.detect(mp_image)

    if not result.face_landmarks:
        return None

    h, w = frame_rgb.shape[:2]
    lm = result.face_landmarks[0]
    landmarks = np.array(
        [[l.x * w, l.y * h] for l in lm],
        dtype=np.float32,
    )
    return landmarks


# ---------------------------------------------------------------------------
# Processor
# ---------------------------------------------------------------------------


class TeaserProcessorOptimized:
    """Bundles MediaPipe face detection and TeaserEncoder inference.

    Usage::

        cfg = TeaserProcessorConfig(checkpoint_path=...)
        proc = TeaserProcessorOptimized(cfg, torch.device("cuda:0"))
        results = proc.process_frames(frames_list, frame_indices)
    """

    def __init__(self, config: TeaserProcessorConfig, device: torch.device) -> None:
        self.config = config
        self.device = device

        # --- Load MediaPipe face detector ---
        self.detector = _init_mediapipe_detector(config.detection_confidence)

        # --- Load TeaserEncoder ---
        from sltk.extraction.teaser_deps.teaser_encoder import TeaserEncoder

        self.model = TeaserEncoder().to(device)
        checkpoint = torch.load(config.checkpoint_path, map_location=device, weights_only=False)
        # Checkpoint contains both encoder and generator; strip prefix
        encoder_state = {
            k.replace("teaser_encoder.", ""): v
            for k, v in checkpoint.items()
            if "teaser_encoder" in k
        }
        self.model.load_state_dict(encoder_state)
        self.model.eval()
        logger.info(
            "TEASER encoder loaded (%d params)", sum(p.numel() for p in self.model.parameters())
        )

    # -----------------------------------------------------------------------
    # Main entry point
    # -----------------------------------------------------------------------

    def process_frames(
        self,
        frames: list[np.ndarray],
        frame_indices: list[int],
    ) -> dict[str, Any]:
        """Run detection + encoding on a list of BGR uint8 frames.

        Args:
            frames: List of ``(H, W, 3)`` BGR uint8 numpy arrays.
            frame_indices: Corresponding absolute frame indices.

        Returns:
            Results dict ready for H5 serialisation (see plan for schema).
        """
        total_frames = len(frames)

        # Accumulators
        all_jaw: list[torch.Tensor] = []
        all_expr: list[torch.Tensor] = []
        all_eyelid: list[torch.Tensor] = []
        all_shape: list[torch.Tensor] = []
        all_pose: list[torch.Tensor] = []
        all_cam: list[torch.Tensor] = []
        all_boxes: list[torch.Tensor] = []
        all_scores: list[torch.Tensor] = []
        successful_frames: list[int] = []

        prev_box: torch.Tensor | None = None

        # --- MediaPipe detection per frame ---
        valid_indices: list[int] = []
        valid_frames: list[np.ndarray] = []
        selected_boxes: list[torch.Tensor] = []
        selected_scores: list[torch.Tensor] = []

        for i, frame_bgr in enumerate(frames):
            frame_rgb = frame_bgr[:, :, ::-1].copy()  # BGR → RGB
            landmarks = _detect_face_mediapipe(self.detector, frame_rgb)

            if landmarks is None:
                continue

            # Convert landmarks to xyxy bounding box
            x_min, y_min = landmarks.min(axis=0)
            x_max, y_max = landmarks.max(axis=0)
            box = torch.tensor([x_min, y_min, x_max, y_max], dtype=torch.float32)

            # IoU tracking for multi-face consistency (shouldn't happen with
            # num_faces=1, but kept for robustness)
            if prev_box is not None:
                iou = _box_iou(prev_box, box)
                # If IoU is very low, face might have jumped — still accept
                # since MediaPipe returns only 1 face
                _ = iou  # kept for future multi-face support

            score = torch.tensor(1.0)  # MediaPipe doesn't expose per-detection scores
            prev_box = box

            valid_indices.append(i)
            valid_frames.append(frame_bgr)
            selected_boxes.append(box)
            selected_scores.append(score)
            successful_frames.append(frame_indices[i])

        if len(valid_indices) == 0:
            return self._empty_results(total_frames)

        # --- Crop & warp using landmark bounding boxes ---
        boxes_tensor = torch.stack(selected_boxes).reshape(-1, 2, 2).to(self.device)
        tforms = crop_face_batched(
            boxes_tensor,
            scale=self.config.crop_scale,
            image_size=self.config.input_image_size,
        )

        images_tensor = (
            torch.from_numpy(np.stack(valid_frames))
            .permute(0, 3, 1, 2)
            .to(torch.float32)
            .to(self.device)[:, [2, 1, 0]]  # BGR → RGB
            / 255.0
        )
        size = self.config.input_image_size
        cropped = warp_batched(images_tensor, tforms, output_size=(size, size))

        # --- Encode (no tokens saved) ---
        with torch.no_grad():
            if self.config.use_amp:
                with torch.autocast(device_type="cuda"):
                    outputs = self.model(cropped)
            else:
                outputs = self.model(cropped)

        all_jaw.append(outputs["jaw_params"].cpu())
        all_expr.append(outputs["expression_params"].cpu())
        all_eyelid.append(outputs["eyelid_params"].cpu())
        all_shape.append(outputs["shape_params"].cpu())
        all_pose.append(outputs["pose_params"].cpu())
        all_cam.append(outputs["cam"].cpu())
        all_boxes.extend(selected_boxes)
        all_scores.extend(selected_scores)

        # --- Assemble results ---
        return self._assemble_results(
            all_jaw,
            all_expr,
            all_eyelid,
            all_shape,
            all_pose,
            all_cam,
            all_boxes,
            all_scores,
            successful_frames,
            total_frames,
        )

    # -----------------------------------------------------------------------
    # Helpers
    # -----------------------------------------------------------------------

    @staticmethod
    def _empty_results(total_frames: int) -> dict[str, Any]:
        return {
            "flame": {
                "jaw_pose": np.zeros((0, 3), dtype=np.float32),
                "expression": np.zeros((0, 50), dtype=np.float32),
                "eyelid_params": np.zeros((0, 2), dtype=np.float32),
                "shape": np.zeros((0, 300), dtype=np.float32),
            },
            "pose_params": np.zeros((0, 3), dtype=np.float16),
            "cam": np.zeros((0, 3), dtype=np.float16),
            "boxes": np.zeros((0, 4), dtype=np.float16),
            "boxes_scores": np.zeros((0,), dtype=np.float16),
            "successful_frames": np.array([], dtype=np.int32),
            "frame_idx": np.zeros((total_frames, 2), dtype=np.int32),
            "num_detections": 0,
        }

    @staticmethod
    def _assemble_results(
        all_jaw: list[torch.Tensor],
        all_expr: list[torch.Tensor],
        all_eyelid: list[torch.Tensor],
        all_shape: list[torch.Tensor],
        all_pose: list[torch.Tensor],
        all_cam: list[torch.Tensor],
        all_boxes: list[torch.Tensor],
        all_scores: list[torch.Tensor],
        successful_frames: list[int],
        total_frames: int,
    ) -> dict[str, Any]:
        jaw = torch.cat(all_jaw).numpy().astype(np.float32)
        expr = torch.cat(all_expr).numpy().astype(np.float32)
        eyelid = torch.cat(all_eyelid).numpy().astype(np.float32)
        shape = torch.cat(all_shape).numpy().astype(np.float32)
        pose = torch.cat(all_pose).numpy().astype(np.float16)
        cam = torch.cat(all_cam).numpy().astype(np.float16)
        boxes = torch.stack(all_boxes).numpy().astype(np.float16)
        scores = torch.tensor(all_scores).numpy().astype(np.float16)
        sf = np.array(successful_frames, dtype=np.int32)

        M = len(sf)

        # Build sparse frame_idx (T, 2) — [start, count] per frame
        frame_idx = np.zeros((total_frames, 2), dtype=np.int32)
        for det_idx, fid in enumerate(successful_frames):
            if 0 <= fid < total_frames:
                frame_idx[fid] = [det_idx, 1]

        return {
            "flame": {
                "jaw_pose": jaw,
                "expression": expr,
                "eyelid_params": eyelid,
                "shape": shape,
            },
            "pose_params": pose,
            "cam": cam,
            "boxes": boxes,
            "boxes_scores": scores,
            "successful_frames": sf,
            "frame_idx": frame_idx,
            "num_detections": M,
        }
