"""
Cross-dataset gloss comparison tools for sign language research.

Enables finding and comparing glosses across multiple datasets to study
how signs are produced in different contexts (e.g., citation forms in
dictionaries vs. naturalistic production in corpora).

Example:
    >>> # Find "HELLO" across BSL datasets
    >>> result = find_gloss_across_datasets(
    ...     "HELLO",
    ...     datasets=["bslcp", "bobsl"],
    ... )
    >>> print(f"Found {result.total_matches} matches")

    >>> # Compare vocabularies across datasets
    >>> comparison = compare_dataset_vocabularies(["bslcp", "bobsl", "asldict"])
    >>> print(f"Shared glosses: {len(comparison.shared_glosses)}")

    >>> # Find parallel examples for comparison
    >>> examples = find_parallel_examples(
    ...     glosses=["HELLO", "THANK-YOU"],
    ...     datasets=["bslcp", "bobsl"],
    ... )
"""

from __future__ import annotations

import logging
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np

if TYPE_CHECKING:
    from sltk.data.datasets.base import BaseDataset

logger = logging.getLogger(__name__)


# =============================================================================
# Data Classes
# =============================================================================


@dataclass
class DurationStats:
    """Statistics about gloss durations."""

    count: int
    mean: float
    std: float
    min: float
    max: float
    median: float
    quartiles: tuple[float, float, float]  # Q1, Q2, Q3

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "count": self.count,
            "mean": self.mean,
            "std": self.std,
            "min": self.min,
            "max": self.max,
            "median": self.median,
            "quartiles": list(self.quartiles),
        }


@dataclass
class GlossMatch:
    """A match for a gloss in a specific dataset."""

    dataset: str
    sample_id: str
    segment_idx: int
    start: float
    end: float
    duration: float
    context_before: list[str]  # Preceding glosses
    context_after: list[str]  # Following glosses
    signer_id: str | None = None
    video_path: str | None = None
    tier: str = "gloss"
    confidence: float | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "dataset": self.dataset,
            "sample_id": self.sample_id,
            "segment_idx": self.segment_idx,
            "start": self.start,
            "end": self.end,
            "duration": self.duration,
            "context_before": self.context_before,
            "context_after": self.context_after,
            "signer_id": self.signer_id,
            "video_path": self.video_path,
            "tier": self.tier,
            "confidence": self.confidence,
        }


@dataclass
class CrossDatasetResult:
    """Results of cross-dataset gloss search."""

    gloss: str
    total_matches: int
    matches_by_dataset: dict[str, list[GlossMatch]]
    dataset_frequencies: dict[str, int]
    duration_stats_by_dataset: dict[str, DurationStats]

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "gloss": self.gloss,
            "total_matches": self.total_matches,
            "matches_by_dataset": {
                ds: [m.to_dict() for m in matches]
                for ds, matches in self.matches_by_dataset.items()
            },
            "dataset_frequencies": self.dataset_frequencies,
            "duration_stats_by_dataset": {
                ds: stats.to_dict() for ds, stats in self.duration_stats_by_dataset.items()
            },
        }

    @property
    def datasets_with_matches(self) -> list[str]:
        """List of datasets that have matches."""
        return [ds for ds, freq in self.dataset_frequencies.items() if freq > 0]


@dataclass
class VocabularyComparison:
    """Comparison of vocabulary across datasets."""

    datasets: list[str]
    shared_glosses: set[str]  # Present in ALL datasets
    unique_glosses: dict[str, set[str]]  # Unique to each dataset
    pairwise_overlap: dict[tuple[str, str], float]  # Jaccard similarity
    frequency_correlation: dict[tuple[str, str], float]  # Spearman correlation
    vocabulary_sizes: dict[str, int]
    total_vocabulary: set[str]  # Union of all vocabularies

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "datasets": self.datasets,
            "shared_glosses": sorted(self.shared_glosses),
            "unique_glosses": {ds: sorted(glosses) for ds, glosses in self.unique_glosses.items()},
            "pairwise_overlap": {
                f"{ds1}_{ds2}": value for (ds1, ds2), value in self.pairwise_overlap.items()
            },
            "frequency_correlation": {
                f"{ds1}_{ds2}": value for (ds1, ds2), value in self.frequency_correlation.items()
            },
            "vocabulary_sizes": self.vocabulary_sizes,
            "total_vocabulary_size": len(self.total_vocabulary),
        }


