"""
Core data structures for SLTK.

These are the canonical formats used throughout the toolkit.
All loaders convert to these formats; all tools expect these formats.
"""

from __future__ import annotations

import logging
from collections.abc import Iterator
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, overload

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class PoseSequence:
    """
    A sequence of poses over time.

    This is the canonical pose format in SLTK. All pose data—regardless
    of source format—converts to/from PoseSequence.

    Attributes:
        data: Pose array of shape (T, N, C) where T=frames, N=keypoints, C=coordinates
        fps: Frames per second
        format: Source format identifier (e.g., 'mediapipe', 'wilor', 'smplx')
        confidence: Optional confidence scores, shape (T, N) or (T,)
        metadata: Optional additional information (video_id, etc.)
        strict_validation: If True, raise ValueError if NaN ratio exceeds nan_threshold

    Example:
        >>> poses = PoseSequence.load("poses.h5", format="wilor", fps=25)
        >>> print(poses.duration)  # Duration in seconds
        4.0
        >>> subset = poses[10:50]  # Slice by frame
        >>> subset_time = poses.slice_time(0.5, 1.5)  # Slice by seconds
    """

    data: np.ndarray  # (T, N, C)
    fps: float
    format: str
    confidence: np.ndarray | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    strict_validation: bool = False
    nan_threshold: float = 0.05  # 5% NaN threshold for strict validation

    def __post_init__(self):
        # Validate fps
        if not isinstance(self.fps, (int, float)):
            raise TypeError(f"fps must be a number, got {type(self.fps).__name__}")
        if self.fps <= 0:
            raise ValueError(f"fps must be positive, got {self.fps}")
        if not np.isfinite(self.fps):
            raise ValueError(f"fps must be finite, got {self.fps}")

        # Validate data array
        if not isinstance(self.data, np.ndarray):
            raise TypeError(f"data must be a numpy array, got {type(self.data).__name__}")

        if self.data.ndim == 2:
            # Assume (T, N*C) -> try to reshape
            # This is a fallback; specific formats should handle properly
            logger.debug(
                "PoseSequence received 2D data with shape %s; expected 3D (T, N, C)",
                self.data.shape,
            )
        elif self.data.ndim != 3:
            raise ValueError(f"Expected 2D or 3D array, got shape {self.data.shape}")

        # Check for NaN values in pose data
        if np.any(np.isnan(self.data)):
            nan_count = np.sum(np.isnan(self.data))
            total_elements = self.data.size
            nan_ratio = nan_count / total_elements

            if self.strict_validation and nan_ratio > self.nan_threshold:
                raise ValueError(
                    f"PoseSequence data contains {nan_count} NaN values "
                    f"({nan_ratio * 100:.1f}% of data), exceeding threshold of "
                    f"{self.nan_threshold * 100:.1f}%. Set strict_validation=False "
                    f"to allow NaN values or increase nan_threshold."
                )

            logger.warning(
                "PoseSequence data contains %d NaN values (%.1f%% of data)",
                nan_count,
                nan_ratio * 100,
            )

        # Validate confidence if provided
        if self.confidence is not None:
            if not isinstance(self.confidence, np.ndarray):
                raise TypeError(
                    f"confidence must be a numpy array, got {type(self.confidence).__name__}"
                )
            # Check confidence values are in valid range
            if self.confidence.size > 0:
                min_conf = np.nanmin(self.confidence)
                max_conf = np.nanmax(self.confidence)
                if min_conf < 0 or max_conf > 1:
                    logger.warning(
                        "Confidence values outside expected [0, 1] range: min=%.3f, max=%.3f",
                        min_conf,
                        max_conf,
                    )

    @property
    def num_frames(self) -> int:
        return self.data.shape[0]

    @property
    def num_keypoints(self) -> int:
        if self.data.ndim == 3:
            return self.data.shape[1]
        return -1  # Unknown for 2D data

    @property
    def num_coords(self) -> int:
        if self.data.ndim == 3:
            return self.data.shape[2]
        return self.data.shape[1]

    @property
    def duration(self) -> float:
        """Duration in seconds."""
        return self.num_frames / self.fps

    @property
    def shape(self) -> tuple:
        """Shape of the underlying data array."""
        return self.data.shape

    def __len__(self) -> int:
        return self.num_frames

    def __getitem__(self, idx) -> PoseSequence:
        """Slice by frame index."""
        if isinstance(idx, slice):
            new_data = self.data[idx]
            new_conf = self.confidence[idx] if self.confidence is not None else None
            return PoseSequence(
                data=new_data,
                fps=self.fps,
                format=self.format,
                confidence=new_conf,
                metadata=self.metadata.copy(),
            )
        elif isinstance(idx, int):
            # Return single frame as PoseSequence with 1 frame
            new_data = self.data[idx : idx + 1]
            new_conf = self.confidence[idx : idx + 1] if self.confidence is not None else None
            return PoseSequence(
                data=new_data,
                fps=self.fps,
                format=self.format,
                confidence=new_conf,
                metadata=self.metadata.copy(),
            )
        else:
            raise TypeError(f"Invalid index type: {type(idx)}")

    def slice_time(self, start_sec: float, end_sec: float) -> PoseSequence:
        """Slice by time in seconds."""
        start_frame = int(start_sec * self.fps)
        end_frame = int(end_sec * self.fps)
        return self[start_frame:end_frame]

    def frame_to_time(self, frame: int) -> float:
        """Convert frame index to time in seconds."""
        return frame / self.fps

    def time_to_frame(self, time_sec: float) -> int:
        """Convert time in seconds to frame index."""
        return int(time_sec * self.fps)

    @classmethod
    def load(
        cls,
        path: str | Path,
        format: str,
        fps: float,
        **kwargs,
    ) -> PoseSequence:
        """
        Load poses from file.

        Args:
            path: Path to pose file (.h5, .npy, .npz, .pkl)
            format: Pose format identifier ('wilor', 'smplx', 'mediapipe')
            fps: Frames per second
            **kwargs: Additional arguments passed to format-specific loader

        Returns:
            PoseSequence instance
        """
        from sltk.data.formats import load_poses

        return load_poses(path, format=format, fps=fps, **kwargs)

    def save(self, path: str | Path, **kwargs):
        """Save poses to file."""
        from sltk.data.formats import save_poses

        save_poses(self, path, **kwargs)

    def to_format(self, target_format: str) -> PoseSequence:
        """Convert to a different pose format."""
        from sltk.data.formats import convert_poses

        return convert_poses(self, target_format)

    def numpy(self) -> np.ndarray:
        """Return raw numpy array."""
        return self.data

    def copy(self) -> PoseSequence:
        """Create a deep copy."""
        return PoseSequence(
            data=self.data.copy(),
            fps=self.fps,
            format=self.format,
            confidence=self.confidence.copy() if self.confidence is not None else None,
            metadata=self.metadata.copy(),
        )


