import time

import torch
import pytorch_lightning as pl
from overrides import overrides
from torch.utils.data.dataloader import DataLoader

from lightning_modules.data_modules.kvu_dataset import KVUDataset, KVUEmbeddingDataset
from lightning_modules.utils import _get_number_samples

class KVUDataModule(pl.LightningDataModule):
    def __init__(self, cfg, tokenizer_layoutxlm, feature_extractor):
        super().__init__()
        self.cfg = cfg
        self.train_loader = None
        self.val_loader = None
        self.tokenizer_layoutxlm = tokenizer_layoutxlm
        self.feature_extractor = feature_extractor
        self.collate_fn = None
        
        self.backbone_type = self.cfg.model.backbone
        self.num_samples_per_epoch = _get_number_samples(cfg.dataset_root_path)
                

    @overrides
    def setup(self, stage=None):
        self.train_loader = self._get_train_loader()
        self.val_loader = self._get_val_loaders()

    @overrides
    def train_dataloader(self):
        return self.train_loader

    @overrides
    def val_dataloader(self):
        return self.val_loader

    def _get_train_loader(self):
        start_time = time.time()
        
        if self.cfg.stage == 1:
            dataset = KVUDataset(
                self.cfg,
                self.tokenizer_layoutxlm,
                self.feature_extractor,
                mode="train",
            )
        elif self.cfg.stage == 2:
            # dataset = KVUEmbeddingDataset(
            #     self.cfg,
            #     mode="train",
            # )
            dataset = KVUDataset(
                self.cfg,
                self.tokenizer_layoutxlm,
                self.feature_extractor,
                mode="train",
            )
        else:
            raise ValueError(
                f"Not supported stage: {self.cfg.stage}"
            )  
         
        print('No. training samples:', len(dataset))

        data_loader = DataLoader(
            dataset,
            batch_size=self.cfg.train.batch_size,
            shuffle=True,
            num_workers=self.cfg.train.num_workers,
            pin_memory=True,
        )

        elapsed_time = time.time() - start_time
        print(f"Elapsed time for loading training data: {elapsed_time}")

        return data_loader

    def _get_val_loaders(self):
        
        if self.cfg.stage == 1:
            dataset = KVUDataset(
                self.cfg,
                self.tokenizer_layoutxlm,
                self.feature_extractor,
                mode="val",
            )
        elif self.cfg.stage == 2:
            # dataset = KVUEmbeddingDataset(
            #     self.cfg,
            #     mode="val",
            # )
            dataset = KVUDataset(
                self.cfg,
                self.tokenizer_layoutxlm,
                self.feature_extractor,
                mode="val",
            )
        else:
            raise ValueError(
                f"Not supported stage: {self.cfg.stage}"
            )  
        
        print('No. validation samples:', len(dataset))

        data_loader = DataLoader(
            dataset,
            batch_size=self.cfg.val.batch_size,
            shuffle=False,
            num_workers=self.cfg.val.num_workers,
            pin_memory=True,
            drop_last=False,
        )

        return data_loader

    @overrides
    def transfer_batch_to_device(self, batch, device, dataloader_idx):
        if isinstance(batch, list):
            for sub_batch in batch:
                for k in sub_batch.keys():
                    if isinstance(sub_batch[k], torch.Tensor):
                        sub_batch[k] = sub_batch[k].to(device)
        else:
            for k in batch.keys():
                if isinstance(batch[k], torch.Tensor):
                    batch[k] = batch[k].to(device)
        return batch
    

        
        
        
    
    
