"""Data loaders — CSV, images, and text in one line."""

from __future__ import annotations

from collections import Counter
from pathlib import Path

import torch


def load_csv(
    path: str | Path,
    target: str,
    features: list[str] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Load a CSV file and return (X, y) tensors.

    Args:
        path: Path to the CSV file.
        target: Name of the target/label column.
        features: Columns to use as features (default: all except target).

    Returns:
        Tuple of (X, y) as float tensors.
    """
    import pandas as pd

    df = pd.read_csv(path)
    y = torch.tensor(df[target].values, dtype=torch.long)

    if features:
        X = torch.tensor(df[features].values, dtype=torch.float32)
    else:
        X = torch.tensor(df.drop(columns=[target]).values, dtype=torch.float32)

    return X, y


def load_images(
    directory: str | Path,
    image_size: int = 224,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Load images from a directory organized by class folders.

    Expected structure::

        directory/
            class_a/
                img1.jpg
                img2.jpg
            class_b/
                img3.jpg

    Args:
        directory: Root directory with class subfolders.
        image_size: Resize images to this square size.

    Returns:
        Tuple of (images, labels) tensors.
    """
    from torchvision import transforms
    from torchvision.datasets import ImageFolder

    transform = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    dataset = ImageFolder(str(directory), transform=transform)
    images = torch.stack([img for img, _ in dataset])
    labels = torch.tensor([label for _, label in dataset])
    return images, labels


def load_text(
    path: str | Path,
    text_col: str,
    label_col: str,
    vocab_size: int = 10000,
    max_len: int = 256,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Load text data from a CSV and return tokenized (X, y) tensors.

    Uses a simple whitespace tokenizer with the most frequent words.

    Args:
        path: Path to the CSV file.
        text_col: Name of the text column.
        label_col: Name of the label column.
        vocab_size: Max vocabulary size.
        max_len: Pad/truncate sequences to this length.

    Returns:
        Tuple of (token_ids, labels) tensors.
    """
    import pandas as pd

    df = pd.read_csv(path)
    texts = df[text_col].astype(str).tolist()
    labels = torch.tensor(df[label_col].values, dtype=torch.long)

    # Build vocabulary from most common words
    word_counts: Counter[str] = Counter()
    for text in texts:
        word_counts.update(text.lower().split())

    vocab = {word: i + 1 for i, (word, _) in enumerate(word_counts.most_common(vocab_size - 1))}
    # 0 is reserved for padding

    # Tokenize and pad
    sequences = []
    for text in texts:
        tokens = [vocab.get(w, 0) for w in text.lower().split()]
        if len(tokens) >= max_len:
            tokens = tokens[:max_len]
        else:
            tokens = tokens + [0] * (max_len - len(tokens))
        sequences.append(tokens)

    X = torch.tensor(sequences, dtype=torch.long)
    return X, labels
