# -----------------------------------------------------------------------------
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
#
# This codebase constitutes NVIDIA proprietary technology and is strictly
# confidential. Any unauthorized reproduction, distribution, or disclosure
# of this code, in whole or in part, outside NVIDIA is strictly prohibited
# without prior written consent.
#
# For inquiries regarding the use of this code in other NVIDIA proprietary
# projects, please contact the Deep Imagination Research Team at
# dir@exchange.nvidia.com.
# -----------------------------------------------------------------------------

# This script contains the code for multi-aspect ratio shard iterator

import random
import time
from collections import defaultdict

import numpy as np
from webdataset.pytorch import IterableDataset
from webdataset.utils import pytorch_worker_info

from cosmos_rl.tools.dataset.wfm.webdataset.config.schema import TarSample
from cosmos_rl.utils.logging import logger


class ShardlistMultiAspectRatioInfinite(IterableDataset):
    r"""
    An iterable dataset that parses and yields tar files.
    This distributor handles the multi-aspect ratio case. For the dataloader to be successful,
    each worker should load only one aspect ratio. Else, there can be a batch where two
    aspect ratios would be present which would raise an error in collate function.
    So, we design data distribution strategy so that each worker sees only one aspect ratio.

    This version only supports infinite loader mode. This enables a simpler code that is faster to initialize
    and produces samples better matching the dataset distribution.
    """

    def __init__(
        self,
        shuffle: bool = True,
        split_by_node: bool = True,
        split_by_worker: bool = True,
        chunk_size: int = 1,
        resume_flag: bool = True,
        verbose: bool = False,
        is_infinite_loader: bool = True,
    ):
        r"""Create a multi-aspect ratio ShardList.
        Args:
            urls (list[TarSample]): a list of tar files along with their metadata
            epoch_shuffle (bool): Shuffles the whole epoch. If disabled, each node will see the same set of urls.
            shuffle (bool): shuffle samples before iterating.
            split_by_node (bool): split shards by node if True
            split_by_worker (bool): split shards by worker if True
            chunk_size (int): Ignored
            resume_flag (bool): Ignored
            verbose (bool): Prints some logs if true
            is_infinite_loader (bool): If true, creates an infinite dataloader.
                So, the dataset will be only one epoch and will not terminate.
        """
        super().__init__()

        self.verbose = verbose
        if self.verbose:
            logger.info("ShardlistMultiAspectRatioInfinite init")
        self.shuffle = shuffle
        self.split_by_node = split_by_node
        self.split_by_worker = split_by_worker
        self.chunk_size = chunk_size  # Ignored
        self.resume_flag = resume_flag  # Ignored
        assert is_infinite_loader is True

    def set_urls(self, urls: list[TarSample]):
        self.url_aspect_split = self._split_urls_by_aspect_ratio(urls)

    def set_chunk_size(self, chunk_size: int):
        """Set chunk size
        For backward compatibility. Ignored.

        Args:
            chunk_size (int): chunk size used in webdataset creation
        """
        self.chunk_size = chunk_size

    def set_epoch(self, epoch: int, start_index: int):
        r"""Set the current epoch. Used for per-node shuffling.
        For backward compatibility. Ignored.

        Args:
            epoch (int): Epoch number
            start_index (int): iteraton number
        """
        self.epoch = epoch
        self.start_index = start_index

    def _split_urls_by_aspect_ratio(self, urls):
        r"""Function for splitting urls by aspect ratio.
        We assume that urls are grouped by dataset_id. That is, data belonging to
        one dataset_id should have all data in the same aspect ratio.
        """

        url_aspect_split = defaultdict(list)

        for url in urls:
            dset_info = url.meta
            if "aspect_ratio" not in dset_info.opts:
                raise ValueError(
                    "aspect_ratio should be specified in dataset_info when using multi aspect distributor"
                )
            aspect_ratio = dset_info.opts["aspect_ratio"]
            url_aspect_split[aspect_ratio].append(url)

        for aspect_ratio in url_aspect_split:
            # Sort the url list
            url_aspect_split[aspect_ratio] = sorted(
                url_aspect_split[aspect_ratio], key=lambda tar: (tar.path, tar.root)
            )

        return url_aspect_split

    def _allocate_workers_to_aspects(
        self, url_aspect_split: dict[str, list[TarSample]], num_workers_all: int
    ) -> list[tuple[str, int]]:
        r"""Allocate workers to each aspect ratio so that:
        1. Each aspect ratio has at least one worker
        2. All the workers have jobs to do

        Args:
            url_aspect_split (dict[list[TarSample]]): TarSample split by aspect ratio
            num_workers_all (int): Total number of dataloader workers

        Returns:
            aspect_worker_allocation (list): List of tuple containing (aspect_key, num_workers)
        """
        if self.verbose:
            logger.info(
                f"#URLs for each aspect ratio: {[len(url_aspect_split[aspect_ratio]) for aspect_ratio in url_aspect_split]}"
            )

        # Must have more global workers than the number of aspect ratios, as each global worker can only load a single
        # aspect ratio.
        num_aspects = len(url_aspect_split)
        assert num_workers_all >= num_aspects

        aspect_keys = list(url_aspect_split.keys())
        # Allocate at least one worker per aspect ratios
        target_ratio = np.array([len(url_aspect_split[key]) for key in aspect_keys])
        target_ratio = target_ratio / target_ratio.sum()
        aspect_worker_allocation = np.ones([num_aspects], dtype=np.int64)
        for _i in range(num_workers_all - num_aspects):
            current_ratio = aspect_worker_allocation / aspect_worker_allocation.sum()
            aspect_worker_allocation[np.argmin(current_ratio - target_ratio)] += 1

        if self.verbose:
            logger.info(f"Aspects: {aspect_keys}")
            logger.info(f"Target ratio: {target_ratio}")
            logger.info(f"Worker allocation: {aspect_worker_allocation}")
            logger.info(
                f"Discrepancy: {aspect_worker_allocation / aspect_worker_allocation.sum() / target_ratio}"
            )
        return [(k, v) for k, v in zip(aspect_keys, aspect_worker_allocation.tolist())]

    def _obtain_node_worker_url_mapping(
        self,
        url_aspect_split: dict[str, list[TarSample]],
        aspect_worker_allocation: list[tuple[str, int]],
        rank: int,
        world_size: int,
        worker_id: int,
        num_workers: int,
    ):
        r"""This function obtains the worker-URL mapping. It assigns the tar list seen by
        each workers.

        Args:
            url_aspect_split (dict[list[TarSample]]: TarSample split by aspect ratio
            aspect_worker_allocation (dict): Number of workers allocated to each aspect ratio
            rank (int): Rank of the current GPU
            world_size (int): Total number of GPUs
            worker_id (int): ID for the current worker in the dataloader
            num_workers (int): Total number of workers in the dataloader

        Returns:
            URL list for the current worker
        """
        assert self.split_by_node is True and self.split_by_worker is True

        # First determine the aspect ratio for the current worker
        global_worker_id = rank * num_workers + worker_id

        cumulative = 0
        for aspect_key, worker_count in aspect_worker_allocation:
            cumulative += worker_count
            if global_worker_id < cumulative:
                chunk_id = global_worker_id - cumulative + worker_count
                break

        if self.verbose:
            logger.info(
                f"GID={global_worker_id}, aspect_key={aspect_key}, chunk_id={chunk_id}"
            )
        # chunk the urls for the target aspect ratio
        urls_asp = url_aspect_split[aspect_key]
        url_chunk = urls_asp[chunk_id::worker_count]

        return url_chunk

    def obtain_url_list(self):
        r"""Return an iterator over the shards."""

        rank, world_size, worker_id, num_workers = pytorch_worker_info()

        # Splitting the shards by worker and node
        if self.verbose:
            logger.info(f"PytorchShardList rank {rank} of {world_size}")
            logger.info(f"PytorchShardList worker {worker_id} of {num_workers}")

        nworkers_all = world_size * num_workers

        # Assigning workers to process each aspect ratio
        aspect_worker_allocation = self._allocate_workers_to_aspects(
            self.url_aspect_split, nworkers_all
        )

        # Form a mapping of url_aspect_split to node and workers
        urls = self._obtain_node_worker_url_mapping(
            self.url_aspect_split,
            aspect_worker_allocation,
            rank,
            world_size,
            worker_id,
            num_workers,
        )

        if self.verbose:
            logger.info("List of urls (before shuffle)")
            logger.info(urls[0:10])

        if self.shuffle:
            global_worker_id = rank * num_workers + worker_id
            random.Random(global_worker_id).shuffle(urls)

        if self.verbose:
            logger.info("List of urls (after shuffle)")
            logger.info(urls[0:10])
            logger.info(f"PytorchShardList got {len(urls)} urls")

        return urls

    def __iter__(self):
        url_list = self.obtain_url_list()
        while True:
            cur_time = time.time_ns()
            random.Random(cur_time).shuffle(url_list)
            for url in url_list:
                yield dict(url=url)
