"""
SignRep inference pipeline for dictionary extraction, continuous features,
and gloss spotting.

Ported from SignRepID into VLT's module structure. Provides three modes:
  1. dictionary  -- isolated sign video -> pooled 768-dim feature
  2. continuous   -- full video -> dense sliding-window features
  3. spot         -- match segments against dictionary features

Usage:
    from sltk.embedding.pipeline import SignRepPipeline

    pipeline = SignRepPipeline()
    result = pipeline.extract_dictionary("sign.mp4", method="middle")
    result = pipeline.extract_continuous("video.mp4", stride=4)
    result = pipeline.spot_from_video(
        "video.mp4", "segments.json", ["dict_feats/"], top_k=20,
    )
"""

from __future__ import annotations

import json
import logging
import re
import threading
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import albumentations as A
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from albumentations.pytorch.transforms import ToTensorV2
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

SEGMENT_LENGTH = 16
FEATURE_DIM = 768
VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}

MODEL_PARAMS = {
    "ckpt_dir": None,
    "backbone_params": {"drop_path_rate": 0.0, "num_cls_tokens": 0},
    "out_dim": FEATURE_DIM,
    "mask_ratio": 0.0,
    "init_weights": False,
    "cls_head_params": {
        "latents_params": {"hidden_dim": 512, "out_dim": 384},
        "total_frames": SEGMENT_LENGTH,
        "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},
        },
    },
}


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass
class DictionaryResult:
    """Result of dictionary feature extraction for a single video."""

    feature: np.ndarray  # (768,) pooled feature vector
    all_features: np.ndarray  # (N_windows, 768) all window features
    active_probs: np.ndarray  # (N_windows,) activity probabilities
    frame_indices: np.ndarray  # (N_windows,) start frame per window
    fps: float
    total_frames: int
    pooling_method: str
    source_path: str

    def save_npz(self, path: str | Path):
        """Save in backward-compatible .npz format."""
        np.savez(
            str(path),
            features=self.all_features,
            best_feature=self.feature,
            best_latent=self.feature,
            active_probs=self.active_probs,
            frame_indices=self.frame_indices,
            fps=np.float32(self.fps),
            total_frames=np.int32(self.total_frames),
            segment_length=np.int32(SEGMENT_LENGTH),
            mode="dictionary",
            pooling_method=self.pooling_method,
        )


@dataclass
class ContinuousResult:
    """Result of continuous feature extraction for a single video."""

    features: np.ndarray  # (N_windows, 768) L2-normalized
    active_probs: np.ndarray  # (N_windows,)
    frame_indices: np.ndarray  # (N_windows,)
    fps: float
    total_frames: int
    stride: int

    def save_npz(self, path: str | Path):
        """Save in backward-compatible .npz format."""
        np.savez(
            str(path),
            features=self.features,
            latents=self.features,
            active_probs=self.active_probs,
            frame_indices=self.frame_indices,
            fps=np.float32(self.fps),
            total_frames=np.int32(self.total_frames),
            stride=np.int32(self.stride),
            segment_length=np.int32(SEGMENT_LENGTH),
            mode="continuous",
        )


@dataclass
class DictionaryIndex:
    """A loaded dictionary ready for retrieval."""

    features: torch.Tensor  # (N_dict, 768) L2-normalized
    glosses: list[str]
    filenames: list[str]
    source_dirs: list[str]
    file_paths: list[str]

    def to(self, device):
        self.features = self.features.to(device)
        return self

    def __len__(self):
        return len(self.glosses)


@dataclass
class SegmentMatch:
    """Top-K matches for a single segment."""

    segment_id: int
    video_id: str
    start_frame: int
    end_frame: int
    start_ms: int
    end_ms: int
    start_timestamp: str
    end_timestamp: str
    num_windows: int
    top_glosses: list[dict]

    def to_dict(self):
        return {
            "video_id": self.video_id,
            "segment_id": self.segment_id,
            "start_frame": self.start_frame,
            "end_frame": self.end_frame,
            "start_ms": self.start_ms,
            "end_ms": self.end_ms,
            "start_timestamp": self.start_timestamp,
            "end_timestamp": self.end_timestamp,
            "num_windows": self.num_windows,
            "top_glosses": self.top_glosses,
        }


