"""
Dataset analysis and comparison tools.

Provides utilities for comparing datasets, analyzing vocabulary,
and computing lexical statistics for sign language corpora.
"""

from __future__ import annotations

import logging
from collections import Counter
from dataclasses import dataclass
from typing import Any

import numpy as np

from sltk.data.core import Sample
from sltk.data.datasets.base import BaseDataset

logger = logging.getLogger(__name__)


@dataclass
class LexicalStatistics:
    """Results from lexical analysis of a dataset."""

    vocabulary_size: int
    total_tokens: int
    type_token_ratio: float
    hapax_legomena: int  # Signs appearing exactly once
    hapax_ratio: float
    dis_legomena: int  # Signs appearing exactly twice
    most_common: list[tuple[str, int]]
    least_common: list[tuple[str, int]]
    zipf_coefficient: float | None
    frequency_distribution: dict[str, int]
    oov_rate_by_split: dict[str, float] | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "vocabulary_size": self.vocabulary_size,
            "total_tokens": self.total_tokens,
            "type_token_ratio": self.type_token_ratio,
            "hapax_legomena": self.hapax_legomena,
            "hapax_ratio": self.hapax_ratio,
            "dis_legomena": self.dis_legomena,
            "most_common": self.most_common,
            "least_common": self.least_common,
            "zipf_coefficient": self.zipf_coefficient,
            "oov_rate_by_split": self.oov_rate_by_split,
        }


@dataclass
class DatasetComparison:
    """Results from comparing multiple datasets."""

    datasets: list[str]
    vocabulary_sizes: dict[str, int]
    vocabulary_overlap: float  # Jaccard similarity
    shared_glosses: list[str]
    unique_glosses: dict[str, list[str]]
    duration_statistics: dict[str, dict[str, float]]
    sample_counts: dict[str, int]
    language_distribution: dict[str, dict[str, int]]

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "datasets": self.datasets,
            "vocabulary_sizes": self.vocabulary_sizes,
            "vocabulary_overlap": self.vocabulary_overlap,
            "shared_glosses": self.shared_glosses,
            "unique_glosses": self.unique_glosses,
            "duration_statistics": self.duration_statistics,
            "sample_counts": self.sample_counts,
            "language_distribution": self.language_distribution,
        }


def lexical_statistics(
    dataset: BaseDataset | list[Sample] | list[str],
    train_vocab: set | None = None,
    top_n: int = 20,
) -> LexicalStatistics:
    """
    Compute lexical statistics for a dataset.

    Args:
        dataset: Dataset, list of samples, or list of glosses
        train_vocab: Optional training vocabulary for OOV calculation
        top_n: Number of most/least common items to return

    Returns:
        LexicalStatistics with vocabulary analysis

    Example:
        >>> stats = lexical_statistics(dataset)
        >>> print(f"Vocabulary size: {stats.vocabulary_size}")
        >>> print(f"Hapax legomena: {stats.hapax_legomena}")
    """
    # Extract glosses from dataset
    glosses = _extract_glosses(dataset)

    if not glosses:
        return LexicalStatistics(
            vocabulary_size=0,
            total_tokens=0,
            type_token_ratio=0.0,
            hapax_legomena=0,
            hapax_ratio=0.0,
            dis_legomena=0,
            most_common=[],
            least_common=[],
            zipf_coefficient=None,
            frequency_distribution={},
        )

    # Count frequencies
    freq = Counter(glosses)
    vocabulary = set(freq.keys())
    total_tokens = sum(freq.values())

    # Hapax and dis legomena
    hapax = [g for g, c in freq.items() if c == 1]
    dis = [g for g, c in freq.items() if c == 2]

    # Type-token ratio
    ttr = len(vocabulary) / total_tokens if total_tokens > 0 else 0.0

    # Most and least common
    most_common = freq.most_common(top_n)
    least_common = (
        freq.most_common()[: -top_n - 1 : -1] if len(freq) > top_n else list(freq.items())
    )

    # Zipf coefficient (log-log regression)
    zipf_coef = _compute_zipf_coefficient(freq)

    # OOV rate if train vocab provided
    oov_rate = None
    if train_vocab is not None:
        oov_tokens = [g for g in glosses if g not in train_vocab]
        oov_rate = {"overall": len(oov_tokens) / len(glosses) if glosses else 0.0}

    return LexicalStatistics(
        vocabulary_size=len(vocabulary),
        total_tokens=total_tokens,
        type_token_ratio=ttr,
        hapax_legomena=len(hapax),
        hapax_ratio=len(hapax) / len(vocabulary) if vocabulary else 0.0,
        dis_legomena=len(dis),
        most_common=most_common,
        least_common=least_common,
        zipf_coefficient=zipf_coef,
        frequency_distribution=dict(freq),
        oov_rate_by_split=oov_rate,
    )


