__all__ = ["MOVEDataset", "make_dataset", "make_dataloader", "split_samples"]

from typing import Optional

import numpy as np

import torch

from torch.utils.data import DataLoader, TensorDataset

from move.core.typing import BoolArray, FloatArray

class MOVEDataset(TensorDataset):

    """

    Characterizes a dataset for PyTorch

    Args:

        cat_all:

            categorical input matrix (N_samples, N_variables x N_max-classes.

        con_all:

            normalized continuous input matrix (N_samples, N_variables).

        cat_shapes:

            list of tuples corresponding to number of features (N_variables,

            N_max-classes) of each categorical class.

        con_shapes:

            list of tuples corresponding to number of features

            (N_variables) of each continuous class.

    Raises:

        ValueError:

            Number of samples between categorical and continuous datasets must

            match.

        ValueError:

            Categorical and continuous data cannot be both empty.

    """

    def __init__(

        self,

        cat_all: Optional[torch.Tensor] = None,

        con_all: Optional[torch.Tensor] = None,

        cat_shapes: Optional[list[tuple[int, ...]]] = None,

        con_shapes: Optional[list[int]] = None,

    ) -> None:

        # Check

        num_samples = None if cat_all is None else cat_all.shape[0]

        if con_all is not None:

            if num_samples is None:

                num_samples = con_all.shape[0]

            elif num_samples != con_all.shape[0]:

                raise ValueError(

                    "Number of samples between categorical and continuous "

                    "datasets must match."

                )

        elif num_samples is None:

            raise ValueError("Categorical and continuous data cannot be both empty.")

        self.num_samples = num_samples

        self.cat_all = cat_all

        self.cat_shapes = cat_shapes

        self.con_all = con_all

        self.con_shapes = con_shapes

    def __len__(self) -> int:

        return self.num_samples

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:

        cat_slice = torch.empty(0) if self.cat_all is None else self.cat_all[idx]

        con_slice = torch.empty(0) if self.con_all is None else self.con_all[idx]

        return cat_slice, con_slice

def concat_cat_list(

    cat_list: list[FloatArray],

) -> tuple[list[tuple[int, ...]], FloatArray]:

    """

    Concatenate a list of categorical data

    Args:

        cat_list: list with each categorical class data

    Returns:

        (tuple): a tuple containing:

            cat_shapes:

                list of categorical data classes shapes (N_variables,

                 N_max-classes)

            cat_all (FloatArray):

                2D array of concatenated patients categorical data

    """

    cat_shapes = []

    cat_flat = []

    for cat in cat_list:

        cat_shape = (cat.shape[1], cat.shape[2])

        cat_shapes.append(cat_shape)

        cat_flat.append(cat.reshape(cat.shape[0], -1))

    cat_all = np.concatenate(cat_flat, axis=1)

    return cat_shapes, cat_all

def concat_con_list(

    con_list: list[FloatArray],

) -> tuple[list[int], FloatArray]:

    """

    Concatenate a list of continuous data

    Args:

        con_list: list with each continuous class data

    Returns:

        (tuple): a tuple containing:

            n_con_shapes:

                list of continuous data classes shapes (in 1D) (N_variables)

            con_all:

                2D array of concatenated patients continuous data

    """

    con_shapes = [con.shape[1] for con in con_list]

    con_all: FloatArray = np.concatenate(con_list, axis=1)

    return con_shapes, con_all

def make_dataset(

    cat_list: Optional[list[FloatArray]] = None,

    con_list: Optional[list[FloatArray]] = None,

    mask: Optional[BoolArray] = None,

) -> MOVEDataset:

    """Creates a dataset that combines categorical and continuous datasets.

    Args:

        cat_list:

            List of categorical datasets (`num_samples` x `num_features`

            x `num_categories`). Defaults to None.

        con_list:

            List of continuous datasets (`num_samples` x `num_features`).

            Defaults to None.

        mask:

            Boolean array to mask samples. Defaults to None.

    Raises:

        ValueError: If both inputs are None

    Returns:

        MOVEDataset

    """

    if not cat_list and not con_list:

        raise ValueError("At least one type of data must be in the input")

    cat_shapes, cat_all = [], None

    if cat_list:

        cat_shapes, cat_all = concat_cat_list(cat_list)

    con_shapes, con_all = [], None

    if con_list:

        con_shapes, con_all = concat_con_list(con_list)

    if cat_all is not None:

        cat_all = torch.from_numpy(cat_all)

        if mask is not None:

            cat_all = cat_all[mask]

    if con_all is not None:

        con_all = torch.from_numpy(con_all)

        if mask is not None:

            con_all = con_all[mask]

    return MOVEDataset(cat_all, con_all, cat_shapes, con_shapes)

def make_dataloader(

    cat_list: Optional[list[FloatArray]] = None,

    con_list: Optional[list[FloatArray]] = None,

    mask: Optional[BoolArray] = None,

    **kwargs

) -> DataLoader:

    """Creates a DataLoader that combines categorical and continuous datasets.

    Args:

        cat_list:

            List of categorical datasets (`num_samples` x `num_features`

            x `num_categories`). Defaults to None.

        con_list:

            List of continuous datasets (`num_samples` x `num_features`).

            Defaults to None.

        mask:

            Boolean array to mask samples. Defaults to None.

        **kwargs:

            Arguments to pass to the DataLoader (e.g., batch size)

    Raises:

        ValueError: If both inputs are None

    Returns:

        DataLoader

    """

    dataset = make_dataset(cat_list, con_list, mask)

    return DataLoader(dataset, **kwargs)

def split_samples(

    num_samples: int,

    train_frac: float,

) -> BoolArray:

    """Generate mask to randomly split samples into training and test sets.

    Args:

        num_samples: Number of samples to split.

        train_frac: Fraction of samples corresponding to training set.

    Returns:

        Boolean array to mask test samples.

    """

    sample_ids = np.arange(num_samples)

    train_size = int(train_frac * num_samples)

    rng = np.random.default_rng()

    train_ids = rng.permutation(sample_ids)[:train_size]

    return np.isin(sample_ids, train_ids)