"""
Research utilities for sign language analysis.

Provides common workflows for researchers:
- Batch processing of videos and poses
- Dataset statistics and comparison
- Cross-dataset evaluation
- Export utilities for training pipelines
"""

from __future__ import annotations

import json
from collections import Counter
from collections.abc import Iterator, Sequence
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

from sltk.exceptions import ConfigurationError, UnsupportedFormatError


@dataclass
class DatasetStats:
    """Statistics about a dataset or collection of files."""

    name: str
    num_samples: int = 0
    num_segments: int = 0
    total_duration_sec: float = 0.0
    vocabulary_size: int = 0
    avg_segment_duration: float = 0.0
    segment_duration_std: float = 0.0
    glosses_per_sample_avg: float = 0.0

    # Detailed breakdowns
    tier_counts: dict[str, int] = field(default_factory=dict)
    gloss_frequencies: dict[str, int] = field(default_factory=dict)
    duration_histogram: list[int] = field(default_factory=list)

    # Metadata
    signers: list[str] = field(default_factory=list)
    num_signers: int = 0

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)

    def to_json(self, path: str | Path) -> None:
        """Save to JSON file."""
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)

    @classmethod
    def from_json(cls, path: str | Path) -> DatasetStats:
        """Load from JSON file."""
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        return cls(**data)


def compute_dataset_statistics(
    elan_files: Sequence[str | Path],
    name: str = "dataset",
    include_gloss_frequencies: bool = True,
    top_k_glosses: int = 100,
) -> DatasetStats:
    """
    Compute comprehensive statistics from a collection of ELAN files.

    Args:
        elan_files: List of paths to ELAN (.eaf) files
        name: Name for the dataset
        include_gloss_frequencies: Whether to compute gloss frequencies
        top_k_glosses: Number of top glosses to include

    Returns:
        DatasetStats object with computed statistics
    """
    from sltk.io.elan import read_eaf

    all_durations = []
    all_glosses = []
    tier_counts: Counter[str] = Counter()
    glosses_per_sample = []
    signers = set()

    for elan_path in elan_files:
        try:
            segments = read_eaf(elan_path)

            sample_glosses = []
            for seg in segments:
                duration = seg.end - seg.start
                all_durations.append(duration)
                if seg.tier:
                    tier_counts[seg.tier] += 1

                if seg.label:
                    all_glosses.append(seg.label)
                    sample_glosses.append(seg.label)

            glosses_per_sample.append(len(sample_glosses))

            # Try to extract signer ID from filename
            stem = Path(elan_path).stem
            if "_" in stem:
                signer = stem.split("_")[0]
                signers.add(signer)
        except Exception:
            continue

    # Compute statistics
    durations_array = np.array(all_durations) if all_durations else np.array([0])
    gloss_counter = Counter(all_glosses)

    # Duration histogram (0.1s bins up to 5s)
    bins = np.arange(0, 5.1, 0.1)
    hist, _ = np.histogram(durations_array, bins=bins)

    return DatasetStats(
        name=name,
        num_samples=len(elan_files),
        num_segments=len(all_durations),
        total_duration_sec=float(durations_array.sum()),
        vocabulary_size=len(gloss_counter),
        avg_segment_duration=float(durations_array.mean()) if len(durations_array) > 0 else 0,
        segment_duration_std=float(durations_array.std()) if len(durations_array) > 0 else 0,
        glosses_per_sample_avg=float(np.mean(glosses_per_sample)) if glosses_per_sample else 0,
        tier_counts=dict(tier_counts),
        gloss_frequencies=(
            dict(gloss_counter.most_common(top_k_glosses)) if include_gloss_frequencies else {}
        ),
        duration_histogram=hist.tolist(),
        signers=sorted(signers),
        num_signers=len(signers),
    )


def compare_vocabularies(
    stats1: DatasetStats,
    stats2: DatasetStats,
) -> dict[str, Any]:
    """
    Compare vocabularies between two datasets.

    Returns dictionary with overlap metrics.
    """
    vocab1 = set(stats1.gloss_frequencies.keys())
    vocab2 = set(stats2.gloss_frequencies.keys())

    intersection = vocab1 & vocab2
    union = vocab1 | vocab2
    only_in_1 = vocab1 - vocab2
    only_in_2 = vocab2 - vocab1

    return {
        "dataset_1": stats1.name,
        "dataset_2": stats2.name,
        "vocab_size_1": len(vocab1),
        "vocab_size_2": len(vocab2),
        "intersection_size": len(intersection),
        "union_size": len(union),
        "jaccard_index": len(intersection) / len(union) if union else 0,
        "overlap_ratio_1": len(intersection) / len(vocab1) if vocab1 else 0,
        "overlap_ratio_2": len(intersection) / len(vocab2) if vocab2 else 0,
        "only_in_1_count": len(only_in_1),
        "only_in_2_count": len(only_in_2),
        "common_glosses": sorted(intersection)[:50],  # Top 50
        "only_in_1_sample": sorted(only_in_1)[:20],
        "only_in_2_sample": sorted(only_in_2)[:20],
    }