@dataclass
class ParallelExample:
    """Parallel examples of a gloss from different datasets."""

    gloss: str
    examples: dict[str, list[GlossMatch]]  # dataset -> matches

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "gloss": self.gloss,
            "examples": {
                ds: [m.to_dict() for m in matches] for ds, matches in self.examples.items()
            },
        }

    @property
    def datasets_with_examples(self) -> list[str]:
        """List of datasets that have examples."""
        return [ds for ds, examples in self.examples.items() if examples]

    @property
    def total_examples(self) -> int:
        """Total number of examples across all datasets."""
        return sum(len(examples) for examples in self.examples.values())


@dataclass
class ContextPattern:
    """A recurring context pattern for a gloss."""

    pattern: tuple[str, ...]  # e.g., ("I", "GLOSS", "YOU")
    frequency: int
    datasets: list[str]
    examples: list[GlossMatch] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "pattern": list(self.pattern),
            "frequency": self.frequency,
            "datasets": self.datasets,
            "examples": [e.to_dict() for e in self.examples],
        }

    def __str__(self) -> str:
        return " -> ".join(self.pattern)


# =============================================================================
# Dataset Loading and Caching
# =============================================================================


class DatasetCache:
    """Cache for loaded datasets and their vocabularies."""

    _instance: DatasetCache | None = None

    def __init__(self):
        self._datasets: dict[str, BaseDataset] = {}
        self._vocabularies: dict[str, dict[str, int]] = {}
        self._samples_by_gloss: dict[str, dict[str, list[tuple[str, int]]]] = {}
        self._annotation_paths: dict[str, Path | None] = {}

    @classmethod
    def get_instance(cls) -> DatasetCache:
        """Get singleton instance."""
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    @classmethod
    def clear(cls):
        """Clear the cache."""
        if cls._instance is not None:
            cls._instance._datasets.clear()
            cls._instance._vocabularies.clear()
            cls._instance._samples_by_gloss.clear()
            cls._instance._annotation_paths.clear()

    def get_dataset(
        self,
        name: str,
        annotation_path: Path | None = None,
        **kwargs,
    ) -> BaseDataset:
        """Get or load a dataset."""
        cache_key = f"{name}_{annotation_path}"

        if cache_key not in self._datasets:
            from sltk.data.datasets import get_dataset

            try:
                if annotation_path:
                    dataset = get_dataset(
                        name,
                        unified_annotations=annotation_path,
                        **kwargs,
                    )
                else:
                    dataset = get_dataset(name, **kwargs)
                self._datasets[cache_key] = dataset
                self._annotation_paths[name] = annotation_path
            except Exception as e:
                logger.warning("Failed to load dataset %s: %s", name, e)
                raise

        return self._datasets[cache_key]

    def get_vocabulary(
        self,
        name: str,
        tier: str = "gloss",
        annotation_path: Path | None = None,
    ) -> dict[str, int]:
        """Get vocabulary with frequency counts for a dataset."""
        cache_key = f"{name}_{tier}_{annotation_path}"

        if cache_key not in self._vocabularies:
            self._build_vocabulary_index(name, tier, annotation_path)

        return self._vocabularies.get(cache_key, {})

    def get_samples_by_gloss(
        self,
        name: str,
        gloss: str,
        tier: str = "gloss",
        case_sensitive: bool = False,
        annotation_path: Path | None = None,
    ) -> list[tuple[str, int]]:
        """
        Get sample IDs and segment indices that contain a gloss.

        Returns:
            List of (sample_id, segment_idx) tuples
        """
        cache_key = f"{name}_{tier}_{annotation_path}"

        if cache_key not in self._samples_by_gloss:
            self._build_vocabulary_index(name, tier, annotation_path)

        index = self._samples_by_gloss.get(cache_key, {})

        # Normalize gloss for lookup
        lookup_gloss = gloss if case_sensitive else gloss.lower()

        return index.get(lookup_gloss, [])

    def _build_vocabulary_index(
        self,
        name: str,
        tier: str,
        annotation_path: Path | None = None,
    ):
        """Build vocabulary and sample index for a dataset."""
        cache_key = f"{name}_{tier}_{annotation_path}"

        try:
            dataset = self.get_dataset(name, annotation_path=annotation_path)
        except Exception as e:
            logger.warning("Failed to load dataset %s for indexing: %s", name, e)
            self._vocabularies[cache_key] = {}
            self._samples_by_gloss[cache_key] = {}
            return

        vocab: dict[str, int] = Counter()
        samples_index: dict[str, list[tuple[str, int]]] = defaultdict(list)

        for i in range(len(dataset)):
            try:
                sample = dataset[i]
            except Exception as e:
                logger.debug("Failed to load sample %d from %s: %s", i, name, e)
                continue

            if sample.segments is None:
                continue

            for seg_idx, segment in enumerate(sample.segments):
                # Check tier matches (case-insensitive tier matching)
                seg_tier = getattr(segment, "tier", "default")
                if not _tier_matches(seg_tier, tier):
                    continue

                label = segment.label
                if label:
                    # Normalize for indexing
                    normalized = label.lower()
                    vocab[normalized] += 1
                    samples_index[normalized].append((sample.id, seg_idx))

        self._vocabularies[cache_key] = dict(vocab)
        self._samples_by_gloss[cache_key] = dict(samples_index)


