"""Sign segmentation using WiLoR features (v2 architecture).

Alternative segmenter that uses only WiLoR hand features (no MediaPipe),
potentially better calibrated for certain datasets.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import torch

from sltk.data.core import Segment, SegmentList
from sltk.exceptions import ConfigurationError

# Token constants (same as original segmenter)
OUT_SIGN_TOKEN = 0  # Background/no sign
IN_SIGN_TOKEN = 1  # Inside sign (active)
BEGIN_TOKEN = 2  # Sign boundary


class SignSegmenterV2:
    """Sign segmentation using WiLoR features (v2 architecture).

    This segmenter uses a different architecture than the original:
    - Uses only WiLoR hand features (192-dim, no MediaPipe needed)
    - Transformer encoder with 4 layers, 4 heads, 512 hidden dims
    - Proven post-processing for proper segment boundaries

    Example:
        >>> from sltk.segmentation import SignSegmenterV2
        >>> from sltk.extraction.weights import get_weight_path
        >>>
        >>> checkpoint = get_weight_path("segmentor_v2")
        >>> segmenter = SignSegmenterV2(model_path=checkpoint)
        >>> segments = segmenter.segment(wilor_features, fps=25.0)
    """

    def __init__(
        self,
        model_path: str | Path | None = None,
        config: dict | None = None,
        device: str = "cuda",
    ):
        """Initialize SignSegmenterV2.

        Args:
            model_path: Path to model checkpoint (.ckpt file)
            config: Optional config dict (uses defaults if None)
            device: Device for inference ('cuda' or 'cpu')
        """
        self.device = device
        self.model: Any = None

        # Default config
        self.config = {
            "window_size": 25,
            "subsample": 1,
            "hidden_size": 512,
            "num_layers": 4,
            "num_heads": 4,
            "ff_size": 1024,
            "dropout": 0.2,
            "temp_conv_kernel": 3,
            "num_classes": 3,
        }

        # Override with user config
        if config is not None:
            self.config.update(config)

        # Load model if path provided
        if model_path is not None:
            self.load_model(model_path)

    def load_model(self, checkpoint_path: str | Path):
        """Load model from checkpoint.

        Args:
            checkpoint_path: Path to PyTorch Lightning checkpoint (.ckpt)

        Raises:
            FileNotFoundError: If checkpoint not found
            ConfigurationError: If checkpoint format is invalid
        """
        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        from sltk.segmentation._signsegmentor_v2.model import TransformerEncoderV2

        # Create model
        self.model = TransformerEncoderV2(
            input_size=192,
            hidden_size=self.config["hidden_size"],
            num_layers=self.config["num_layers"],
            num_heads=self.config["num_heads"],
            ff_size=self.config["ff_size"],
            dropout=self.config["dropout"],
            temp_conv_kernel=self.config["temp_conv_kernel"],
            num_classes=self.config["num_classes"],
        )

        # Load checkpoint (handle Lightning format)
        # weights_only=False needed: Lightning checkpoints contain PosixPath
        checkpoint = torch.load(
            checkpoint_path,
            map_location=self.device,
            weights_only=False,
        )

        # Extract state dict (Lightning saves under "state_dict" key)
        if "state_dict" in checkpoint:
            state_dict = checkpoint["state_dict"]

            # Remove Lightning module prefix if present
            # Keys in Lightning checkpoint may have "model." or "encoder." prefix
            cleaned_state_dict = {}
            for key, value in state_dict.items():
                # Remove prefixes like "model.", "encoder.", etc.
                clean_key = key
                for prefix in ["model.", "encoder."]:
                    if clean_key.startswith(prefix):
                        clean_key = clean_key[len(prefix) :]

                # Map old keys to new model structure
                # The original model has complex structure, we need to map to our simplified version
                # Skip keys that don't belong to our model (like gloss_layer, sign_loss_function, etc.)
                if any(
                    skip in clean_key
                    for skip in [
                        "gloss_layer",
                        "loss_function",
                        "input_layer",
                        "mixer",
                        "decoder",
                        "gloss_",
                        "sign_weight",
                        "repetition",
                    ]
                ):
                    continue

                cleaned_state_dict[clean_key] = value

            state_dict = cleaned_state_dict

        else:
            state_dict = checkpoint

        # Load state dict (strict=False to allow partial loading)
        missing, unexpected = self.model.load_state_dict(state_dict, strict=False)

        # Move to device and set to eval mode
        self.model.to(self.device)
        self.model.eval()  # CRITICAL: set to eval mode for inference

        # Log loading info
        if missing or unexpected:
            print(f"Model loaded from {checkpoint_path}")
            if missing:
                print(f"  Missing keys: {len(missing)} (expected for simplified model)")
            if unexpected:
                print(f"  Unexpected keys: {len(unexpected)}")

    def segment(
        self,
        wilor_features: np.ndarray,
        fps: float = 25.0,
        min_segment_duration: float = 0.1,
    ) -> SegmentList:
        """Segment a video given WiLoR features.

        Args:
            wilor_features: WiLoR features of shape (T, 192)
            fps: Video frame rate
            min_segment_duration: Minimum segment duration in seconds

        Returns:
            SegmentList with detected sign boundaries

        Raises:
            ConfigurationError: If model not loaded
        """
        if self.model is None:
            raise ConfigurationError("Model not loaded. Call load_model() first.")

        from sltk.segmentation._signsegmentor_v2.inference import inference_windowed
        from sltk.segmentation._signsegmentor_v2.postprocess import label_fixes

        # Convert to tensor
        features = torch.from_numpy(wilor_features).float()

        # Subsample features to match training resolution
        # (model was trained on subsampled data)
        subsample: int = int(self.config["subsample"])
        if subsample > 1:
            features = features[::subsample]

        # Run inference
        predictions = inference_windowed(
            model=self.model,
            features=features,
            window_size=self.config["window_size"],
            subsample=subsample,
            device=self.device,
        )

        # Apply post-processing
        predictions = label_fixes(predictions)

        # Convert to numpy
        predictions = predictions.numpy()

        # Convert predictions to segments
        segments = self._predictions_to_segments(predictions, fps, min_segment_duration, subsample)

        return segments

    def _predictions_to_segments(
        self,
        predictions: np.ndarray,
        fps: float,
        min_duration: float,
        subsample: int,
    ) -> SegmentList:
        """Convert frame-wise predictions to segment list.

        Args:
            predictions: Frame-level class predictions (T,)
            fps: Video frame rate
            min_duration: Minimum segment duration in seconds
            subsample: Subsample factor (for upsampling)

        Returns:
            SegmentList with detected segments
        """
        segments = []

        # Upsample predictions if subsampled
        if subsample > 1:
            predictions = np.repeat(predictions, subsample)

        # Find contiguous active regions (class 1 or 2)
        active_mask = (predictions == IN_SIGN_TOKEN) | (predictions == BEGIN_TOKEN)

        # Handle empty predictions
        if len(active_mask) == 0:
            return SegmentList([])

        # Find boundaries
        changes = np.diff(active_mask.astype(int))
        starts = np.where(changes == 1)[0] + 1
        ends = np.where(changes == -1)[0] + 1

        # Handle edge cases
        if active_mask[0]:
            starts = np.concatenate([[0], starts])
        if active_mask[-1]:
            ends = np.concatenate([ends, [len(active_mask)]])

        # Create segments
        for start_frame, end_frame in zip(starts, ends):
            start_sec = start_frame / fps
            end_sec = end_frame / fps

            # Filter by minimum duration
            if end_sec - start_sec >= min_duration:
                segments.append(
                    Segment(
                        start=start_sec,
                        end=end_sec,
                        label="SIGN",
                        confidence=None,  # Could add confidence from model logits later
                    )
                )

        return SegmentList(segments)