def find_gloss_examples(
    gloss: str,
    elan_files: Sequence[str | Path],
    max_examples: int = 20,
    context_window: int = 2,
) -> list[dict[str, Any]]:
    """
    Find examples of a specific gloss across multiple ELAN files.

    Args:
        gloss: Target gloss to find
        elan_files: List of ELAN file paths
        max_examples: Maximum number of examples to return
        context_window: Number of surrounding glosses to include

    Returns:
        List of example dictionaries with timing and context
    """
    from sltk.io.elan import read_eaf

    examples: list[dict[str, Any]] = []
    gloss_lower = gloss.lower()

    for elan_path in elan_files:
        if len(examples) >= max_examples:
            break

        try:
            segments = read_eaf(elan_path)

            for i, seg in enumerate(segments):
                if len(examples) >= max_examples:
                    break

                if seg.label and seg.label.lower() == gloss_lower:
                    # Get context
                    context_before = []
                    context_after = []

                    for j in range(max(0, i - context_window), i):
                        if segments[j].label:
                            context_before.append(segments[j].label)

                    for j in range(i + 1, min(len(segments), i + context_window + 1)):
                        if segments[j].label:
                            context_after.append(segments[j].label)

                    examples.append(
                        {
                            "file": Path(elan_path).stem,
                            "gloss": seg.label,
                            "start": seg.start,
                            "end": seg.end,
                            "duration": seg.end - seg.start,
                            "tier": seg.tier,
                            "context_before": context_before,
                            "context_after": context_after,
                        }
                    )
        except Exception:
            continue

    return examples