def _tier_matches(segment_tier: str, target_tier: str) -> bool:
    """Check if segment tier matches target tier (case-insensitive, handles variants)."""
    seg = segment_tier.lower().replace("-", "").replace("_", "")
    target = target_tier.lower().replace("-", "").replace("_", "")

    # Direct match
    if seg == target:
        return True

    # Handle common variants
    tier_aliases = {
        "gloss": ["gloss", "rhgloss", "lhgloss", "rhidgloss", "lhidgloss", "glossrh", "glosslh"],
        "glossrh": ["rhgloss", "rhidgloss", "glossrh"],
        "glosslh": ["lhgloss", "lhidgloss", "glosslh"],
        "translation": ["translation", "trans", "english"],
        "mouthing": ["mouthing", "mouth"],
    }

    target_aliases = tier_aliases.get(target, [target])
    return seg in target_aliases


def _get_available_datasets() -> list[str]:
    """Get list of available dataset names."""
    from sltk.data.datasets import list_datasets

    return list_datasets()


# =============================================================================
# Core Functions
# =============================================================================


def find_gloss_across_datasets(
    gloss: str,
    datasets: list[str] | None = None,
    case_sensitive: bool = False,
    context_window: int = 3,
    tier: str = "gloss",
    max_results_per_dataset: int = 100,
    annotation_paths: dict[str, Path] | None = None,
) -> CrossDatasetResult:
    """
    Find all occurrences of a gloss across multiple datasets.

    Args:
        gloss: The gloss/sign label to search for
        datasets: List of dataset names to search (None = all available)
        case_sensitive: Whether to match case exactly
        context_window: Number of surrounding glosses to include
        tier: Annotation tier to search (e.g., "gloss", "gloss-rh")
        max_results_per_dataset: Maximum matches to return per dataset
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        CrossDatasetResult with matches and statistics

    Example:
        >>> result = find_gloss_across_datasets("HELLO", datasets=["bslcp", "bobsl"])
        >>> print(f"Found {result.total_matches} matches")
        >>> for ds, matches in result.matches_by_dataset.items():
        ...     print(f"  {ds}: {len(matches)} matches")
    """
    if datasets is None:
        datasets = _get_available_datasets()

    annotation_paths = annotation_paths or {}
    cache = DatasetCache.get_instance()

    matches_by_dataset: dict[str, list[GlossMatch]] = {}
    dataset_frequencies: dict[str, int] = {}
    duration_stats_by_dataset: dict[str, DurationStats] = {}

    search_gloss = gloss if case_sensitive else gloss.lower()

    for ds_name in datasets:
        try:
            ann_path = annotation_paths.get(ds_name)
            dataset = cache.get_dataset(ds_name, annotation_path=ann_path)

            # Get samples containing this gloss
            sample_refs = cache.get_samples_by_gloss(
                ds_name,
                search_gloss,
                tier=tier,
                case_sensitive=case_sensitive,
                annotation_path=ann_path,
            )

            if not sample_refs:
                matches_by_dataset[ds_name] = []
                dataset_frequencies[ds_name] = 0
                continue

            # Collect matches
            matches: list[GlossMatch] = []
            durations: list[float] = []

            # Group by sample_id for efficient loading
            samples_to_load: dict[str, list[int]] = defaultdict(list)
            for sample_id, seg_idx in sample_refs:
                samples_to_load[sample_id].append(seg_idx)

            for sample_id, segment_indices in samples_to_load.items():
                if len(matches) >= max_results_per_dataset:
                    break

                try:
                    sample = dataset.get_by_id(sample_id)
                except Exception as e:
                    logger.debug("Failed to load sample %s: %s", sample_id, e)
                    continue

                if sample.segments is None:
                    continue

                # Build segment list for context
                tier_segments = [
                    (i, seg)
                    for i, seg in enumerate(sample.segments)
                    if _tier_matches(getattr(seg, "tier", "default"), tier)
                ]

                for seg_idx in segment_indices:
                    if len(matches) >= max_results_per_dataset:
                        break

                    if seg_idx >= len(sample.segments):
                        continue

                    segment = sample.segments[seg_idx]

                    # Verify label matches
                    label = segment.label
                    if label is None:
                        continue
                    if case_sensitive and label != gloss:
                        continue
                    if not case_sensitive and label.lower() != search_gloss:
                        continue

                    # Find position in tier_segments for context
                    tier_position = None
                    for pos, (idx, _) in enumerate(tier_segments):
                        if idx == seg_idx:
                            tier_position = pos
                            break

                    # Get context
                    context_before = []
                    context_after = []

                    if tier_position is not None:
                        # Before context
                        start_ctx = max(0, tier_position - context_window)
                        for i in range(start_ctx, tier_position):
                            ctx_seg = tier_segments[i][1]
                            if ctx_seg.label:
                                context_before.append(ctx_seg.label)

                        # After context
                        end_ctx = min(len(tier_segments), tier_position + context_window + 1)
                        for i in range(tier_position + 1, end_ctx):
                            ctx_seg = tier_segments[i][1]
                            if ctx_seg.label:
                                context_after.append(ctx_seg.label)

                    duration = segment.end - segment.start
                    durations.append(duration)

                    match = GlossMatch(
                        dataset=ds_name,
                        sample_id=sample_id,
                        segment_idx=seg_idx,
                        start=segment.start,
                        end=segment.end,
                        duration=duration,
                        context_before=context_before,
                        context_after=context_after,
                        signer_id=sample.signer_id,
                        video_path=str(sample.video_path) if sample.video_path else None,
                        tier=getattr(segment, "tier", "default"),
                        confidence=segment.confidence,
                    )
                    matches.append(match)

            matches_by_dataset[ds_name] = matches
            dataset_frequencies[ds_name] = len(sample_refs)  # Total frequency

            # Compute duration stats
            if durations:
                duration_stats_by_dataset[ds_name] = _compute_duration_stats(durations)

        except Exception as e:
            logger.warning("Error searching dataset %s: %s", ds_name, e)
            matches_by_dataset[ds_name] = []
            dataset_frequencies[ds_name] = 0

    total_matches = sum(len(m) for m in matches_by_dataset.values())

    return CrossDatasetResult(
        gloss=gloss,
        total_matches=total_matches,
        matches_by_dataset=matches_by_dataset,
        dataset_frequencies=dataset_frequencies,
        duration_stats_by_dataset=duration_stats_by_dataset,
    )