@dataclass
class Segment:
    """
    A temporal segment with optional label.

    Attributes:
        start: Start time in seconds
        end: End time in seconds
        label: Optional label (gloss, class, etc.)
        confidence: Optional confidence score
        tier: ELAN tier name (for export)
        metadata: Additional information

    Example:
        >>> seg = Segment(start=1.5, end=2.3, label="HELLO")
        >>> print(seg.duration)
        0.8
    """

    start: float
    end: float
    label: str | None = None
    confidence: float | None = None
    tier: str = "default"
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        # Validate start and end times
        if not isinstance(self.start, (int, float)):
            raise TypeError(f"start must be a number, got {type(self.start).__name__}")
        if not isinstance(self.end, (int, float)):
            raise TypeError(f"end must be a number, got {type(self.end).__name__}")

        if not np.isfinite(self.start):
            raise ValueError(f"start must be finite, got {self.start}")
        # Note: end can be inf to represent "until end of video" when duration unknown
        if np.isnan(self.end):
            raise ValueError(f"end must not be NaN, got {self.end}")

        if self.start < 0:
            raise ValueError(f"start must be non-negative, got {self.start}")

        if self.end < self.start:
            raise ValueError(f"end ({self.end}) must be >= start ({self.start})")

        # Validate confidence if provided
        if self.confidence is not None:
            if not isinstance(self.confidence, (int, float)):
                raise TypeError(
                    f"confidence must be a number, got {type(self.confidence).__name__}"
                )
            if not np.isfinite(self.confidence):
                raise ValueError(f"confidence must be finite, got {self.confidence}")
            if self.confidence < 0 or self.confidence > 1:
                logger.warning(
                    "Segment confidence outside expected [0, 1] range: %.3f",
                    self.confidence,
                )

    @property
    def duration(self) -> float:
        return self.end - self.start

    @property
    def midpoint(self) -> float:
        return (self.start + self.end) / 2

    def overlaps(self, other: Segment) -> bool:
        """Check if this segment overlaps with another."""
        return self.start < other.end and other.start < self.end

    def contains(self, time: float) -> bool:
        """Check if a time point is within this segment."""
        return self.start <= time <= self.end

    def intersection(self, other: Segment) -> Segment | None:
        """Return overlapping region, or None if no overlap."""
        if not self.overlaps(other):
            return None
        return Segment(
            start=max(self.start, other.start),
            end=min(self.end, other.end),
            tier=self.tier,
        )

    def iou(self, other: Segment) -> float:
        """Intersection over union with another segment."""
        inter = self.intersection(other)
        if inter is None:
            return 0.0
        union_start = min(self.start, other.start)
        union_end = max(self.end, other.end)
        union_duration = union_end - union_start
        return inter.duration / union_duration

    def to_frames(self, fps: float) -> tuple[int, int]:
        """Convert to frame indices."""
        return int(self.start * fps), int(self.end * fps)

    @classmethod
    def from_frames(
        cls,
        start_frame: int,
        end_frame: int,
        fps: float,
        **kwargs,
    ) -> Segment:
        """Create segment from frame indices."""
        return cls(
            start=start_frame / fps,
            end=end_frame / fps,
            **kwargs,
        )


