"""
Dataset utilities untuk loading dan preprocessing.

Membaca dataset .npy menjadi Dataset PyTorch sesuai README.
"""

import numpy as np
import torch
from torch.utils.data import Dataset
from pathlib import Path
from typing import Optional, List, Tuple


class LandmarkDataset(Dataset):
    """
    PyTorch Dataset untuk landmark sequences.

    Sesuai README:
    - Input shape: (30, feature_size) = (30, 306)
    """

    def __init__(
        self,
        data_dir: Path,
        category: str,
        split: str = "train"
    ):
        """
        Inisialisasi dataset.

        Args:
            data_dir: Root directory (dataset/train/ atau dataset/test/)
            category: Nama kategori
            split: 'train' atau 'test'
        """
        # Path sesuai README: dataset/{train|test}/{category}/{label}/
        self.data_dir = data_dir / split / category
        self.samples = []
        self.labels = []
        self.label_to_idx = {}

        self._load_data()

    def _load_data(self):
        """Load semua data."""
        if not self.data_dir.exists():
            return

        # Build label mapping dari subdirectories
        label_dirs = sorted([d for d in self.data_dir.iterdir() if d.is_dir()])
        self.label_to_idx = {d.name: idx for idx, d in enumerate(label_dirs)}

        # Load samples
        for label_dir in label_dirs:
            label_idx = self.label_to_idx[label_dir.name]

            for npy_file in label_dir.glob("*.npy"):
                try:
                    sequence = np.load(npy_file)
                    self.samples.append(sequence)
                    self.labels.append(label_idx)
                except Exception as e:
                    print(f"Error loading {npy_file}: {e}")

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

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        """
        Get satu sample.

        Returns:
            Tuple (features, label)
        """
        sequence = self.samples[idx]
        label = self.labels[idx]

        return torch.FloatTensor(sequence), torch.LongTensor([label]).squeeze()


def load_landmark_sequence(sequence_path: Path) -> np.ndarray:
    """
    Load landmark sequence dari .npy file.

    Args:
        sequence_path: Path ke file .npy

    Returns:
        Array shape (num_frames, feature_size)
    """
    return np.load(sequence_path)


def get_category_labels(category: str, data_dir: Path, split: str = "train") -> List[str]:
    """
    Get semua label untuk kategori.

    Args:
        category: Nama kategori
        data_dir: Root data directory
        split: 'train' atau 'test'

    Returns:
        List of label names
    """
    split_dir = data_dir / split / category

    if not split_dir.exists():
        return []

    labels = sorted([d.name for d in split_dir.iterdir() if d.is_dir()])
    return labels


def get_num_classes(category: str, data_dir: Path, split: str = "train") -> int:
    """
    Get jumlah kelas untuk kategori.

    Args:
        category: Nama kategori
        data_dir: Root data directory
        split: 'train' atau 'test'

    Returns:
        Jumlah kelas
    """
    return len(get_category_labels(category, data_dir, split))