""" Data module for pytorch lightning """

from glob import glob

from lightning.pytorch import LightningDataModule
from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
from ocf_data_sampler.sample.base import (
    NumpyBatch,
    SampleBase,
    TensorBatch,
    batch_to_tensor,
)
from torch.utils.data import DataLoader, Dataset


def collate_fn(samples: list[NumpyBatch]) -> TensorBatch:
    """Convert a list of NumpySample samples to a tensor batch"""
    return batch_to_tensor(stack_np_samples_into_batch(samples))


class PremadeSamplesDataset(Dataset):
    """Dataset to load samples from

    Args:
        sample_dir: Path to the directory of pre-saved samples.
        sample_class: sample class type to use for save/load/to_numpy
    """

    def __init__(self, sample_dir: str, sample_class: SampleBase):
        """Initialise PremadeSamplesDataset"""
        self.sample_paths = glob(f"{sample_dir}/*")
        self.sample_class = sample_class

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

    def __getitem__(self, idx):
        sample = self.sample_class.load(self.sample_paths[idx])
        return sample.to_numpy()


class BaseDataModule(LightningDataModule):
    """Base Datamodule for training pvnet and using pvnet pipeline in ocf-data-sampler."""

    def __init__(
        self,
        configuration: str | None = None,
        sample_dir: str | None = None,
        batch_size: int = 16,
        num_workers: int = 0,
        prefetch_factor: int | None = None,
        train_period: list[str | None] = [None, None],
        val_period: list[str | None] = [None, None],
    ):
        """Base Datamodule for training pvnet architecture.

        Can also be used with pre-made batches if `sample_dir` is set.

        Args:
            configuration: Path to ocf-data-sampler configuration file.
            sample_dir: Path to the directory of pre-saved samples. Cannot be used together with
                `configuration` or '[train/val]_period'.
            batch_size: Batch size.
            num_workers: Number of workers to use in multiprocess batch loading.
            prefetch_factor: Number of data will be prefetched at the end of each worker process.
            train_period: Date range filter for train dataloader.
            val_period: Date range filter for val dataloader.

        """
        super().__init__()

        if not ((sample_dir is not None) ^ (configuration is not None)):
            raise ValueError("Exactly one of `sample_dir` or `configuration` must be set.")

        if sample_dir is not None:
            if any([period != [None, None] for period in [train_period, val_period]]):
                raise ValueError("Cannot set `(train/val)_period` with presaved samples")

        self.configuration = configuration
        self.sample_dir = sample_dir
        self.train_period = train_period
        self.val_period = val_period

        self._common_dataloader_kwargs = dict(
            batch_size=batch_size,
            sampler=None,
            batch_sampler=None,
            num_workers=num_workers,
            collate_fn=collate_fn,
            pin_memory=False,
            drop_last=False,
            timeout=0,
            worker_init_fn=None,
            prefetch_factor=prefetch_factor,
            persistent_workers=False,
        )

    def _get_streamed_samples_dataset(self, start_time, end_time) -> Dataset:
        raise NotImplementedError

    def _get_premade_samples_dataset(self, subdir) -> Dataset:
        raise NotImplementedError

    def train_dataloader(self) -> DataLoader:
        """Construct train dataloader"""
        if self.sample_dir is not None:
            dataset = self._get_premade_samples_dataset("train")
        else:
            dataset = self._get_streamed_samples_dataset(*self.train_period)
        return DataLoader(dataset, shuffle=True, **self._common_dataloader_kwargs)

    def val_dataloader(self) -> DataLoader:
        """Construct val dataloader"""
        if self.sample_dir is not None:
            dataset = self._get_premade_samples_dataset("val")
        else:
            dataset = self._get_streamed_samples_dataset(*self.val_period)
        return DataLoader(dataset, shuffle=False, **self._common_dataloader_kwargs)