@dataclass
class SpottingResult:
    """Complete spotting results."""

    segments: list[SegmentMatch]
    metadata: dict[str, Any]

    def save_json(self, path: str | Path):
        """Save results to JSON."""
        output = {
            "metadata": self.metadata,
            "segments": [s.to_dict() for s in self.segments],
        }
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(output, f, indent=2)

    def to_dict(self):
        return {
            "metadata": self.metadata,
            "segments": [s.to_dict() for s in self.segments],
        }

    def save_eaf(
        self,
        path: str | Path,
        media_path: str | Path | None = None,
    ) -> Path:
        """Save spotting results as an ELAN EAF file with one tier per rank.

        Creates tiers Rank-1 through Rank-K, each containing the gloss name
        as the annotation value, plus matching Score-1 through Score-K tiers.

        Args:
            path: Output EAF file path.
            media_path: Optional video path for MEDIA_DESCRIPTOR.

        Returns:
            Path to the saved EAF file.
        """
        from xml.etree.ElementTree import Element, ElementTree, SubElement
        from xml.etree.ElementTree import indent as xml_indent

        path = Path(path)
        top_k = self.metadata.get("top_k", 0)
        if self.segments and self.segments[0].top_glosses:
            top_k = len(self.segments[0].top_glosses)

        root = Element("ANNOTATION_DOCUMENT")
        root.set("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance")
        root.set("xsi:noNamespaceSchemaLocation", "http://www.mpi.nl/tools/elan/EAFv3.0.xsd")
        root.set("AUTHOR", "sltk_spotter")
        root.set("FORMAT", "3.0")

        header = SubElement(root, "HEADER")
        header.set("MEDIA_FILE", "")
        header.set("TIME_UNITS", "milliseconds")

        if media_path is not None:
            media_path = Path(media_path).resolve()
            md = SubElement(header, "MEDIA_DESCRIPTOR")
            md.set("MEDIA_URL", media_path.as_uri())
            ext = media_path.suffix.lower()
            mime = {
                ".mp4": "video/mp4",
                ".mov": "video/quicktime",
                ".avi": "video/x-msvideo",
                ".mkv": "video/x-matroska",
            }.get(ext, "video/mp4")
            md.set("MIME_TYPE", mime)
            try:
                rel = media_path.relative_to(path.resolve().parent)
                md.set("RELATIVE_MEDIA_URL", f"./{rel}")
            except ValueError:
                pass

        time_order = SubElement(root, "TIME_ORDER")

        # Create tier pairs: Rank-N and Score-N
        gloss_tiers = []
        score_tiers = []
        for rank in range(1, top_k + 1):
            gt = SubElement(root, "TIER")
            gt.set("LINGUISTIC_TYPE_REF", "default-lt")
            gt.set("TIER_ID", f"Rank-{rank}")
            gloss_tiers.append(gt)

            st = SubElement(root, "TIER")
            st.set("LINGUISTIC_TYPE_REF", "default-lt")
            st.set("TIER_ID", f"Score-{rank}")
            score_tiers.append(st)

        ts_id = 1
        for seg in self.segments:
            # One pair of time slots per segment, shared across all rank tiers
            ts_start_id = f"ts{ts_id}"
            ts_start = SubElement(time_order, "TIME_SLOT")
            ts_start.set("TIME_SLOT_ID", ts_start_id)
            ts_start.set("TIME_VALUE", str(seg.start_ms))
            ts_id += 1

            ts_end_id = f"ts{ts_id}"
            ts_end = SubElement(time_order, "TIME_SLOT")
            ts_end.set("TIME_SLOT_ID", ts_end_id)
            ts_end.set("TIME_VALUE", str(seg.end_ms))
            ts_id += 1

            for g in seg.top_glosses:
                rank = g["rank"]
                if rank > top_k:
                    break

                # Gloss tier
                ann = SubElement(gloss_tiers[rank - 1], "ANNOTATION")
                aa = SubElement(ann, "ALIGNABLE_ANNOTATION")
                aa.set("ANNOTATION_ID", f"a_r{rank}_g_{seg.segment_id}")
                aa.set("TIME_SLOT_REF1", ts_start_id)
                aa.set("TIME_SLOT_REF2", ts_end_id)
                val = SubElement(aa, "ANNOTATION_VALUE")
                val.text = g["gloss"]

                # Score tier
                ann2 = SubElement(score_tiers[rank - 1], "ANNOTATION")
                aa2 = SubElement(ann2, "ALIGNABLE_ANNOTATION")
                aa2.set("ANNOTATION_ID", f"a_r{rank}_s_{seg.segment_id}")
                aa2.set("TIME_SLOT_REF1", ts_start_id)
                aa2.set("TIME_SLOT_REF2", ts_end_id)
                val2 = SubElement(aa2, "ANNOTATION_VALUE")
                val2.text = f"{g['similarity']:.4f}"

        lt = SubElement(root, "LINGUISTIC_TYPE")
        lt.set("LINGUISTIC_TYPE_ID", "default-lt")
        lt.set("TIME_ALIGNABLE", "true")
        lt.set("GRAPHIC_REFERENCES", "false")

        xml_indent(root)
        tree = ElementTree(root)
        path.parent.mkdir(parents=True, exist_ok=True)
        tree.write(path, encoding="unicode", xml_declaration=True)
        return path


# ---------------------------------------------------------------------------
# Frame transform
# ---------------------------------------------------------------------------


class FrameTransform:
    """Per-frame resize + ImageNet normalization."""

    def __init__(self):
        self.transform = A.Compose(
            [
                A.Resize(224, 224, p=1.0),
                A.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                    max_pixel_value=255.0,
                    p=1.0,
                ),
                ToTensorV2(p=1.0),
            ]
        )

    def __call__(self, frame):
        return self.transform(image=frame)["image"]


# ---------------------------------------------------------------------------
# Video segment datasets
# ---------------------------------------------------------------------------