def compare_datasets(
    datasets: dict[str, BaseDataset | list[Sample]],
) -> DatasetComparison:
    """
    Compare multiple datasets.

    Args:
        datasets: Dictionary mapping dataset names to datasets

    Returns:
        DatasetComparison with overlap and statistics

    Example:
        >>> comparison = compare_datasets({
        ...     "BOBSL": bobsl_dataset,
        ...     "How2Sign": h2s_dataset,
        ... })
        >>> print(f"Vocabulary overlap: {comparison.vocabulary_overlap:.2%}")
    """
    dataset_names = list(datasets.keys())

    # Extract vocabularies
    vocabularies = {}
    sample_counts = {}
    durations = {}
    languages = {}

    for name, dataset in datasets.items():
        glosses = _extract_glosses(dataset)
        vocabularies[name] = set(glosses)

        # Count samples
        if isinstance(dataset, BaseDataset):
            sample_counts[name] = len(dataset)
        else:
            sample_counts[name] = len(dataset)

        # Duration statistics
        durations[name] = _compute_duration_stats(dataset)

        # Language distribution
        languages[name] = _compute_language_distribution(dataset)

    # Vocabulary overlap (Jaccard similarity over all datasets)
    all_vocabs = list(vocabularies.values())
    if len(all_vocabs) >= 2:
        intersection = set.intersection(*all_vocabs)
        union = set.union(*all_vocabs)
        overlap = len(intersection) / len(union) if union else 0.0
        shared = sorted(intersection)
    else:
        overlap = 1.0 if all_vocabs else 0.0
        shared = sorted(all_vocabs[0]) if all_vocabs else []

    # Unique glosses per dataset
    unique = {}
    for name, vocab in vocabularies.items():
        other_vocabs = [v for n, v in vocabularies.items() if n != name]
        if other_vocabs:
            other_union = set.union(*other_vocabs)
            unique[name] = sorted(vocab - other_union)
        else:
            unique[name] = sorted(vocab)

    return DatasetComparison(
        datasets=dataset_names,
        vocabulary_sizes={n: len(v) for n, v in vocabularies.items()},
        vocabulary_overlap=overlap,
        shared_glosses=shared[:100],  # Limit to first 100
        unique_glosses={n: v[:50] for n, v in unique.items()},  # Limit to first 50
        duration_statistics=durations,
        sample_counts=sample_counts,
        language_distribution=languages,
    )


def pos_analysis(
    glosses: list[str],
    pos_mapping: dict[str, str] | None = None,
) -> dict[str, Any]:
    """
    Analyze parts of speech distribution in glosses.

    Args:
        glosses: List of gloss strings
        pos_mapping: Optional mapping from glosses to POS tags

    Returns:
        Dictionary with POS distribution and statistics

    Example:
        >>> pos_stats = pos_analysis(glosses, pos_mapping=gloss_to_pos)
        >>> print(pos_stats["distribution"])
    """
    if pos_mapping is None:
        # Use simple heuristic POS tagging based on gloss patterns
        pos_mapping = _heuristic_pos_tagging(glosses)

    pos_counts: Counter[str] = Counter()
    for gloss in glosses:
        pos = pos_mapping.get(gloss, "UNKNOWN")
        pos_counts[pos] += 1

    total = sum(pos_counts.values())
    distribution = {
        pos: {"count": count, "ratio": count / total if total > 0 else 0.0}
        for pos, count in pos_counts.items()
    }

    return {
        "distribution": distribution,
        "total_tokens": total,
        "num_pos_tags": len(pos_counts),
        "most_common_pos": pos_counts.most_common(5),
    }


