"""
SignRep model wrapper for computing sign language video embeddings.

This is a backward-compatible wrapper around the full SignRepPipeline.
The original embed_video() / embed_frames() / embed_videos() API is preserved.
For full pipeline capabilities (dictionary extraction, continuous extraction,
gloss spotting), use SignRepPipeline directly or access it via the .pipeline
property.

Reference:
    Sherpa et al. "SignRep: Enhancing Self-Supervised Sign Representations"
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)


def _get_default_checkpoint() -> Path:
    """Get default SignRep checkpoint from weight system or bundled location."""
    from sltk.extraction.weights import get_weight_path

    path = get_weight_path("signrep_checkpoint")
    if path:
        return path

    bundled = Path(__file__).parent.parent / "weights" / "signrep" / "ckpt.pt"
    return bundled


DEFAULT_CHECKPOINT = _get_default_checkpoint()


@dataclass
class SignRepConfig:
    """Configuration for SignRep model."""

    checkpoint_path: str | Path = field(default_factory=lambda: DEFAULT_CHECKPOINT)
    device: str = "cuda"
    segment_length: int = 16  # Number of frames per segment
    stride: int = 2  # Stride between segments
    out_dim: int = 768
    mask_ratio: float = 0.0
    batch_size: int = 8
    num_workers: int = 4
    amp: bool = True
    preload: bool = False

    # Model architecture parameters (kept for backward compat)
    backbone_params: dict[str, Any] = field(
        default_factory=lambda: {
            "drop_path_rate": 0.0,
            "num_cls_tokens": 0,
        }
    )

    cls_head_params: dict[str, Any] = field(
        default_factory=lambda: {
            "latents_params": {"hidden_dim": 512, "out_dim": 384},
            "total_frames": 16,
            "projection_params": {
                "full_body_angles": {"hidden_dim": None, "output_dim": 22 * 2},
                "full_right_angles": {"hidden_dim": None, "output_dim": 41 * 2},
                "full_left_angles": {"hidden_dim": None, "output_dim": 41 * 2},
                "body_pose": {"hidden_dim": None, "output_dim": 61 * 3},
                "right_pose": {"hidden_dim": None, "output_dim": 21 * 3},
                "left_pose": {"hidden_dim": None, "output_dim": 21 * 3},
                "body_dist": {"hidden_dim": None, "output_dim": 12 * 22 * 3},
                "right_dist": {"hidden_dim": None, "output_dim": 5 * 11 * 3},
                "left_dist": {"hidden_dim": None, "output_dim": 5 * 11 * 3},
            },
        }
    )


# Mapping from old aggregation names to pipeline pooling methods
_AGGREGATION_TO_POOLING = {
    "mean": "mean_pool",
    "max": "max_pool",
    "first": "middle",
    "last": "middle",
}


class SignRepEmbedder:
    """
    SignRep model wrapper for computing video embeddings.

    This class provides a high-level interface for:
    - Loading the SignRep model
    - Extracting embeddings from video files
    - Extracting embeddings from pre-loaded frames
    - Batch processing multiple videos

    For advanced features (dictionary extraction, gloss spotting),
    use the .pipeline property to access the full SignRepPipeline.

    Example:
        embedder = SignRepEmbedder()

        # Single video
        embedding = embedder.embed_video("video.mp4")

        # Multiple videos
        embeddings = embedder.embed_videos(["video1.mp4", "video2.mp4"])

        # From frames (T, H, W, C) numpy array
        embedding = embedder.embed_frames(frames)

        # Access full pipeline
        result = embedder.pipeline.extract_dictionary("sign.mp4")
    """

    def __init__(
        self,
        config: SignRepConfig | None = None,
        lazy_load: bool = True,
    ):
        self.config = config or SignRepConfig()
        self._pipeline: Any = None
        self._transform: Any = None

        if not lazy_load:
            self._load_pipeline()

    def _load_pipeline(self):
        """Load the SignRep pipeline."""
        if self._pipeline is not None:
            return

        from sltk.embedding.pipeline import SignRepPipeline

        self._pipeline = SignRepPipeline(
            checkpoint=str(self.config.checkpoint_path),
            device=self.config.device,
            amp=self.config.amp,
            batch_size=self.config.batch_size,
            num_workers=self.config.num_workers,
            preload=self.config.preload,
        )

        logger.info("SignRep pipeline loaded successfully")

    @property
    def pipeline(self):
        """Access the full SignRepPipeline for advanced features."""
        self._load_pipeline()
        return self._pipeline

    def embed_video(
        self,
        video_path: str | Path,
        aggregation: str = "mean",
        return_all_segments: bool = False,
    ) -> np.ndarray | dict[str, np.ndarray]:
        """
        Compute embedding for a video file.

        Args:
            video_path: Path to video file
            aggregation: How to aggregate segment embeddings ('mean', 'max', 'first', 'last')
            return_all_segments: If True, return all segment embeddings

        Returns:
            If return_all_segments: dict with 'embedding' and 'segments' keys
            Otherwise: Single embedding vector (768-dim by default)
        """
        self._load_pipeline()

        method = _AGGREGATION_TO_POOLING.get(aggregation, "mean_pool")
        result = self._pipeline.extract_dictionary(
            video_path,
            method=method,
            stride=self.config.stride,
        )

        if return_all_segments:
            return {
                "embedding": result.feature,
                "segments": result.all_features,
                "latents": result.all_features,
                "num_segments": len(result.all_features),
            }

        return result.feature

    def embed_frames(
        self,
        frames: np.ndarray,
        aggregation: str = "mean",
        return_all_segments: bool = False,
    ) -> np.ndarray | dict[str, np.ndarray]:
        """
        Compute embedding from video frames.

        Args:
            frames: Video frames as (T, H, W, C) numpy array (RGB)
            aggregation: How to aggregate segment embeddings
            return_all_segments: If True, return all segment embeddings

        Returns:
            Embedding vector or dict with all segments
        """
        self._load_pipeline()
        import torch

        from sltk.embedding.pipeline import SEGMENT_LENGTH, FrameTransform

        if self._transform is None:
            self._transform = FrameTransform()

        # Transform each frame
        transformed_list: list[Any] = []
        for frame in frames:
            transformed_list.append(self._transform(frame))
        transformed = torch.stack(transformed_list)  # (T, C, H, W)

        segment_features = []
        stride = self.config.stride
        seg_len = SEGMENT_LENGTH

        with torch.inference_mode():
            for i in range(0, transformed.shape[0], stride):
                segment = transformed[i : i + seg_len]
                if segment.shape[0] < seg_len:
                    continue

                # (T, C, H, W) -> (1, C, T, H, W)
                segment = segment.unsqueeze(0).permute(0, 2, 1, 3, 4)
                segment = segment.to(self._pipeline.device)

                if self._pipeline.amp:
                    with torch.amp.autocast("cuda"):
                        output = self._pipeline.model.inference_forward(segment)
                else:
                    output = self._pipeline.model.inference_forward(segment)

                latent = output["latent"][0].cpu().numpy()
                segment_features.append(latent)

        if not segment_features:
            raise ValueError("Video too short to extract any segments")

        features_array = np.array(segment_features)

        # Aggregate
        if aggregation == "mean":
            embedding = features_array.mean(axis=0)
        elif aggregation == "max":
            embedding = features_array.max(axis=0)
        elif aggregation == "first":
            embedding = features_array[0]
        elif aggregation == "last":
            embedding = features_array[-1]
        else:
            raise ValueError(f"Unknown aggregation: {aggregation}")

        # L2-normalize
        norm = np.linalg.norm(embedding)
        if norm > 0:
            embedding = embedding / norm

        if return_all_segments:
            return {
                "embedding": embedding,
                "segments": features_array,
                "latents": features_array,
                "num_segments": len(segment_features),
            }

        return embedding

    def embed_videos(
        self,
        video_paths: list[str | Path],
        aggregation: str = "mean",
        show_progress: bool = True,
    ) -> dict[str, np.ndarray | dict[str, np.ndarray] | None]:
        """
        Compute embeddings for multiple videos.

        Args:
            video_paths: List of video file paths
            aggregation: How to aggregate segment embeddings
            show_progress: Whether to show progress bar

        Returns:
            Dictionary mapping video path to embedding
        """
        results: dict[str, np.ndarray | dict[str, np.ndarray] | None] = {}

        if show_progress:
            try:
                from tqdm import tqdm

                iterator = tqdm(video_paths, desc="Embedding videos")
            except ImportError:
                iterator = video_paths
        else:
            iterator = video_paths

        for path in iterator:
            path_str = str(path)
            try:
                results[path_str] = self.embed_video(path, aggregation)
            except Exception as e:
                logger.warning("Failed to embed %s: %s", path, e)
                results[path_str] = None

        return results

    @property
    def embedding_dim(self) -> int:
        """Get embedding dimension."""
        return self.config.out_dim

    def is_loaded(self) -> bool:
        """Check if model is loaded."""
        return self._pipeline is not None
