"""Define custom Dataset class for handling missing data."""

import torch
import numpy as np
import pandas as pd
from .processing import _format_cols, _format_cols_test


class Dataset(torch.utils.data.Dataset):
    """
    Custom PyTorch Dataset class for handling missing data.

    If col_types is not provided, column types are inferred from the data.

    """

    def __init__(
        self,
        data: pd.DataFrame,
        col_types: list = None,
        type_dict: dict = None,
        col_names: list[str] = None,
        original_names: list[str] = None,
        test_format: bool = False,
    ):
        super().__init__()
        if col_types is None:
            original_cols = list(data.columns)
            data, self.type_dict, self.col_types = _format_cols(data)
            self.col_names = list(data.columns)  # formatted (after dummy expansion)
            self.original_col_names = original_cols  # logical/original column names
        else:
            self.col_types = col_types
            self.type_dict = type_dict
            self.col_names = col_names  # formatted column names
            self.original_col_names = (
                original_names if original_names is not None else col_names
            )

            if test_format:
                data = _format_cols_test(
                    data, self.col_types, self.type_dict, self.col_names
                )

        mask_expand_np = ~data.isnull().to_numpy()
        data_np = data.to_numpy(dtype=np.float32, na_value=np.nan, copy=False)
        data_np[~mask_expand_np] = 0.0

        self.data = torch.as_tensor(data_np, dtype=torch.float32)
        self.mask_expand = mask_expand_np

    def __len__(self):
        return self.data.size(0)

    def __getitem__(self, index):
        return (
            self.data[index],
            self.mask_expand[index],
        )