def vocabulary_growth(
    samples: list[Sample],
    window_size: int = 100,
) -> dict[str, list[int]]:
    """
    Analyze vocabulary growth over samples.

    Useful for understanding how vocabulary size increases as more data is seen.

    Args:
        samples: List of samples in order
        window_size: Number of samples per measurement point

    Returns:
        Dictionary with growth curves
    """
    seen_vocab = set()
    running_freq: Counter[str] = Counter()
    hapax_counts = []
    vocab_sizes = []
    sample_indices = []

    for i, sample in enumerate(samples):
        if sample.glosses:
            seen_vocab.update(sample.glosses)
            running_freq.update(sample.glosses)

        if (i + 1) % window_size == 0:
            # Count current hapax using running frequency counter
            hapax = sum(1 for c in running_freq.values() if c == 1)

            vocab_sizes.append(len(seen_vocab))
            hapax_counts.append(hapax)
            sample_indices.append(i + 1)

    return {
        "sample_indices": sample_indices,
        "vocabulary_sizes": vocab_sizes,
        "hapax_counts": hapax_counts,
    }


def _extract_glosses(
    dataset: BaseDataset | list[Sample] | list[str],
) -> list[str]:
    """Extract all glosses from a dataset."""
    glosses = []

    if isinstance(dataset, list):
        for item in dataset:
            if isinstance(item, str):
                glosses.append(item)
            elif isinstance(item, Sample) and item.glosses:
                glosses.extend(item.glosses)
    elif isinstance(dataset, BaseDataset):
        for i in range(len(dataset)):
            try:
                sample = dataset[i]
                if sample.glosses:
                    glosses.extend(sample.glosses)
            except (KeyError, IndexError, FileNotFoundError) as e:
                logger.debug("Could not load sample %d from dataset: %s", i, e)
                continue
            except Exception as e:
                logger.warning("Unexpected error loading sample %d from dataset: %s", i, e)
                continue

    return glosses


def _compute_zipf_coefficient(freq: Counter) -> float | None:
    """
    Compute Zipf's law coefficient using log-log regression.

    Returns None if there are insufficient data points or if the
    calculation cannot be performed (e.g., division by zero risk).
    """
    if len(freq) < 10:
        logger.debug(
            "Insufficient data for Zipf coefficient: %d items (need at least 10)", len(freq)
        )
        return None

    # Sort by frequency
    sorted_freq = sorted(freq.values(), reverse=True)

    # Filter out zero frequencies (cannot take log of zero)
    sorted_freq = [f for f in sorted_freq if f > 0]
    if len(sorted_freq) < 10:
        logger.debug(
            "Insufficient non-zero frequencies for Zipf coefficient: %d items",
            len(sorted_freq),
        )
        return None

    # Log-log regression
    ranks = np.arange(1, len(sorted_freq) + 1)
    log_ranks = np.log(ranks)
    log_freqs = np.log(sorted_freq)

    # Check for NaN or Inf values
    if not np.all(np.isfinite(log_ranks)) or not np.all(np.isfinite(log_freqs)):
        logger.warning("Non-finite values encountered in Zipf coefficient calculation")
        return None

    # Simple linear regression
    n = len(log_ranks)
    sum_x = np.sum(log_ranks)
    sum_y = np.sum(log_freqs)
    sum_xy = np.sum(log_ranks * log_freqs)
    sum_xx = np.sum(log_ranks * log_ranks)

    denominator = n * sum_xx - sum_x * sum_x
    if abs(denominator) < 1e-10:
        logger.debug("Zipf coefficient denominator too small: %e", denominator)
        return None

    slope = (n * sum_xy - sum_x * sum_y) / denominator

    return -slope  # Zipf coefficient is negative of slope