class VideoSegmentDataset(Dataset):
    """Yields (C, T, H, W) segments via frame seeking. Memory-efficient."""

    def __init__(
        self,
        video_path,
        stride,
        segment_length=SEGMENT_LENGTH,
        transform=None,
        pad_short=True,
    ):
        self.video_path = str(video_path)
        self.stride = stride
        self.segment_length = segment_length
        self.transform = transform or FrameTransform()
        self.pad_short = pad_short

        cap = cv2.VideoCapture(self.video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Cannot open video: {video_path}")
        self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
        cap.release()

        self.start_indices = self._compute_start_indices()

    def _compute_start_indices(self):
        if self.total_frames >= self.segment_length:
            indices = list(range(0, self.total_frames - self.segment_length + 1, self.stride))
            last_valid = self.total_frames - self.segment_length
            if indices and indices[-1] != last_valid:
                indices.append(last_valid)
            return indices
        elif self.pad_short and self.total_frames > 0:
            return [0]
        return []

    def __len__(self):
        return len(self.start_indices)

    def __getitem__(self, idx):
        start = self.start_indices[idx]
        cap = cv2.VideoCapture(self.video_path)
        cap.set(cv2.CAP_PROP_POS_FRAMES, start)

        frames = []
        for _ in range(self.segment_length):
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(self.transform(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
        cap.release()

        if len(frames) == 0:
            raise RuntimeError(f"No frames read from {self.video_path} at pos {start}")

        while len(frames) < self.segment_length:
            frames.append(frames[-1].clone())

        segment = torch.stack(frames).permute(1, 0, 2, 3)  # (C, T, H, W)
        return segment, start


class SegmentCropDataset(Dataset):
    """Reads the middle SEGMENT_LENGTH frames of each provided segment.

    For segments shorter than SEGMENT_LENGTH, pads by repeating the last frame.
    For segments longer, takes the centered SEGMENT_LENGTH window.
    """

    def __init__(self, video_path, segments, segment_length=SEGMENT_LENGTH, transform=None):
        self.video_path = str(video_path)
        self.segment_length = segment_length
        self.transform = transform or FrameTransform()
        self.segments = segments

        cap = cv2.VideoCapture(self.video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Cannot open video: {video_path}")
        self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
        cap.release()

    def __len__(self):
        return len(self.segments)

    def __getitem__(self, idx):
        seg = self.segments[idx]
        seg_start = seg["start_frame"]
        seg_end = seg["end_frame"]
        seg_len = seg_end - seg_start + 1

        # Always center a full segment_length window on the segment midpoint,
        # even when the segment is shorter than segment_length.
        mid = seg_start + seg_len // 2
        start = mid - self.segment_length // 2
        # Clamp to video bounds
        start = max(0, min(start, self.total_frames - self.segment_length))

        cap = cv2.VideoCapture(self.video_path)
        cap.set(cv2.CAP_PROP_POS_FRAMES, start)

        frames = []
        for _ in range(self.segment_length):
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(self.transform(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
        cap.release()

        if len(frames) == 0:
            raise RuntimeError(
                f"No frames read from {self.video_path} at pos {start} for segment {idx}"
            )

        # Pad to SEGMENT_LENGTH by repeating last frame
        while len(frames) < self.segment_length:
            frames.append(frames[-1].clone())

        segment = torch.stack(frames).permute(1, 0, 2, 3)  # (C, T, H, W)
        return segment, idx


class PreloadedVideoSegmentDataset(Dataset):
    """Preloads all frames then yields segments. Reliable seeking."""

    def __init__(
        self,
        video_path,
        stride,
        segment_length=SEGMENT_LENGTH,
        transform=None,
        pad_short=True,
    ):
        self.segment_length = segment_length
        self.transform = transform or FrameTransform()

        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            raise RuntimeError(f"Cannot open video: {video_path}")
        self.fps = cap.get(cv2.CAP_PROP_FPS) or 25.0

        self.frames = []
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            self.frames.append(self.transform(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
        cap.release()

        self.total_frames = len(self.frames)
        if self.total_frames == 0:
            raise RuntimeError(f"No frames read from {video_path}")

        while len(self.frames) < segment_length and pad_short:
            self.frames.append(self.frames[-1].clone())

        effective_len = len(self.frames)
        if effective_len >= segment_length:
            self.start_indices = list(range(0, effective_len - segment_length + 1, stride))
            last_valid = effective_len - segment_length
            if self.start_indices and self.start_indices[-1] != last_valid:
                self.start_indices.append(last_valid)
        else:
            self.start_indices = []

    def __len__(self):
        return len(self.start_indices)

    def __getitem__(self, idx):
        start = self.start_indices[idx]
        segment = torch.stack(self.frames[start : start + self.segment_length])
        segment = segment.permute(1, 0, 2, 3)
        return segment, start


# ---------------------------------------------------------------------------
# Dictionary pooling strategies
# ---------------------------------------------------------------------------


def pool_middle(features, active_probs=None, **kwargs):
    """Take the feature from the center window."""
    mid_idx = len(features) // 2
    return features[mid_idx]


def pool_max(features, active_probs=None, **kwargs):
    """Element-wise max across all windows."""
    return features.max(dim=0).values


def pool_mean(features, active_probs=None, **kwargs):
    """Simple average across all windows."""
    return features.mean(dim=0)


POOLING_METHODS = {
    "middle": pool_middle,
    "max_pool": pool_max,
    "mean_pool": pool_mean,
}


# ---------------------------------------------------------------------------
# Segment pooling strategies (for spotting)
# ---------------------------------------------------------------------------


def segment_pool_max(sims):
    """Max similarity across windows per dictionary entry."""
    return sims.max(dim=0).values


def segment_pool_mean(sims):
    """Mean similarity across windows per dictionary entry."""
    return sims.mean(dim=0)


def segment_pool_softmax(sims, temperature=0.1):
    """Softmax-weighted similarity across windows."""
    weights = torch.softmax(sims / temperature, dim=0)
    return (sims * weights).sum(dim=0)


SEGMENT_POOLING = {
    "max": segment_pool_max,
    "mean": segment_pool_mean,
    "softmax_weighted": segment_pool_softmax,
}


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


def parse_gloss_name(filename_stem):
    """Strip numeric ID suffix: ADDRESS2-5675 -> ADDRESS2."""
    match = re.match(r"^(.+)-\d+$", filename_stem)
    return match.group(1) if match else filename_stem


def find_overlapping_windows(frame_indices, segment_length, seg_start, seg_end):
    """Find window indices overlapping [seg_start, seg_end). Vectorized."""
    starts = frame_indices
    ends = starts + segment_length
    mask = (starts < seg_end) & (ends > seg_start)
    return np.where(mask)[0]


def frame_to_ms(frame_idx, fps):
    """Convert frame index to milliseconds."""
    return int(round(frame_idx / fps * 1000))


def load_segments(json_path):
    """Load segments JSON grouped by original_video_id."""
    with open(json_path) as f:
        all_segments = json.load(f)

    video_segments = defaultdict(list)
    for seg in all_segments:
        video_segments[seg["original_video_id"]].append(seg)

    for vid in video_segments:
        video_segments[vid].sort(key=lambda s: s["segment_id"])

    return video_segments


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------


def _remap_classifier_checkpoint(state_dict):
    """Remap classifier checkpoint keys to FINAL model structure."""
    remapped = {}
    skipped = []
    for k, v in state_dict.items():
        if k.startswith("backbone.backbone."):
            remapped[k[len("backbone.") :]] = v
        elif k.startswith("backbone."):
            remapped[k[len("backbone.") :]] = v
        else:
            skipped.append(k)
    if skipped:
        logger.info("Skipped %d classifier-specific keys", len(skipped))
    return remapped


def load_model(checkpoint_path, device):
    """Load model with auto-detected checkpoint format."""
    from sltk.embedding.signrep_models.final_model import Model

    logger.info("Initializing SignRep model...")
    model = Model(**MODEL_PARAMS)

    logger.info("Loading checkpoint: %s", checkpoint_path)
    try:
        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    except Exception:
        logger.warning("weights_only=True failed, retrying with weights_only=False")
        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)

    if isinstance(checkpoint, dict) and "model" in checkpoint:
        state_dict = checkpoint["model"]
        fmt = "dict with 'model' key"
    elif isinstance(checkpoint, dict) and "model_state" in checkpoint:
        state_dict = checkpoint["model_state"]
        fmt = "dict with 'model_state' key"
    else:
        state_dict = checkpoint
        fmt = "flat state dict"

    sample_keys = list(state_dict.keys())[:5]
    is_classifier_ckpt = any(k.startswith("backbone.backbone.") for k in sample_keys)

    if is_classifier_ckpt:
        logger.info("Checkpoint format: %s (classifier model -- remapping keys)", fmt)
        state_dict = _remap_classifier_checkpoint(state_dict)
        has_trained_heads = any(
            k.startswith(("fc.", "active_head.", "cls_head.")) for k in state_dict
        )
    else:
        logger.info("Checkpoint format: %s", fmt)
        has_trained_heads = any(k.startswith(("fc.", "active_head.")) for k in state_dict)

    result = model.load_state_dict(state_dict, strict=False)
    if result.missing_keys:
        logger.info(
            "Missing keys (%d): %s...",
            len(result.missing_keys),
            result.missing_keys[:5],
        )
    if result.unexpected_keys:
        logger.info(
            "Unexpected keys (%d): %s...",
            len(result.unexpected_keys),
            result.unexpected_keys[:5],
        )

    if not has_trained_heads:
        logger.warning("Trained heads not found. Dictionary mode will use center-window heuristic.")

    param_count = sum(p.numel() for p in model.parameters())
    logger.info("Parameters: %s", f"{param_count:,}")

    model = model.to(device)
    model.eval()
    model._has_trained_heads = has_trained_heads
    logger.info("Model ready on %s", device)
    return model


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------


class SignRepPipeline:
    """Unified inference pipeline for SignRep.

    Three modes:
      1. extract_dictionary() -- isolated sign -> single pooled feature
      2. extract_continuous() -- full video -> dense features
      3. spot() / spot_from_video() -- match against dictionary

    Example:
        pipeline = SignRepPipeline()
        result = pipeline.extract_dictionary("sign.mp4", method="middle")
        result = pipeline.extract_continuous("video.mp4", stride=4)
        result = pipeline.spot_from_video(
            "video.mp4", "segments.json", ["dict_feats/"], top_k=20,
        )
    """

    def __init__(
        self,
        checkpoint: str | Path | None = None,
        device: str = "cuda",
        amp: bool = True,
        batch_size: int = 8,
        num_workers: int = 4,
        preload: bool = False,
        active_index: int = 1,
    ):
        if device == "cuda" and not torch.cuda.is_available():
            logger.warning("CUDA not available, falling back to CPU")
            device = "cpu"

        self.device = torch.device(device)
        self.amp = amp and (device != "cpu")
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.preload = preload
        self.active_index = active_index

        # Resolve checkpoint path
        if checkpoint is None:
            from sltk.extraction.weights import get_weight_path

            checkpoint = get_weight_path("signrep_checkpoint")
            if checkpoint is None:
                raise FileNotFoundError(
                    "SignRep checkpoint not found. Place ckpt.pt at "
                    "sltk/weights/signrep/ckpt.pt or set SLTK_SIGNREP_CHECKPOINT."
                )

        self.model = load_model(str(checkpoint), self.device)

    # ----- core extraction loop -----

    def _extract_batched(self, video_path, stride, need_active=True):
        """Core batched extraction. Returns tensors on GPU."""
        DatasetClass = PreloadedVideoSegmentDataset if self.preload else VideoSegmentDataset
        dataset = DatasetClass(video_path=str(video_path), stride=stride)

        if len(dataset) == 0:
            return None

        effective_workers = 0 if self.preload else self.num_workers
        dataloader = DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=effective_workers,
            pin_memory=(str(self.device) != "cpu"),
            drop_last=False,
            persistent_workers=effective_workers > 0,
        )

        all_features = []
        all_active_probs = []
        all_frame_indices = []

        with torch.inference_mode():
            for batch_segments, batch_starts in dataloader:
                batch_segments = batch_segments.to(self.device, non_blocking=True)

                if self.amp:
                    with torch.amp.autocast("cuda"):
                        outputs = self.model.inference_forward(
                            batch_segments, need_active=need_active
                        )
                else:
                    outputs = self.model.inference_forward(batch_segments, need_active=need_active)

                all_features.append(outputs["latent"])
                if need_active and "cls_active" in outputs:
                    probs = torch.softmax(outputs["cls_active"], dim=-1)[:, self.active_index]
                    all_active_probs.append(probs)
                all_frame_indices.append(batch_starts)

        features = torch.cat(all_features, dim=0).float()  # ensure fp32 after AMP
        frame_indices = torch.cat(all_frame_indices, dim=0)
        active_probs = torch.cat(all_active_probs, dim=0).float() if all_active_probs else None

        return features, active_probs, frame_indices, dataset.fps, dataset.total_frames

    # ----- Mode 1: dictionary -----

    def extract_dictionary(
        self,
        video_path: str | Path,
        method: str = "middle",
        stride: int = 1,
    ) -> DictionaryResult:
        """Extract a representative feature from an isolated sign video.

        Args:
            video_path: Path to video file.
            method: Pooling strategy. One of: middle, max_pool, mean_pool.
            stride: Frame stride between windows.

        Returns:
            DictionaryResult with the pooled feature and metadata.
        """
        result = self._extract_batched(video_path, stride, need_active=True)
        if result is None:
            raise RuntimeError(f"No valid segments in {video_path}")

        features_gpu, active_probs_gpu, frame_indices, fps, total_frames = result

        pool_fn = POOLING_METHODS.get(method)
        if pool_fn is None:
            raise ValueError(
                f"Unknown pooling method '{method}'. "
                f"Choose from: {list(POOLING_METHODS.keys())}"
            )

        pooled = pool_fn(features_gpu, active_probs=active_probs_gpu)

        # L2-normalize after pooling
        pooled = F.normalize(pooled.unsqueeze(0), dim=1).squeeze(0)

        return DictionaryResult(
            feature=pooled.cpu().numpy(),
            all_features=features_gpu.cpu().numpy(),
            active_probs=(
                active_probs_gpu.cpu().numpy()
                if active_probs_gpu is not None
                else np.zeros(len(features_gpu))
            ),
            frame_indices=frame_indices.cpu().numpy(),
            fps=fps,
            total_frames=total_frames,
            pooling_method=method,
            source_path=str(video_path),
        )

    def extract_dictionary_batch(
        self,
        input_dir: str | Path,
        output_dir: str | Path | None = None,
        method: str = "middle",
        stride: int = 1,
        resume: bool = False,
    ) -> dict[str, DictionaryResult]:
        """Batch extract dictionary features from a directory of videos.

        Returns dict mapping filename_stem -> DictionaryResult.
        Optionally saves .npz files to output_dir.
        """
        input_dir = Path(input_dir)
        video_paths = sorted(
            p for p in input_dir.rglob("*") if p.suffix.lower() in VIDEO_EXTENSIONS and p.is_file()
        )
        logger.info("Found %d videos in %s", len(video_paths), input_dir)

        if output_dir is not None:
            output_dir = Path(output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)

        results = {}
        for video_path in tqdm(video_paths, desc="Dictionary extraction"):
            stem = video_path.stem

            if resume and output_dir is not None:
                out_path = output_dir / video_path.relative_to(input_dir).with_suffix(".npz")
                if out_path.exists():
                    continue

            try:
                result = self.extract_dictionary(video_path, method=method, stride=stride)
                results[stem] = result

                if output_dir is not None:
                    out_path = output_dir / video_path.relative_to(input_dir).with_suffix(".npz")
                    out_path.parent.mkdir(parents=True, exist_ok=True)
                    result.save_npz(out_path)

            except Exception as e:
                logger.warning("Failed to extract %s: %s", video_path.name, e)

        logger.info("Processed %d videos", len(results))
        return results

    # ----- Mode 2: continuous -----

    def extract_continuous(
        self,
        video_path: str | Path,
        stride: int = 4,
    ) -> ContinuousResult:
        """Extract features from all sliding windows of a continuous video.

        Args:
            video_path: Path to video file.
            stride: Frame stride between windows (default: 4).

        Returns:
            ContinuousResult with L2-normalized features.
        """
        result = self._extract_batched(video_path, stride, need_active=True)
        if result is None:
            raise RuntimeError(f"No valid segments in {video_path}")

        features_gpu, active_probs_gpu, frame_indices, fps, total_frames = result

        # L2-normalize on GPU
        features_gpu = F.normalize(features_gpu, dim=1)

        return ContinuousResult(
            features=features_gpu.cpu().numpy(),
            active_probs=(
                active_probs_gpu.cpu().numpy()
                if active_probs_gpu is not None
                else np.zeros(len(features_gpu))
            ),
            frame_indices=frame_indices.cpu().numpy(),
            fps=fps,
            total_frames=total_frames,
            stride=stride,
        )

    # ----- Dictionary loading -----

    def load_dictionary(
        self,
        paths: Sequence[str | Path],
        feature_key: str = "best_latent",
        deduplicate: bool = True,
    ) -> DictionaryIndex:
        """Load dictionary features from .pt files, .npz directories, or a mix.

        Auto-detects format per path:
          - File ending in .pt  -> consolidated tensor (fast)
          - Directory            -> walks for .npz files (flexible)

        Multiple paths are merged into a single index.

        Args:
            paths: List of .pt file paths and/or directories of .npz files.
            feature_key: Key to read from .npz files (default: best_latent).
            deduplicate: If True, keep only one entry per unique gloss name
                (the first occurrence). Prevents duplicate matches in results.

        Returns:
            DictionaryIndex ready for spotting.
        """
        all_features = []
        glosses = []
        filenames = []
        source_dirs = []
        file_paths = []

        for path in paths:
            p = Path(path)

            if p.is_file() and p.suffix == ".pt":
                data = torch.load(p, map_location="cpu", weights_only=True)
                feats = data["features"].float()
                names = data["filenames"]
                src = data.get("source_dir", p.stem)

                all_features.append(feats)
                for name in names:
                    filenames.append(name)
                    glosses.append(parse_gloss_name(name))
                    source_dirs.append(src)
                    file_paths.append(f"{src}/{name}.npz")

                logger.info("Loaded %s: %d entries (.pt)", p.name, feats.shape[0])

            elif p.is_dir():
                npz_files = sorted(p.rglob("*.npz"))
                if len(npz_files) == 0:
                    logger.warning("No .npz files found in %s", p)
                    continue

                dir_features = []
                for npz_path in tqdm(npz_files, desc=f"Loading dictionary ({p.name})"):
                    data = np.load(npz_path, allow_pickle=True)

                    feat = None
                    for key in [
                        feature_key,
                        "best_latent",
                        "best_feature",
                        "features",
                    ]:
                        if key in data:
                            feat = data[key].astype(np.float32)
                            break
                    if feat is None:
                        logger.warning("Skip %s: no feature key found", npz_path.name)
                        continue

                    if feat.ndim > 1:
                        feat = feat.mean(axis=0)

                    dir_features.append(feat)
                    filenames.append(npz_path.stem)
                    glosses.append(parse_gloss_name(npz_path.stem))
                    source_dirs.append(p.name)
                    file_paths.append(str(npz_path))

                if dir_features:
                    all_features.append(torch.from_numpy(np.stack(dir_features)).float())
                    logger.info("Loaded %s: %d entries (.npz)", p.name, len(dir_features))

            else:
                logger.warning("Skipping %s (not a .pt file or directory)", p)

        if len(all_features) == 0:
            raise FileNotFoundError(f"No valid dictionary entries in: {[str(p) for p in paths]}")

        features = torch.cat(all_features, dim=0)
        features = F.normalize(features, dim=1)

        total_before = len(glosses)

        if deduplicate:
            seen = {}
            keep = []
            for i, gloss in enumerate(glosses):
                if gloss not in seen:
                    seen[gloss] = i
                    keep.append(i)
            if len(keep) < total_before:
                keep_t = torch.tensor(keep, dtype=torch.long)
                features = features[keep_t]
                glosses = [glosses[i] for i in keep]
                filenames = [filenames[i] for i in keep]
                source_dirs = [source_dirs[i] for i in keep]
                file_paths = [file_paths[i] for i in keep]
                logger.info(
                    "Deduplicated: %d -> %d unique glosses",
                    total_before,
                    len(glosses),
                )

        logger.info(
            "Dictionary loaded: %d entries from %d source(s), %d-dim",
            len(glosses),
            len(paths),
            features.shape[1],
        )

        return DictionaryIndex(
            features=features.to(self.device),
            glosses=glosses,
            filenames=filenames,
            source_dirs=source_dirs,
            file_paths=file_paths,
        )

    # ----- Mode 3: spotting -----

    def _match_segments(
        self,
        all_sims,
        frame_indices_np,
        fps,
        video_id,
        segments,
        dictionary,
        pool_fn,
        top_k,
    ):
        """Match segments against dictionary using a pre-computed similarity matrix."""
        results = []
        for seg in segments:
            seg_start = seg["start_frame"]
            seg_end = seg["end_frame"]

            win_idx = find_overlapping_windows(frame_indices_np, SEGMENT_LENGTH, seg_start, seg_end)

            if len(win_idx) == 0:
                results.append(
                    SegmentMatch(
                        segment_id=seg["segment_id"],
                        video_id=video_id,
                        start_frame=seg_start,
                        end_frame=seg_end,
                        start_ms=frame_to_ms(seg_start, fps),
                        end_ms=frame_to_ms(seg_end, fps),
                        start_timestamp=seg.get("start_timestamp", ""),
                        end_timestamp=seg.get("end_timestamp", ""),
                        num_windows=0,
                        top_glosses=[],
                    )
                )
                continue

            seg_sims = all_sims[win_idx]  # (N_overlap, N_dict)
            pooled = pool_fn(seg_sims)  # (N_dict,)

            k = min(top_k, pooled.shape[0])
            top_scores, top_indices = torch.topk(pooled, k)

            top_scores_cpu = top_scores.cpu()
            top_indices_cpu = top_indices.cpu()
            top_glosses = []
            for r in range(k):
                idx = int(top_indices_cpu[r])
                top_glosses.append(
                    {
                        "rank": r + 1,
                        "gloss": dictionary.glosses[idx],
                        "similarity": round(top_scores_cpu[r].item(), 4),
                        "dict_file": dictionary.filenames[idx],
                        "file_path": dictionary.file_paths[idx],
                        "source_dir": dictionary.source_dirs[idx],
                    }
                )

            results.append(
                SegmentMatch(
                    segment_id=seg["segment_id"],
                    video_id=video_id,
                    start_frame=seg_start,
                    end_frame=seg_end,
                    start_ms=frame_to_ms(seg_start, fps),
                    end_ms=frame_to_ms(seg_end, fps),
                    start_timestamp=seg.get("start_timestamp", ""),
                    end_timestamp=seg.get("end_timestamp", ""),
                    num_windows=len(win_idx),
                    top_glosses=top_glosses,
                )
            )

        return results

    @staticmethod
    def _build_metadata(dictionary, top_k, segment_pooling):
        """Build metadata dict for spotting results."""
        return {
            "dictionary_size": len(dictionary),
            "dictionary_dirs": list(dict.fromkeys(dictionary.source_dirs)),
            "top_k": top_k,
            "segment_pooling": segment_pooling,
            "feature_dim": FEATURE_DIM,
        }

    def spot(
        self,
        features: ContinuousResult | str | Path | np.ndarray,
        segments: str | Path | list[dict],
        dictionary: DictionaryIndex,
        top_k: int = 20,
        segment_pooling: str = "max",
    ) -> SpottingResult:
        """Match continuous features against dictionary for each segment.

        Args:
            features: ContinuousResult, path to .npz, or (N, 768) array.
            segments: Path to segments JSON, or list of segment dicts for
                      a single video.
            dictionary: Loaded DictionaryIndex.
            top_k: Number of top matches per segment.
            segment_pooling: How to aggregate across windows: max, mean,
                             softmax_weighted.

        Returns:
            SpottingResult with matches per segment.
        """
        # Resolve features
        if isinstance(features, ContinuousResult):
            feats_np = features.features
            frame_indices_np = features.frame_indices
            fps = features.fps
        elif isinstance(features, (str, Path)):
            data = np.load(str(features), allow_pickle=True)
            for key in ["latents", "features"]:
                if key in data:
                    feats_np = data[key].astype(np.float32)
                    break
            frame_indices_np = data["frame_indices"].astype(np.int64)
            fps = float(data.get("fps", 25.0))
        elif isinstance(features, np.ndarray):
            feats_np = features
            frame_indices_np = np.arange(len(features)) * 4
            fps = 25.0
        else:
            raise TypeError(f"Unsupported features type: {type(features)}")

        # Move features to GPU, ensure fp32, and L2-normalize
        feats_gpu = torch.from_numpy(feats_np).float().to(self.device)
        feats_gpu = F.normalize(feats_gpu, dim=1)

        # Resolve segments
        if isinstance(segments, (str, Path)):
            all_video_segments = load_segments(segments)
        elif isinstance(segments, list):
            if len(segments) > 0 and "original_video_id" in segments[0]:
                vid = segments[0]["original_video_id"]
                all_video_segments = {vid: segments}
            else:
                all_video_segments = {"unknown": segments}
        else:
            raise TypeError(f"Unsupported segments type: {type(segments)}")

        # Compute full similarity matrix once (fused)
        with torch.inference_mode():
            all_sims = feats_gpu @ dictionary.features.T  # (N_windows, N_dict)

        pool_fn = SEGMENT_POOLING.get(segment_pooling)
        if pool_fn is None:
            raise ValueError(
                f"Unknown segment pooling '{segment_pooling}'. "
                f"Choose from: {list(SEGMENT_POOLING.keys())}"
            )

        result_segments = []
        for video_id, segs in sorted(all_video_segments.items()):
            result_segments.extend(
                self._match_segments(
                    all_sims,
                    frame_indices_np,
                    fps,
                    video_id,
                    segs,
                    dictionary,
                    pool_fn,
                    top_k,
                )
            )

        metadata = self._build_metadata(dictionary, top_k, segment_pooling)
        return SpottingResult(segments=result_segments, metadata=metadata)

    def spot_segments(
        self,
        video_path: str | Path,
        segments: list[dict],
        dictionary: DictionaryIndex,
        top_k: int = 20,
        video_id: str = "unknown",
    ) -> SpottingResult:
        """Spot by extracting features from the middle 16 frames of each segment.

        Instead of continuous sliding-window extraction over the whole video,
        this reads only the centered SEGMENT_LENGTH frames per segment, extracts
        one feature vector each, and matches directly against the dictionary.

        Args:
            video_path: Path to the video file.
            segments: List of segment dicts with start_frame/end_frame keys.
            dictionary: Loaded DictionaryIndex.
            top_k: Number of top matches per segment.
            video_id: Identifier for the video in results.

        Returns:
            SpottingResult with matches per segment.
        """
        video_path = Path(video_path)
        dataset = SegmentCropDataset(
            video_path=str(video_path),
            segments=segments,
            segment_length=SEGMENT_LENGTH,
        )
        fps = dataset.fps

        effective_workers = 0
        dataloader = DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=effective_workers,
            pin_memory=(str(self.device) != "cpu"),
            drop_last=False,
        )

        # Extract one feature per segment
        all_features = []
        all_seg_indices = []
        with torch.inference_mode():
            for batch_segments, batch_indices in dataloader:
                batch_segments = batch_segments.to(self.device, non_blocking=True)
                if self.amp:
                    with torch.amp.autocast("cuda"):
                        outputs = self.model.inference_forward(batch_segments, need_active=False)
                else:
                    outputs = self.model.inference_forward(batch_segments, need_active=False)
                all_features.append(outputs["latent"])
                all_seg_indices.append(batch_indices)

        features = torch.cat(all_features, dim=0).float()
        seg_indices = torch.cat(all_seg_indices, dim=0)

        # L2-normalize and compute similarities
        features = F.normalize(features, dim=1)
        with torch.inference_mode():
            sims = features @ dictionary.features.T  # (N_segments, N_dict)

        # Build results per segment
        result_segments = []
        for i in range(len(segments)):
            seg = segments[seg_indices[i]]
            seg_sims = sims[i]  # (N_dict,)

            k = min(top_k, seg_sims.shape[0])
            top_scores, top_indices = torch.topk(seg_sims, k)

            top_scores_cpu = top_scores.cpu()
            top_indices_cpu = top_indices.cpu()
            top_glosses = []
            for r in range(k):
                idx = int(top_indices_cpu[r])
                top_glosses.append(
                    {
                        "rank": r + 1,
                        "gloss": dictionary.glosses[idx],
                        "similarity": round(top_scores_cpu[r].item(), 4),
                        "dict_file": dictionary.filenames[idx],
                        "file_path": dictionary.file_paths[idx],
                        "source_dir": dictionary.source_dirs[idx],
                    }
                )

            result_segments.append(
                SegmentMatch(
                    segment_id=seg.get("segment_id", i),
                    video_id=video_id,
                    start_frame=seg["start_frame"],
                    end_frame=seg["end_frame"],
                    start_ms=frame_to_ms(seg["start_frame"], fps),
                    end_ms=frame_to_ms(seg["end_frame"], fps),
                    start_timestamp=seg.get("start_timestamp", ""),
                    end_timestamp=seg.get("end_timestamp", ""),
                    num_windows=1,
                    top_glosses=top_glosses,
                )
            )

        metadata = self._build_metadata(dictionary, top_k, "segment_middle")
        metadata["extraction_method"] = "middle_frames"
        metadata["segment_length"] = SEGMENT_LENGTH
        return SpottingResult(segments=result_segments, metadata=metadata)

    def spot_from_video(
        self,
        video_path: str | Path,
        segments_json: str | Path,
        dictionary_dirs: list[str | Path],
        top_k: int = 20,
        stride: int = 4,
        segment_pooling: str = "max",
        dict_feature_key: str = "best_latent",
    ) -> SpottingResult:
        """End-to-end: extract features then spot against dictionary.

        Keeps features on GPU throughout -- no intermediate disk I/O.

        Args:
            video_path: Path to continuous sign video.
            segments_json: Path to segments JSON.
            dictionary_dirs: List of .pt files and/or .npz directories.
            top_k: Number of top matches per segment.
            stride: Frame stride for continuous extraction.
            segment_pooling: Segment-level pooling strategy.
            dict_feature_key: Key in dictionary .npz files.

        Returns:
            SpottingResult with matches per segment.
        """
        dictionary = self.load_dictionary(dictionary_dirs, feature_key=dict_feature_key)
        continuous = self.extract_continuous(video_path, stride=stride)

        all_video_segments = load_segments(segments_json)
        video_stem = Path(video_path).stem
        if video_stem in all_video_segments:
            segments = all_video_segments[video_stem]
        else:
            all_segs = []
            for segs in all_video_segments.values():
                all_segs.extend(segs)
            segments = all_segs

        return self.spot(
            features=continuous,
            segments=segments,
            dictionary=dictionary,
            top_k=top_k,
            segment_pooling=segment_pooling,
        )

    def spot_batch(
        self,
        features_dir: str | Path,
        segments_json: str | Path,
        dictionary_dirs: list[str | Path],
        output_path: str | Path,
        top_k: int = 20,
        segment_pooling: str = "max",
        feature_key: str = "latents",
        dict_feature_key: str = "best_latent",
    ) -> SpottingResult:
        """Batch spotting over pre-extracted features.

        Loads pre-extracted continuous features from features_dir and
        matches against the dictionary for all segments.
        """
        dictionary = self.load_dictionary(dictionary_dirs, feature_key=dict_feature_key)
        all_video_segments = load_segments(segments_json)

        features_dir = Path(features_dir)
        all_segment_matches = []

        pool_fn = SEGMENT_POOLING.get(segment_pooling)
        if pool_fn is None:
            raise ValueError(
                f"Unknown segment pooling '{segment_pooling}'. "
                f"Choose from: {list(SEGMENT_POOLING.keys())}"
            )

        for video_id in sorted(all_video_segments.keys()):
            npz_path = features_dir / f"{video_id}.npz"
            if not npz_path.exists():
                logger.warning("%s not found, skipping %s", npz_path, video_id)
                continue

            data = np.load(npz_path, allow_pickle=True)
            for key in [feature_key, "latents", "features"]:
                if key in data:
                    feats_np = data[key].astype(np.float32)
                    break
            frame_indices_np = data["frame_indices"].astype(np.int64)
            fps = float(data.get("fps", 25.0))

            feats_gpu = torch.from_numpy(feats_np).float().to(self.device)
            feats_gpu = F.normalize(feats_gpu, dim=1)

            with torch.inference_mode():
                all_sims = feats_gpu @ dictionary.features.T

            all_segment_matches.extend(
                self._match_segments(
                    all_sims,
                    frame_indices_np,
                    fps,
                    video_id,
                    all_video_segments[video_id],
                    dictionary,
                    pool_fn,
                    top_k,
                )
            )

            logger.info(
                "%s: %d segments processed",
                video_id,
                len(all_video_segments[video_id]),
            )

        metadata = self._build_metadata(dictionary, top_k, segment_pooling)
        result = SpottingResult(segments=all_segment_matches, metadata=metadata)
        result.save_json(output_path)
        logger.info("Results saved to %s", output_path)
        return result


# ---------------------------------------------------------------------------
# Thread-safe singleton
# ---------------------------------------------------------------------------

_pipeline: SignRepPipeline | None = None
_pipeline_lock = threading.Lock()
_pipeline_error: str | None = None


def get_pipeline(device: str = "cuda", **kwargs) -> SignRepPipeline:
    """Thread-safe lazy singleton for the SignRep pipeline."""
    global _pipeline, _pipeline_error

    if _pipeline is not None:
        return _pipeline

    with _pipeline_lock:
        if _pipeline is not None:
            return _pipeline
        if _pipeline_error is not None:
            raise RuntimeError(f"Pipeline init failed: {_pipeline_error}")
        try:
            _pipeline = SignRepPipeline(device=device, **kwargs)
            return _pipeline
        except Exception as e:
            _pipeline_error = str(e)
            raise
