from torch.utils.data import Dataset

from custom.dataset.anydata_dataset import AnyDataset
from imaginaire.utils import log


class RankDataset(Dataset):
    """Per-GPU dataset assignment for multi-dataset training.

    Given a list of dataset config paths, each GPU loads only one dataset
    based on its data-parallel rank (round-robin assignment).
    Reports infinite length; the trainer's max_iter controls when to stop.
    Shorter datasets wrap around via modulo.

    Usage in experiment config:
        dataset_train = dict(
            config=[
                'custom/config/anydata/local/train/driving.yaml',
                'custom/config/anydata/local/train/robotics.yaml',
            ],
            ...
        )
    """

    def __init__(
        self,
        configs,
        phase='train',
        data_defaults=None,
        data_overrides=None,
        rank=None,
        world_size=None,
        shuffle=None,
    ):
        super().__init__()
        from megatron.core import parallel_state

        if rank is None:
            rank = parallel_state.get_data_parallel_rank()
        if world_size is None:
            world_size = parallel_state.get_data_parallel_world_size()

        num_datasets = len(configs)
        assigned_idx = rank % num_datasets
        assigned_config = configs[assigned_idx]

        log.info(
            f'[RankDataset] rank {rank}/{world_size}, '
            f'assigned dataset {assigned_idx}/{num_datasets}: {assigned_config}'
        )

        self.dataset = AnyDataset(
            dataset_dir=assigned_config,
            phase=phase,
            data_defaults=data_defaults,
            data_overrides=data_overrides,
            single_gpu=True,  # each rank has its own dataset, don't split shards by rank
            shuffle=shuffle,
        )
        self._local_len = len(self.dataset)

        # Surface dataset sources in the wandb run's Config tab for run-level audit
        # (one entry per (phase, dataset_idx))
        if rank == 0:
            try:
                import wandb
                if wandb.run is not None:
                    s3_paths = [
                        p[0] if isinstance(p, (list, tuple)) and len(p) == 2 else p
                        for p in self.dataset.config.get('path', [])
                    ]
                    wandb.config.update(
                        {
                            f'dataset/{phase}/yaml_{assigned_idx}': assigned_config,
                            f'dataset/{phase}/s3_paths_{assigned_idx}': s3_paths,
                        },
                        allow_val_change=True,
                    )
            except Exception as e:
                log.warning(f'[RankDataset] could not push paths to wandb.config: {e}')

    def __len__(self):
        # NOTE(bvh): 100M is large but bounded (2^31 would hang DistributedSampler).
        # NOTE(bvh): keep this value consistent between WebbedDataset (AnyData) and RankDataset (Any4D / custom).
        return 100_000_000

    def __getitem__(self, index):
        return self.dataset[index % self._local_len]

    def __str__(self):
        return f'RankDataset({self.dataset})'