@dataclass
class SegmentList:
    """
    A list of segments with convenience methods.

    Example:
        >>> segments = SegmentList([Segment(0, 1, "A"), Segment(1.5, 2, "B")])
        >>> segments.to_elan("output.eaf", video_path="source.mp4")
    """

    segments: list[Segment] = field(default_factory=list)

    def __len__(self) -> int:
        return len(self.segments)

    def __iter__(self) -> Iterator[Segment]:
        return iter(self.segments)

    @overload
    def __getitem__(self, idx: int) -> Segment: ...
    @overload
    def __getitem__(self, idx: slice) -> SegmentList: ...
    def __getitem__(self, idx: int | slice) -> Segment | SegmentList:
        if isinstance(idx, slice):
            return SegmentList(self.segments[idx])
        return self.segments[idx]

    def append(self, segment: Segment):
        self.segments.append(segment)

    def extend(self, segments: list[Segment]):
        self.segments.extend(segments)

    def sort(self, key=None, reverse: bool = False):
        """Sort segments in place."""
        if key is None:

            def key(s):
                return s.start

        self.segments.sort(key=key, reverse=reverse)

    def sorted(self, key=None, reverse: bool = False) -> SegmentList:
        """Return a new sorted SegmentList."""
        if key is None:

            def key(s):
                return s.start

        return SegmentList(sorted(self.segments, key=key, reverse=reverse))

    def filter_by_tier(self, tier: str) -> SegmentList:
        """Return segments from a specific tier."""
        return SegmentList([s for s in self.segments if s.tier == tier])

    def filter_by_confidence(self, min_confidence: float) -> SegmentList:
        """Return segments above confidence threshold."""
        return SegmentList(
            [s for s in self.segments if s.confidence is None or s.confidence >= min_confidence]
        )

    def filter_by_label(self, labels: str | list[str]) -> SegmentList:
        """Return segments with specific label(s)."""
        if isinstance(labels, str):
            labels = [labels]
        return SegmentList([s for s in self.segments if s.label in labels])

    def filter_by_duration(
        self, min_duration: float = 0.0, max_duration: float = float("inf")
    ) -> SegmentList:
        """Return segments within duration range."""
        return SegmentList([s for s in self.segments if min_duration <= s.duration <= max_duration])

    @property
    def tiers(self) -> list[str]:
        """List of unique tier names."""
        return list(sorted(set(s.tier for s in self.segments)))

    @property
    def labels(self) -> list[str | None]:
        """List of labels in order."""
        return [s.label for s in self.segments]

    @property
    def unique_labels(self) -> list[str]:
        """List of unique labels (excluding None)."""
        return list(sorted(set(s.label for s in self.segments if s.label is not None)))

    @property
    def total_duration(self) -> float:
        """Total duration of all segments (may overlap)."""
        return sum(s.duration for s in self.segments)

    @property
    def span(self) -> tuple[float, float]:
        """Start of first segment to end of last segment."""
        if not self.segments:
            return (0.0, 0.0)
        starts = [s.start for s in self.segments]
        ends = [s.end for s in self.segments]
        return (min(starts), max(ends))

    def to_elan(
        self,
        output_path: str | Path,
        video_path: str | Path | None = None,
        **kwargs,
    ):
        """Export to ELAN .eaf file."""
        from sltk.io.elan import write_eaf

        write_eaf(self, output_path, video_path=video_path, **kwargs)

    @classmethod
    def from_elan(cls, eaf_path: str | Path, **kwargs) -> SegmentList:
        """Load from ELAN .eaf file."""
        from sltk.io.elan import read_eaf

        return read_eaf(eaf_path, **kwargs)

    def to_dataframe(self):
        """Convert to pandas DataFrame (requires pandas)."""
        try:
            import pandas as pd
        except ImportError:
            raise ImportError("pandas required: pip install pandas")

        return pd.DataFrame(
            [
                {
                    "start": s.start,
                    "end": s.end,
                    "duration": s.duration,
                    "label": s.label,
                    "confidence": s.confidence,
                    "tier": s.tier,
                    **s.metadata,
                }
                for s in self.segments
            ]
        )