def compare_dataset_vocabularies(
    datasets: list[str],
    tier: str = "gloss",
    min_frequency: int = 1,
    annotation_paths: dict[str, Path] | None = None,
) -> VocabularyComparison:
    """
    Compare vocabularies across multiple datasets.

    Args:
        datasets: List of dataset names to compare
        tier: Annotation tier to use for vocabulary
        min_frequency: Minimum frequency for a gloss to be included
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        VocabularyComparison with overlap statistics

    Example:
        >>> comparison = compare_dataset_vocabularies(["bslcp", "bobsl", "asldict"])
        >>> print(f"Shared: {len(comparison.shared_glosses)}")
        >>> print(f"BSLCP unique: {len(comparison.unique_glosses['bslcp'])}")
    """
    annotation_paths = annotation_paths or {}
    cache = DatasetCache.get_instance()

    # Get vocabularies for each dataset
    vocabularies: dict[str, set[str]] = {}
    frequencies: dict[str, dict[str, int]] = {}
    vocabulary_sizes: dict[str, int] = {}

    for ds_name in datasets:
        ann_path = annotation_paths.get(ds_name)
        raw_vocab = cache.get_vocabulary(ds_name, tier=tier, annotation_path=ann_path)

        # Filter by minimum frequency
        filtered_vocab = {g for g, freq in raw_vocab.items() if freq >= min_frequency}
        vocabularies[ds_name] = filtered_vocab
        frequencies[ds_name] = {g: f for g, f in raw_vocab.items() if f >= min_frequency}
        vocabulary_sizes[ds_name] = len(filtered_vocab)

    # Compute shared glosses (present in ALL datasets)
    if vocabularies:
        if len(vocabularies) > 1:
            shared_glosses = set.intersection(*vocabularies.values())
        else:
            shared_glosses = set()
    else:
        shared_glosses = set()

    # Compute unique glosses (present in only one dataset)
    unique_glosses: dict[str, set[str]] = {}
    for ds_name, vocab in vocabularies.items():
        other_vocabs = [v for n, v in vocabularies.items() if n != ds_name]
        if other_vocabs:
            other_union = set.union(*other_vocabs)
            unique_glosses[ds_name] = vocab - other_union
        else:
            unique_glosses[ds_name] = vocab

    # Compute pairwise Jaccard similarity
    pairwise_overlap: dict[tuple[str, str], float] = {}
    for i, ds1 in enumerate(datasets):
        for ds2 in datasets[i + 1 :]:
            vocab1 = vocabularies.get(ds1, set())
            vocab2 = vocabularies.get(ds2, set())

            if vocab1 and vocab2:
                intersection = len(vocab1 & vocab2)
                union = len(vocab1 | vocab2)
                jaccard = intersection / union if union > 0 else 0.0
            else:
                jaccard = 0.0

            pairwise_overlap[(ds1, ds2)] = jaccard

    # Compute frequency correlation (Spearman)
    frequency_correlation: dict[tuple[str, str], float] = {}
    for i, ds1 in enumerate(datasets):
        for ds2 in datasets[i + 1 :]:
            freq1 = frequencies.get(ds1, {})
            freq2 = frequencies.get(ds2, {})

            # Only compute for shared glosses
            shared = set(freq1.keys()) & set(freq2.keys())
            if len(shared) >= 3:
                try:
                    from scipy.stats import spearmanr

                    shared_list = sorted(shared)
                    f1 = [freq1[g] for g in shared_list]
                    f2 = [freq2[g] for g in shared_list]
                    corr, _ = spearmanr(f1, f2)
                    frequency_correlation[(ds1, ds2)] = float(corr) if np.isfinite(corr) else 0.0
                except ImportError:
                    # Fall back to simple correlation
                    frequency_correlation[(ds1, ds2)] = _simple_correlation(freq1, freq2, shared)
            else:
                frequency_correlation[(ds1, ds2)] = 0.0

    # Total vocabulary
    total_vocabulary = set.union(*vocabularies.values()) if vocabularies else set()

    return VocabularyComparison(
        datasets=datasets,
        shared_glosses=shared_glosses,
        unique_glosses=unique_glosses,
        pairwise_overlap=pairwise_overlap,
        frequency_correlation=frequency_correlation,
        vocabulary_sizes=vocabulary_sizes,
        total_vocabulary=total_vocabulary,
    )


