from collections import abc
from loguru import logger

import pytorch_lightning as pl
from torch import distributed as dist
from torch.utils.data import DataLoader, ConcatDataset, DistributedSampler

from src.datasets.pretrain_dataset import PretrainDataset


class PretrainDataModule(pl.LightningDataModule):
    """
    For distributed training, each training process is assgined
    only a part of the training scenes to reduce memory overhead.
    """

    def __init__(self, args, config):
        super().__init__()

        # 1. data config
        # Train and Val should from the same data source
        self.train_data_source = config.DATASET.TRAIN_DATA_SOURCE
        self.val_data_source = config.DATASET.VAL_DATA_SOURCE
        # training and validating
        self.train_data_root = config.DATASET.TRAIN_DATA_ROOT
        self.val_data_root = config.DATASET.VAL_DATA_ROOT

        # 2. dataset config']

        # dataset options
        self.pretrain_img_resize = config.DATASET.PRETRAIN_IMG_RESIZE  # 840
        self.pretrain_img_pad = config.DATASET.PRETRAIN_IMG_PAD  # True
        self.pretrain_df = config.DATASET.PRETRAIN_DF  # 8
        self.coarse_scale = 1 / config.XOFTR.RESOLUTION[0]  # 0.125. for training xoftr.
        self.frame_gap = config.DATASET.PRETRAIN_FRAME_GAP

        # 3.loader parameters
        self.train_loader_params = {
            "batch_size": args.batch_size,
            "num_workers": args.num_workers,
            "pin_memory": getattr(args, "pin_memory", True),
        }
        self.val_loader_params = {
            "batch_size": 1,
            "shuffle": False,
            "num_workers": args.num_workers,
            "pin_memory": getattr(args, "pin_memory", True),
        }

    def setup(self, stage=None):
        """
        Setup train / val / test dataset. This method will be called by PL automatically.
        Args:
            stage (str): 'fit' in training phase, and 'test' in testing phase.
        """

        assert stage in ["fit", "test"], "stage must be either fit or test"

        try:
            self.world_size = dist.get_world_size()
            self.rank = dist.get_rank()
            logger.info(f"[rank:{self.rank}] world_size: {self.world_size}")
        except AssertionError as ae:
            self.world_size = 1
            self.rank = 0
            logger.warning(str(ae) + " (set wolrd_size=1 and rank=0)")

        if stage == "fit":
            self.train_dataset = self._setup_dataset(self.train_data_root, mode="train")
            # setup multiple (optional) validation subsets
            self.val_dataset = []
            self.val_dataset.append(self._setup_dataset(self.val_data_root, mode="val"))
            logger.info(f"[rank:{self.rank}] Train & Val Dataset loaded!")
        else:  # stage == 'test
            raise ValueError(f"only 'fit' implemented")

    def _setup_dataset(self, data_root, mode="train"):
        """Setup train / val / test set"""

        dataset_builder = self._build_concat_dataset
        return dataset_builder(data_root, mode=mode)

    def _build_concat_dataset(self, data_root, mode):
        datasets = []

        datasets.append(
            PretrainDataset(
                data_root,
                mode=mode,
                img_resize=self.pretrain_img_resize,
                df=self.pretrain_df,
                img_padding=self.pretrain_img_pad,
                coarse_scale=self.coarse_scale,
                frame_gap=self.frame_gap,
            )
        )

        return ConcatDataset(datasets)

    def train_dataloader(self):
        """Build training dataloader for KAIST dataset."""
        sampler = DistributedSampler(self.train_dataset, shuffle=True)
        dataloader = DataLoader(self.train_dataset, sampler=sampler, **self.train_loader_params)
        return dataloader

    def val_dataloader(self):
        """Build validation dataloader KAIST dataset."""
        if not isinstance(self.val_dataset, abc.Sequence):
            return DataLoader(self.val_dataset, sampler=sampler, **self.val_loader_params)
        else:
            dataloaders = []
            for dataset in self.val_dataset:
                sampler = DistributedSampler(dataset, shuffle=False)
                dataloaders.append(DataLoader(dataset, sampler=sampler, **self.val_loader_params))
            return dataloaders