@dataclass
class Sample:
    """
    A single sample from a dataset.

    Unified container that holds all information about one data point.
    Not all fields are required—datasets populate what they have.

    Attributes:
        id: Unique identifier
        video_path: Path to source video (if available)
        poses: Pose sequence (if available/loaded)
        segments: Temporal segments with labels
        glosses: List of gloss labels
        text: Sentence-level translation
        signer_id: Signer identifier
        dataset: Dataset name
        split: Data split (train/val/test)
        metadata: Additional information
    """

    id: str
    video_path: Path | None = None
    poses: PoseSequence | None = None
    segments: SegmentList | None = None
    glosses: list[str] | None = None
    text: str | None = None
    signer_id: str | None = None
    dataset: str | None = None
    split: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def has_poses(self) -> bool:
        return self.poses is not None

    def has_video(self) -> bool:
        if self.video_path is None:
            return False
        try:
            return self.video_path.exists()
        except (OSError, PermissionError):
            return False

    def has_segments(self) -> bool:
        return self.segments is not None and len(self.segments) > 0

    def has_text(self) -> bool:
        return self.text is not None and len(self.text) > 0

    @property
    def duration(self) -> float | None:
        """Duration from poses or segments."""
        if self.poses is not None:
            return self.poses.duration
        if self.segments is not None and len(self.segments) > 0:
            return self.segments.span[1]
        return None
