"""TEASER face processor: YOLO 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``.
"""

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 = ""
    face_detector_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 bounding-box corners.

    Args:
        landmarks: ``(B, N, 2)`` — either 5-point face landmarks or 2-corner
            bounding boxes ``(B, 2, 2)`` from YOLO ``xyxy`` 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


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

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

    Usage::

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

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

        # --- Load YOLOv8 face detector ---
        from ultralytics import YOLO

        self.detector = YOLO(config.face_detector_path).to(device)

        # --- 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

        # --- YOLO detection on all frames at once ---
        with torch.no_grad():
            detections = self.detector(frames, verbose=False)

        # --- Per-frame: select best face, track via IoU ---
        valid_indices: list[int] = []
        valid_frames: list[np.ndarray] = []
        selected_boxes: list[torch.Tensor] = []
        selected_scores: list[torch.Tensor] = []

        for i, result in enumerate(detections):
            if len(result.boxes) == 0:
                continue

            if len(result.boxes) == 1:
                best_idx = 0
            elif prev_box is None:
                best_idx = int(torch.argmax(result.boxes.conf))
            else:
                # IoU tracking — pick box closest to previous
                ious = [
                    _box_iou(prev_box, result.boxes.xyxy[j].detach().cpu())
                    for j in range(len(result.boxes))
                ]
                best_idx = int(np.argmax(ious))

            box = result.boxes.xyxy[best_idx].detach().cpu()
            score = result.boxes.conf[best_idx].detach().cpu()
            prev_box = box

            valid_indices.append(i)
            valid_frames.append(frames[i])
            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 ---
        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,
        }