def export_segments_for_training(
    elan_files: Sequence[str | Path],
    output_path: str | Path,
    format: str = "csv",
    include_context: bool = True,
    context_window: int = 2,
) -> Path:
    """
    Export segments from multiple ELAN files for ML training.

    Args:
        elan_files: List of ELAN file paths
        output_path: Output file path
        format: Output format ('csv', 'json', 'jsonl')
        include_context: Include surrounding gloss context
        context_window: Number of context glosses on each side

    Returns:
        Path to created file
    """
    import csv

    from sltk.io.elan import read_eaf

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    all_records = []

    for elan_path in elan_files:
        try:
            segments = read_eaf(elan_path)
            file_stem = Path(elan_path).stem

            for i, seg in enumerate(segments):
                record = {
                    "file": file_stem,
                    "segment_idx": i,
                    "start": seg.start,
                    "end": seg.end,
                    "duration": seg.end - seg.start,
                    "label": seg.label or "",
                    "tier": seg.tier,
                }

                if include_context:
                    context_before = []
                    context_after = []

                    for j in range(max(0, i - context_window), i):
                        if segments[j].label:
                            context_before.append(segments[j].label)

                    for j in range(i + 1, min(len(segments), i + context_window + 1)):
                        if segments[j].label:
                            context_after.append(segments[j].label)

                    record["context_before"] = "|".join(g for g in context_before if g)
                    record["context_after"] = "|".join(g for g in context_after if g)

                all_records.append(record)
        except Exception:
            continue

    # Write output
    if format == "csv":
        if all_records:
            with open(output_path, "w", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(f, fieldnames=all_records[0].keys())
                writer.writeheader()
                writer.writerows(all_records)

    elif format == "json":
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(all_records, f, indent=2, ensure_ascii=False)

    elif format == "jsonl":
        with open(output_path, "w", encoding="utf-8") as f:
            for record in all_records:
                f.write(json.dumps(record, ensure_ascii=False) + "\n")

    else:
        raise UnsupportedFormatError(format, supported_formats=["csv", "json", "jsonl"])

    return output_path


def create_vocabulary_mapping(
    elan_files: Sequence[str | Path],
    min_frequency: int = 1,
    special_tokens: list[str] | None = None,
) -> dict[str, Any]:
    """
    Create a vocabulary mapping from ELAN files.

    Args:
        elan_files: List of ELAN file paths
        min_frequency: Minimum gloss frequency to include
        special_tokens: Special tokens to add (e.g., ['<PAD>', '<UNK>'])

    Returns:
        Dictionary with gloss_to_id, id_to_gloss, and frequencies
    """
    from sltk.io.elan import read_eaf

    gloss_counter: Counter[str] = Counter()

    for elan_path in elan_files:
        try:
            segments = read_eaf(elan_path)
            for seg in segments:
                if seg.label:
                    gloss_counter[seg.label] += 1
        except Exception:
            continue

    # Filter by frequency
    filtered_glosses = [g for g, c in gloss_counter.items() if c >= min_frequency]
    filtered_glosses.sort()  # Deterministic ordering

    # Build mapping
    gloss_to_id: dict[str, int] = {}
    id_to_gloss: dict[int, str] = {}

    # Add special tokens first
    if special_tokens:
        for token in special_tokens:
            idx = len(gloss_to_id)
            gloss_to_id[token] = idx
            id_to_gloss[idx] = token

    # Add glosses
    for gloss in filtered_glosses:
        idx = len(gloss_to_id)
        gloss_to_id[gloss] = idx
        id_to_gloss[idx] = gloss

    return {
        "gloss_to_id": gloss_to_id,
        "id_to_gloss": {int(k): v for k, v in id_to_gloss.items()},
        "frequencies": dict(gloss_counter),
        "vocab_size": len(gloss_to_id),
        "min_frequency": min_frequency,
        "special_tokens": special_tokens or [],
    }


def batch_extract_poses(
    video_paths: list[str | Path],
    output_dir: str | Path,
    extractor: str = "mediapipe",
    skip_existing: bool = True,
    **extractor_kwargs,
) -> dict[str, Any]:
    """
    Extract poses from multiple videos.

    Args:
        video_paths: List of video file paths
        output_dir: Output directory for pose files
        extractor: Extractor to use ('mediapipe', 'wilor', 'nlf')
        skip_existing: Skip videos with existing output
        **extractor_kwargs: Arguments passed to extractor

    Returns:
        Dictionary with results and statistics
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    successful: list[str] = []
    failed: list[dict[str, Any]] = []
    skipped: list[str] = []

    # Initialize extractor
    if extractor == "mediapipe":
        from sltk.extraction import MediaPipeConfig, MediaPipeExtractor

        config = MediaPipeConfig(**extractor_kwargs)
        ext = MediaPipeExtractor(config)
    elif extractor == "wilor":
        from sltk.extraction import WiLoRConfig, WiLoRExtractor

        config = WiLoRConfig(**extractor_kwargs)
        ext = WiLoRExtractor(config)
    else:
        raise ConfigurationError(f"Unknown extractor: {extractor}. Supported: mediapipe, wilor")

    try:
        ext.load_model()

        for video_path in video_paths:
            video_path = Path(video_path)
            output_path = output_dir / f"{video_path.stem}_{extractor}.h5"

            if skip_existing and output_path.exists():
                skipped.append(str(video_path))
                continue

            try:
                result = ext.extract_from_video(video_path, output_path)
                if result.success:
                    successful.append(str(video_path))
                else:
                    failed.append(
                        {
                            "path": str(video_path),
                            "errors": result.errors,
                        }
                    )
            except Exception as e:
                failed.append(
                    {
                        "path": str(video_path),
                        "errors": [str(e)],
                    }
                )
    finally:
        if hasattr(ext, "close"):
            ext.close()

    return {
        "total": len(video_paths),
        "successful": successful,
        "failed": failed,
        "skipped": skipped,
        "success_rate": len(successful) / max(len(video_paths), 1),
    }


def compute_pose_coverage(
    pose_path: str | Path,
    threshold: float = 0.5,
) -> dict[str, Any]:
    """
    Compute pose coverage statistics (frames with valid detections).

    Args:
        pose_path: Path to pose H5 file
        threshold: Confidence threshold for valid detection

    Returns:
        Dictionary with coverage metrics
    """
    from sltk.io.h5 import load_poses_h5

    poses = load_poses_h5(pose_path)

    # Valid frames (any non-zero keypoint)
    valid_mask = np.any(poses.data != 0, axis=(1, 2))

    # High-confidence frames
    if poses.confidence is not None:
        high_conf_mask = np.mean(poses.confidence, axis=-1) >= threshold
    else:
        high_conf_mask = valid_mask

    # Compute gaps (sequences of invalid frames)
    gaps = []
    gap_start = None
    for i, valid in enumerate(valid_mask):
        if not valid and gap_start is None:
            gap_start = i
        elif valid and gap_start is not None:
            gaps.append(i - gap_start)
            gap_start = None
    if gap_start is not None:
        gaps.append(len(valid_mask) - gap_start)

    return {
        "total_frames": poses.num_frames,
        "valid_frames": int(valid_mask.sum()),
        "coverage_ratio": float(valid_mask.mean()),
        "high_conf_frames": int(high_conf_mask.sum()),
        "high_conf_ratio": float(high_conf_mask.mean()),
        "num_gaps": len(gaps),
        "avg_gap_length": float(np.mean(gaps)) if gaps else 0,
        "max_gap_length": int(max(gaps)) if gaps else 0,
        "duration_sec": poses.duration,
        "fps": poses.fps,
    }


def align_poses_to_segments(
    poses,
    segments,
    padding_frames: int = 0,
) -> Iterator[tuple[Any, np.ndarray]]:
    """
    Yield pose clips aligned to segment boundaries.

    Args:
        poses: PoseSequence object
        segments: SegmentList or list of Segment objects
        padding_frames: Number of frames to pad on each side

    Yields:
        Tuples of (segment, pose_clip)
    """
    for seg in segments:
        start_frame = max(0, int(seg.start * poses.fps) - padding_frames)
        end_frame = min(poses.num_frames, int(seg.end * poses.fps) + padding_frames)

        if end_frame > start_frame:
            pose_clip = poses.data[start_frame:end_frame]
            yield seg, pose_clip