def find_parallel_examples(
    glosses: list[str],
    datasets: list[str],
    tier: str = "gloss",
    max_examples_per_dataset: int = 5,
    annotation_paths: dict[str, Path] | None = None,
) -> list[ParallelExample]:
    """
    Find parallel examples of glosses across datasets.

    Useful for comparing citation forms (dictionaries) with
    naturalistic production (corpora).

    Args:
        glosses: List of glosses to find examples for
        datasets: List of dataset names to search
        tier: Annotation tier to search
        max_examples_per_dataset: Maximum examples per dataset per gloss
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        List of ParallelExample objects

    Example:
        >>> examples = find_parallel_examples(
        ...     glosses=["HELLO", "THANK-YOU"],
        ...     datasets=["bslcp", "bobsl"],
        ... )
        >>> for ex in examples:
        ...     print(f"{ex.gloss}: {ex.total_examples} examples")
    """
    results: list[ParallelExample] = []

    for gloss in glosses:
        cross_result = find_gloss_across_datasets(
            gloss=gloss,
            datasets=datasets,
            tier=tier,
            max_results_per_dataset=max_examples_per_dataset,
            annotation_paths=annotation_paths,
        )

        examples: dict[str, list[GlossMatch]] = {}
        for ds_name in datasets:
            matches = cross_result.matches_by_dataset.get(ds_name, [])
            examples[ds_name] = matches[:max_examples_per_dataset]

        results.append(ParallelExample(gloss=gloss, examples=examples))

    return results


