"""High-level FLAME mesh reconstruction and expression classification.

Usage:
    from sltk.extraction.teaser_deps.flame import FLAMEProcessor

    proc = FLAMEProcessor(device="cuda:0")
    result = proc.reconstruct_from_h5("face_teaser.h5")
    clusters = proc.classify_expression(result["vertices"])
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

import numpy as np
import torch

from sltk.extraction.weights import get_weight_path

logger = logging.getLogger(__name__)


class FLAMEProcessor:
    """High-level FLAME mesh reconstruction from TEASER H5 parameters.

    Loads the FLAME model, indices of interest, and K-means centroids
    for expression classification. Provides batched forward pass and
    nearest-centroid classification.

    Args:
        device: Torch device string.
    """

    def __init__(self, device: str = "cuda:0") -> None:
        self.device = torch.device(device)
        self._flame: Any = None
        self._indices: torch.Tensor | None = None
        self._centroids: torch.Tensor | None = None
        self._initialized = False

    def _ensure_initialized(self) -> None:
        """Lazy-load FLAME model, indices, and centroids."""
        if self._initialized:
            return

        from sltk.extraction.teaser_deps.flame.flame_model import FLAME

        logger.info("Loading FLAME model on %s", self.device)
        self._flame = FLAME().to(self.device).eval()

        # Indices of interest (lower-face vertices for classification)
        indices_path = get_weight_path("flame_indices")
        if indices_path is None:
            raise FileNotFoundError("FLAME indices_of_interest.pt not found")
        self._indices = torch.load(str(indices_path), weights_only=True).to(self.device)
        logger.debug("Loaded %d indices of interest", len(self._indices))

        # K-means centroids: (20, 6261) = 20 clusters x 2087 vertices x 3
        centroids_path = get_weight_path("flame_centroids")
        if centroids_path is None:
            raise FileNotFoundError("FLAME centroids not found")
        self._centroids = torch.load(str(centroids_path), weights_only=True).to(
            self.device
        )
        logger.debug("Loaded centroids: %s", self._centroids.shape)

        self._initialized = True

    def load_h5(self, h5_path: str | Path) -> dict[str, Any]:
        """Read TEASER H5 file and return FLAME parameters.

        Supports both flat layout (keys at root) and grouped layout
        (keys inside numbered groups like '1.0').

        Args:
            h5_path: Path to face_teaser.h5 file.

        Returns:
            Dict with keys: flame (dict of param arrays), successful_frames,
            cam, pose_params, boxes, boxes_scores, meta.
        """
        import h5py

        h5_path = Path(h5_path)
        if not h5_path.exists():
            raise FileNotFoundError(f"H5 file not found: {h5_path}")

        required_keys = ["successful_frames", "cam", "pose_params"]
        flame_keys = ["expression", "jaw_pose", "eyelid_params"]

        def _read_group(hf_group: h5py.Group) -> dict[str, Any]:
            out: dict[str, Any] = {}
            for name in required_keys:
                if name not in hf_group:
                    raise KeyError(f"Missing dataset: {name}")
                out[name] = hf_group[name][()]

            out["boxes"] = hf_group["boxes"][()]
            if "boxes_scores" in hf_group:
                out["boxes_scores"] = hf_group["boxes_scores"][()]

            if "flame" not in hf_group:
                raise KeyError("Missing group: flame")
            out["flame"] = {}
            for name in flame_keys:
                if name in hf_group["flame"]:
                    out["flame"][name] = hf_group["flame"][name][()]
            return out

        result: dict[str, Any] = {}
        with h5py.File(str(h5_path), "r") as hf:
            result["meta"] = {key: hf.attrs[key] for key in hf.attrs.keys()}

            if "successful_frames" in hf:
                # Flat layout
                group_data = _read_group(hf)
                result.update(group_data)
            else:
                # Grouped layout (multiple face tracks)
                for group_name in hf.keys():
                    result[group_name] = _read_group(hf[group_name])

        return result

    def reconstruct_from_h5(
        self,
        h5_path: str | Path,
        batch_size: int = 2048,
    ) -> dict[str, Any]:
        """Reconstruct FLAME meshes from TEASER H5 file.

        Args:
            h5_path: Path to face_teaser.h5.
            batch_size: Batch size for FLAME forward pass.

        Returns:
            Dict with vertices (N, 5023, 3), landmarks_fan_3d (N, 68, 3),
            successful_frames, and the original H5 data.
        """
        self._ensure_initialized()
        h5_data = self.load_h5(h5_path)

        # Get flame params from the right level
        if "flame" in h5_data:
            flame_params = h5_data["flame"]
        else:
            # Grouped layout: use first group
            first_key = [k for k in h5_data if k != "meta"][0]
            flame_params = h5_data[first_key]["flame"]

        result = self.reconstruct_from_params(flame_params, batch_size=batch_size)
        result["h5_data"] = h5_data
        return result

    def reconstruct_from_params(
        self,
        flame_params: dict[str, np.ndarray | torch.Tensor],
        batch_size: int = 2048,
    ) -> dict[str, Any]:
        """Reconstruct FLAME meshes from parameter arrays.

        Args:
            flame_params: Dict with keys like expression, jaw_pose, eyelid_params.
            batch_size: Batch size for forward pass.

        Returns:
            Dict with vertices, landmarks_fan_3d tensors on the processor device.
        """
        self._ensure_initialized()
        assert self._flame is not None

        # Determine total number of frames
        num_data = 0
        for val in flame_params.values():
            if hasattr(val, "shape"):
                num_data = max(num_data, val.shape[0])

        all_vertices: list[torch.Tensor] = []
        all_landmarks: list[torch.Tensor] = []

        for i in range(0, num_data, batch_size):
            batch_params = {
                key: torch.as_tensor(val[i : i + batch_size], dtype=torch.float32).to(
                    self.device
                )
                for key, val in flame_params.items()
                if hasattr(val, "shape")
            }
            with torch.no_grad():
                out = self._flame.forward(batch_params)
            all_vertices.append(out["vertices"])
            all_landmarks.append(out["landmarks_fan_3d"])

        vertices = torch.cat(all_vertices, dim=0)
        landmarks = torch.cat(all_landmarks, dim=0)

        return {
            "vertices": vertices,
            "landmarks_fan_3d": landmarks,
        }

    def classify_expression(self, vertices: torch.Tensor) -> torch.Tensor:
        """Classify expression by nearest K-means centroid.

        Selects lower-face vertices (indices_of_interest), flattens,
        and finds nearest centroid via torch.cdist.

        Args:
            vertices: FLAME mesh vertices (N, 5023, 3).

        Returns:
            Cluster IDs (N,) as long tensor, values in [0, 19].
        """
        self._ensure_initialized()
        assert self._indices is not None
        assert self._centroids is not None

        # Select lower-face vertices: (N, 2087, 3)
        selected = vertices[:, self._indices]
        # Flatten to (N, 6261)
        flat = selected.reshape(selected.shape[0], -1)

        # Normalize features (L2 norm per sample)
        flat = flat / (flat.norm(dim=1, keepdim=True) + 1e-8)
        centroids_norm = self._centroids / (
            self._centroids.norm(dim=1, keepdim=True) + 1e-8
        )

        # Nearest centroid via cdist: (N, 20)
        dists = torch.cdist(flat.unsqueeze(0), centroids_norm.unsqueeze(0)).squeeze(0)
        cluster_ids = dists.argmin(dim=1)

        return cluster_ids

    def get_landmarks_2d(
        self, vertices: torch.Tensor, cam: np.ndarray | None = None
    ) -> np.ndarray:
        """Get 68-point FAN landmarks projected to 2D.

        If cam (camera) params are not provided, returns the XY coordinates
        of the 3D landmarks directly (orthographic projection).

        Args:
            vertices: FLAME mesh vertices (N, 5023, 3).
            cam: Optional camera parameters (N, 3) [scale, tx, ty].

        Returns:
            2D landmarks (N, 68, 2) as numpy array.
        """
        self._ensure_initialized()
        assert self._flame is not None

        with torch.no_grad():
            landmarks3d = self._flame.seletec_3d68(vertices)

        # FLAME Y-axis points up, image Y points down — flip Y
        lmk = landmarks3d[:, :, :2].clone()  # (N, 68, 2) — take XY
        lmk[:, :, 1] = -lmk[:, :, 1]

        if cam is not None:
            cam_t = torch.as_tensor(cam, dtype=torch.float32, device=self.device)
            scale = cam_t[:, 0:1].unsqueeze(-1)  # (N, 1, 1)
            tx = cam_t[:, 1:2].unsqueeze(-1)  # (N, 1, 1)
            ty = cam_t[:, 2:3].unsqueeze(-1)  # (N, 1, 1)
            lmk = lmk * scale + torch.cat([tx, ty], dim=-1)

        return lmk.cpu().numpy()
