import numpy as np
from torch.utils.data import DataLoader
from .PIFuDataset import PIFuDataset
import pytorch_lightning as pl


class PIFuDataModule(pl.LightningDataModule):
    def __init__(self, cfg):
        super(PIFuDataModule, self).__init__()
        self.cfg = cfg
        self.overfit = self.cfg.overfit

        if self.overfit:
            self.batch_size = 1
        else:
            self.batch_size = self.cfg.batch_size

        self.data_size = {}

    def prepare_data(self):

        pass

    @staticmethod
    def worker_init_fn(worker_id):
        np.random.seed(np.random.get_state()[1][0] + worker_id)

    def setup(self, stage):

        if stage == 'fit':
            self.train_dataset = PIFuDataset(cfg=self.cfg, split="train")
            self.val_dataset = PIFuDataset(cfg=self.cfg, split="val")
            self.data_size = {'train': len(self.train_dataset),
                              'val': len(self.val_dataset)}

        if stage == 'test':
            self.test_dataset = PIFuDataset(cfg=self.cfg, split="test")

    def train_dataloader(self):

        train_data_loader = DataLoader(
            self.train_dataset,
            batch_size=self.batch_size, shuffle=True,
            num_workers=self.cfg.num_threads, pin_memory=True,
            worker_init_fn=self.worker_init_fn)

        return train_data_loader

    def val_dataloader(self):

        if self.overfit:
            current_dataset = self.train_dataset
        else:
            current_dataset = self.val_dataset

        val_data_loader = DataLoader(
            current_dataset,
            batch_size=1, shuffle=False,
            num_workers=self.cfg.num_threads, pin_memory=True,
            worker_init_fn=self.worker_init_fn)

        return val_data_loader

    def test_dataloader(self):

        test_data_loader = DataLoader(
            self.test_dataset,
            batch_size=1, shuffle=False,
            num_workers=self.cfg.num_threads, pin_memory=True)

        return test_data_loader