def _compute_duration_stats(
    dataset: BaseDataset | list[Sample],
) -> dict[str, float]:
    """Compute duration statistics for a dataset."""
    durations = []

    if isinstance(dataset, BaseDataset):
        for i in range(min(len(dataset), 1000)):  # Sample up to 1000
            try:
                sample = dataset[i]
                if sample.duration is not None:
                    durations.append(sample.duration)
            except (KeyError, IndexError, FileNotFoundError) as e:
                logger.debug("Could not load sample %d for duration stats: %s", i, e)
                continue
            except Exception as e:
                logger.warning("Unexpected error loading sample %d for duration stats: %s", i, e)
                continue
    else:
        for sample in dataset:
            if isinstance(sample, Sample) and sample.duration is not None:
                durations.append(sample.duration)

    if not durations:
        return {"mean": 0.0, "std": 0.0, "min": 0.0, "max": 0.0, "total": 0.0}

    return {
        "mean": float(np.mean(durations)),
        "std": float(np.std(durations)),
        "min": float(np.min(durations)),
        "max": float(np.max(durations)),
        "total": float(np.sum(durations)),
    }


def _compute_language_distribution(
    dataset: BaseDataset | list[Sample],
) -> dict[str, int]:
    """Compute sign language distribution in dataset."""
    lang_counts: Counter[str] = Counter()

    if isinstance(dataset, BaseDataset):
        if hasattr(dataset, "LANGUAGES"):
            for lang in dataset.LANGUAGES:
                lang_counts[lang] = 0  # Initialize known languages

        for i in range(min(len(dataset), 1000)):
            try:
                sample = dataset[i]
                if hasattr(sample, "metadata") and "sign_language" in sample.metadata:
                    lang_counts[sample.metadata["sign_language"]] += 1
            except (KeyError, IndexError, FileNotFoundError) as e:
                logger.debug("Could not load sample %d for language distribution: %s", i, e)
                continue
            except Exception as e:
                logger.warning(
                    "Unexpected error loading sample %d for language distribution: %s", i, e
                )
                continue
    else:
        for sample in dataset:
            if isinstance(sample, Sample) and hasattr(sample, "metadata"):
                lang = sample.metadata.get("sign_language", "unknown")
                lang_counts[lang] += 1

    return dict(lang_counts)


def _heuristic_pos_tagging(glosses: list[str]) -> dict[str, str]:
    """
    Simple heuristic POS tagging for sign language glosses.

    This is a basic heuristic and should be replaced with proper
    linguistic annotation when available.
    """
    # Common patterns in sign language glosses
    pronouns = {"I", "ME", "YOU", "HE", "SHE", "IT", "WE", "THEY", "INDEX", "SELF"}
    wh_words = {"WHAT", "WHO", "WHERE", "WHEN", "WHY", "HOW", "WHICH"}
    negation = {"NOT", "NO", "NONE", "NOTHING", "NEVER", "DONT"}
    auxiliaries = {"CAN", "WILL", "MUST", "SHOULD", "COULD", "WOULD", "HAVE"}

    pos_map = {}

    for gloss in set(glosses):
        upper = gloss.upper()

        if upper in pronouns or upper.startswith("IX-"):
            pos_map[gloss] = "PRONOUN"
        elif upper in wh_words:
            pos_map[gloss] = "WH-WORD"
        elif upper in negation:
            pos_map[gloss] = "NEGATION"
        elif upper in auxiliaries:
            pos_map[gloss] = "AUXILIARY"
        elif upper.endswith("-ED") or upper.endswith("-ING"):
            pos_map[gloss] = "VERB"
        elif "-" in gloss and gloss.count("-") >= 2:
            pos_map[gloss] = "COMPOUND"
        else:
            pos_map[gloss] = "CONTENT"  # Default

    return pos_map