def compare_gloss_durations(
    gloss: str,
    datasets: list[str] | None = None,
    tier: str = "gloss",
    annotation_paths: dict[str, Path] | None = None,
) -> dict[str, DurationStats]:
    """
    Compare duration of a gloss across datasets.

    Useful for studying citation vs. connected signing duration differences.

    Args:
        gloss: The gloss to analyze
        datasets: List of dataset names (None = all available)
        tier: Annotation tier to search
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        Dictionary mapping dataset name to DurationStats

    Example:
        >>> stats = compare_gloss_durations("HELLO", datasets=["bslcp", "bobsl"])
        >>> for ds, s in stats.items():
        ...     print(f"{ds}: mean={s.mean:.3f}s, std={s.std:.3f}s")
    """
    result = find_gloss_across_datasets(
        gloss=gloss,
        datasets=datasets,
        tier=tier,
        max_results_per_dataset=1000,  # Get enough for statistics
        annotation_paths=annotation_paths,
    )

    return result.duration_stats_by_dataset


def analyze_gloss_contexts(
    gloss: str,
    datasets: list[str] | None = None,
    context_size: int = 1,
    tier: str = "gloss",
    min_frequency: int = 2,
    max_patterns: int = 50,
    annotation_paths: dict[str, Path] | None = None,
) -> list[ContextPattern]:
    """
    Analyze what contexts a gloss appears in across datasets.

    Args:
        gloss: The gloss to analyze
        datasets: List of dataset names (None = all available)
        context_size: Context window size (1 = bigrams, 2 = trigrams with gloss in middle)
        tier: Annotation tier to search
        min_frequency: Minimum frequency for a pattern to be included
        max_patterns: Maximum number of patterns to return
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        List of ContextPattern objects sorted by frequency

    Example:
        >>> patterns = analyze_gloss_contexts("HELLO", context_size=1)
        >>> for p in patterns[:10]:
        ...     print(f"{p}: {p.frequency} occurrences")
    """
    result = find_gloss_across_datasets(
        gloss=gloss,
        datasets=datasets,
        context_window=context_size,
        tier=tier,
        max_results_per_dataset=500,
        annotation_paths=annotation_paths,
    )

    # Count patterns
    pattern_counts: Counter[tuple[str, ...]] = Counter()
    pattern_datasets: dict[tuple[str, ...], set[str]] = defaultdict(set)
    pattern_examples: dict[tuple[str, ...], list[GlossMatch]] = defaultdict(list)

    placeholder = "<GLOSS>"

    for ds_name, matches in result.matches_by_dataset.items():
        for match in matches:
            # Build pattern: context_before + gloss + context_after
            if context_size == 1:
                # Bigram patterns
                for ctx_gloss in match.context_before[-1:]:
                    pat: tuple[str, ...] = (ctx_gloss, placeholder)
                    pattern_counts[pat] += 1
                    pattern_datasets[pat].add(ds_name)
                    if len(pattern_examples[pat]) < 3:
                        pattern_examples[pat].append(match)

                for ctx_gloss in match.context_after[:1]:
                    pat = (placeholder, ctx_gloss)
                    pattern_counts[pat] += 1
                    pattern_datasets[pat].add(ds_name)
                    if len(pattern_examples[pat]) < 3:
                        pattern_examples[pat].append(match)
            else:
                # N-gram patterns with gloss in context
                before = match.context_before[-context_size:]
                after = match.context_after[:context_size]
                pat = tuple(before) + (placeholder,) + tuple(after)

                if len(pat) >= 2:  # At least one context word
                    pattern_counts[pat] += 1
                    pattern_datasets[pat].add(ds_name)
                    if len(pattern_examples[pat]) < 3:
                        pattern_examples[pat].append(match)

    # Filter and sort
    patterns: list[ContextPattern] = []
    for pattern, count in pattern_counts.most_common(max_patterns):
        if count >= min_frequency:
            # Replace placeholder with actual gloss in pattern
            final_pattern = tuple(gloss if p == placeholder else p for p in pattern)
            patterns.append(
                ContextPattern(
                    pattern=final_pattern,
                    frequency=count,
                    datasets=sorted(pattern_datasets[pattern]),
                    examples=pattern_examples[pattern],
                )
            )

    return patterns


