"""Shared helpers for classification evaluation flows."""

from __future__ import annotations

from pathlib import Path

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms


def build_imagefolder_loaders(
    data_root: str,
    batch_size: int,
    workers: int,
    *,
    image_size: int = 224,
    resize_size: int = 256,
    include_train_loader: bool = True,
) -> tuple[torch.utils.data.DataLoader | None, torch.utils.data.DataLoader | None, torch.utils.data.DataLoader | None]:
    """Build train/val/test loaders for ImageFolder directory layout.

    Args:
        data_root: Root directory that may contain ``train``, ``val``, and ``test``.
        batch_size: Batch size used for created dataloaders.
        workers: Number of worker processes for dataloading.
        image_size: Final crop size for model input.
        resize_size: Resize dimension used before center crop for eval splits.
        include_train_loader: Whether to create a train loader when directory exists.

    Returns:
        Tuple of ``(train_loader, val_loader, test_loader)`` where missing splits are ``None``.
    """
    root = Path(data_root)
    traindir = root / "train"
    valdir = root / "val"
    testdir = root / "test"

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    train_transform = transforms.Compose(
        [
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
    )
    eval_transform = transforms.Compose(
        [
            transforms.Resize(resize_size),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            normalize,
        ]
    )

    train_loader = None
    if include_train_loader and traindir.exists():
        train_dataset = datasets.ImageFolder(traindir, train_transform)
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=workers,
            pin_memory=True,
        )

    val_loader = None
    if valdir.exists():
        val_dataset = datasets.ImageFolder(valdir, eval_transform)
        val_loader = torch.utils.data.DataLoader(
            val_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=workers,
            pin_memory=True,
        )

    test_loader = None
    if testdir.exists():
        test_dataset = datasets.ImageFolder(testdir, eval_transform)
        test_loader = torch.utils.data.DataLoader(
            test_dataset,
            batch_size=batch_size,
            shuffle=False,
            num_workers=workers,
            pin_memory=True,
        )

    return train_loader, val_loader, test_loader
