from __future__ import annotations
from torch.utils.data import DataLoader, DistributedSampler, Sampler
from .nifti_dataset import NiftiSegDataset
from sacflow.utils.distributed import get_world_size, get_rank


class DistributedEvalSamplerNoPad(Sampler):
    """Shard evaluation data across ranks without padding/duplication.

    PyTorch's DistributedSampler pads samples so every rank has equal length.
    That is useful for training but biases validation metrics because some cases
    are duplicated. This sampler uses rank::world_size indices exactly once.
    """
    def __init__(self, dataset):
        self.dataset = dataset
        self.rank = get_rank()
        self.world_size = get_world_size()
        self.indices = list(range(self.rank, len(dataset), self.world_size))

    def __iter__(self):
        return iter(self.indices)

    def __len__(self):
        return len(self.indices)


def build_loader(cfg, split: str, training: bool, require_label: bool = False, distributed: bool | None = None):
    """Build a NIfTI segmentation loader.

    Important DDP behavior:
    - Training loaders use DistributedSampler when world_size > 1.
    - Evaluation/validation loaders default to *no* DistributedSampler. This is deliberate:
      training-time validation is run only on rank 0, and standalone eval usually uses one rank.
      Using a DistributedSampler for validation without metric all-gather biases metrics to a
      rank-local subset.
    """
    data_cfg = cfg["data"]
    ds = NiftiSegDataset(data_cfg["manifest"], split=split, cfg=data_cfg, training=training, require_label=require_label)
    if distributed is None:
        distributed = bool(training and get_world_size() > 1)
    if distributed:
        sampler = DistributedSampler(ds, shuffle=True) if training else DistributedEvalSamplerNoPad(ds)
    else:
        sampler = None
    loader = DataLoader(
        ds,
        batch_size=data_cfg.get("batch_size" if training else "val_batch_size", 1),
        shuffle=(training and sampler is None),
        sampler=sampler,
        num_workers=cfg.get("num_workers", 4),
        pin_memory=cfg.get("pin_memory", True),
        persistent_workers=cfg.get("num_workers", 4) > 0,
    )
    return loader