def batch_find_glosses(
    glosses: list[str],
    datasets: list[str] | None = None,
    tier: str = "gloss",
    max_results_per_dataset: int = 10,
    annotation_paths: dict[str, Path] | None = None,
) -> dict[str, CrossDatasetResult]:
    """
    Find multiple glosses across datasets efficiently.

    Args:
        glosses: List of glosses to search for
        datasets: List of dataset names (None = all available)
        tier: Annotation tier to search
        max_results_per_dataset: Maximum results per dataset per gloss
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        Dictionary mapping gloss to CrossDatasetResult

    Example:
        >>> results = batch_find_glosses(
        ...     ["HELLO", "THANK-YOU", "GOODBYE"],
        ...     datasets=["bslcp", "bobsl"],
        ... )
        >>> for gloss, result in results.items():
        ...     print(f"{gloss}: {result.total_matches} matches")
    """
    results: dict[str, CrossDatasetResult] = {}

    for gloss in glosses:
        results[gloss] = find_gloss_across_datasets(
            gloss=gloss,
            datasets=datasets,
            tier=tier,
            max_results_per_dataset=max_results_per_dataset,
            annotation_paths=annotation_paths,
        )

    return results


# =============================================================================
# Helper Functions
# =============================================================================


def _compute_duration_stats(durations: list[float]) -> DurationStats:
    """Compute duration statistics from a list of durations."""
    if not durations:
        return DurationStats(
            count=0,
            mean=0.0,
            std=0.0,
            min=0.0,
            max=0.0,
            median=0.0,
            quartiles=(0.0, 0.0, 0.0),
        )

    arr = np.array(durations)

    return DurationStats(
        count=len(durations),
        mean=float(arr.mean()),
        std=float(arr.std()),
        min=float(arr.min()),
        max=float(arr.max()),
        median=float(np.median(arr)),
        quartiles=(
            float(np.percentile(arr, 25)),
            float(np.percentile(arr, 50)),
            float(np.percentile(arr, 75)),
        ),
    )


