from math import ceil
import numpy as np
import pandas as pd
from torch.utils.data import Sampler
import torch.distributed as dist
from src.custom_args import CustomTrainingArguments
import src.custom_utils as utils
from tqdm import tqdm
from src.datasets import CustomDataset

logger = utils.get_logger(__name__)


def normal_dataset_infinite_sampler(grp_list, bsz, mu=4, sigma=1):
    def infit(x):
        np.random.shuffle(x)
        it = iter(x)
        while True:
            try:
                yield next(it)
            except StopIteration:
                np.random.shuffle(x)
                it = iter(x)

    ND = len(grp_list)
    cats = np.arange(ND)

    # init iterators
    all_iters = [infit(x) for x in grp_list]
    from itertools import cycle

    while True:
        # draw how many datasets
        D = max(1, (min(ND, round(np.random.normal(mu, sigma)))))

        # draw datasets
        didxs = np.random.choice(cats, size=D, replace=False)
        diters = cycle(didxs)

        # uniform sampling within each dataset
        idxs = [next(all_iters[next(diters)]) for _ in range(bsz)]
        yield idxs


class MbeirDatataskBalanceSampler(Sampler):

    def __init__(self, train_dataset: CustomDataset, args: CustomTrainingArguments) -> None:
        super().__init__(None)
        if 'datatask_id' in train_dataset.data:
            self.mapping = train_dataset.data
        elif 'dataset_id' in train_dataset.data:
            self.mapping = train_dataset.data[['dataset_id']]
            self.mapping.rename(columns={'dataset_id': 'datatask_id'}, inplace=True)
        else:
            with open(utils.load_dataset_paths()['mbeir_data_task_mapping']) as f:
                mapping = pd.read_json(f, orient='index')
            utils.rank0_log(
                f"Datatask mapping loaded with {len(mapping)} lines", logger.info)
            mapping.reset_index(inplace=True)
            mapping.columns = ['data_id', 'datatask_id']
            self.mapping = train_dataset.data.merge(mapping, on='data_id')[
                ['data_id', 'datatask_id']]
            utils.rank0_log('Datatask mapping created:', logger.info)
        self.args = args

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

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

    def get_indices(self):
        if utils.is_main_process():
            batch_size = dist.get_world_size() * self.args.per_device_train_batch_size
            n_batches = ceil(len(self.mapping) / batch_size)
            it = normal_dataset_infinite_sampler([x[1].index.to_numpy() for x in self.mapping.groupby(by='datatask_id', sort=False)], batch_size)
            indices = [next(it) for _ in tqdm(
                range(n_batches), desc='Sampling indices balanced by M-BEIR datatask ...')]
            np.random.shuffle(indices)
            utils.rank0_log('Batched indices shuffled', logger.info)
            indices = [y for x in indices for y in x]
            utils.rank0_log(
                f"Missing samples in epoch: {len(self.mapping) - len(set(indices))}\n",
                logger.info
            )
        else:
            indices = []

        all_indices = [None] * dist.get_world_size()
        dist.all_gather_object(all_indices, indices)
        return all_indices[0]