def _simple_correlation(
    freq1: dict[str, int],
    freq2: dict[str, int],
    shared: set[str],
) -> float:
    """Compute simple Pearson correlation for frequency counts."""
    if len(shared) < 3:
        return 0.0

    shared_list = sorted(shared)
    f1 = np.array([freq1[g] for g in shared_list], dtype=float)
    f2 = np.array([freq2[g] for g in shared_list], dtype=float)

    # Normalize
    f1_norm = f1 - f1.mean()
    f2_norm = f2 - f2.mean()

    std1 = f1_norm.std()
    std2 = f2_norm.std()

    if std1 == 0 or std2 == 0:
        return 0.0

    corr = np.sum(f1_norm * f2_norm) / (len(shared) * std1 * std2)
    return float(corr) if np.isfinite(corr) else 0.0


def get_dataset_vocabulary_summary(
    datasets: list[str] | None = None,
    tier: str = "gloss",
    annotation_paths: dict[str, Path] | None = None,
) -> dict[str, dict[str, Any]]:
    """
    Get vocabulary summary for multiple datasets.

    Args:
        datasets: List of dataset names (None = all available)
        tier: Annotation tier to use
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        Dictionary with vocabulary statistics per dataset

    Example:
        >>> summary = get_dataset_vocabulary_summary(["bslcp", "bobsl"])
        >>> for ds, stats in summary.items():
        ...     print(f"{ds}: {stats['vocabulary_size']} unique glosses")
    """
    if datasets is None:
        datasets = _get_available_datasets()

    annotation_paths = annotation_paths or {}
    cache = DatasetCache.get_instance()

    summaries: dict[str, dict[str, Any]] = {}

    for ds_name in datasets:
        try:
            ann_path = annotation_paths.get(ds_name)
            vocab = cache.get_vocabulary(ds_name, tier=tier, annotation_path=ann_path)

            if vocab:
                freqs = list(vocab.values())
                summaries[ds_name] = {
                    "vocabulary_size": len(vocab),
                    "total_tokens": sum(freqs),
                    "max_frequency": max(freqs),
                    "min_frequency": min(freqs),
                    "mean_frequency": np.mean(freqs),
                    "top_10": sorted(vocab.items(), key=lambda x: -x[1])[:10],
                }
            else:
                summaries[ds_name] = {
                    "vocabulary_size": 0,
                    "total_tokens": 0,
                    "error": "No vocabulary found",
                }
        except Exception as e:
            summaries[ds_name] = {
                "vocabulary_size": 0,
                "total_tokens": 0,
                "error": str(e),
            }

    return summaries


def find_common_glosses(
    datasets: list[str],
    tier: str = "gloss",
    min_frequency_per_dataset: int = 5,
    top_n: int = 100,
    annotation_paths: dict[str, Path] | None = None,
) -> list[tuple[str, dict[str, int]]]:
    """
    Find glosses that are common across all specified datasets.

    Args:
        datasets: List of dataset names
        tier: Annotation tier to use
        min_frequency_per_dataset: Minimum frequency in each dataset
        top_n: Maximum number of results to return
        annotation_paths: Optional dict of dataset name -> annotation path

    Returns:
        List of (gloss, frequency_dict) tuples sorted by total frequency

    Example:
        >>> common = find_common_glosses(["bslcp", "bobsl"], top_n=20)
        >>> for gloss, freqs in common:
        ...     print(f"{gloss}: {freqs}")
    """
    comparison = compare_dataset_vocabularies(
        datasets=datasets,
        tier=tier,
        min_frequency=min_frequency_per_dataset,
        annotation_paths=annotation_paths,
    )

    cache = DatasetCache.get_instance()
    annotation_paths = annotation_paths or {}

    # Get frequencies for shared glosses
    results: list[tuple[str, dict[str, int]]] = []

    for gloss in comparison.shared_glosses:
        freq_dict: dict[str, int] = {}
        for ds_name in datasets:
            ann_path = annotation_paths.get(ds_name)
            vocab = cache.get_vocabulary(ds_name, tier=tier, annotation_path=ann_path)
            freq_dict[ds_name] = vocab.get(gloss, 0)

        # Only include if meets minimum in all datasets
        if all(f >= min_frequency_per_dataset for f in freq_dict.values()):
            results.append((gloss, freq_dict))

    # Sort by total frequency
    results.sort(key=lambda x: sum(x[1].values()), reverse=True)

    return results[:top_n]
