Dataset controller

@abstractmethod
def get_sample(self, index: int) -> Dict[str, Any]:
    """Retrieve a data sample without applying transforms"""
    pass

def get_types(self):
    """Get `type` directly from values in `dict` returned by `get_sample`"""
    return {_k: type(_v).__name__ for _k, _v in self.get_sample(0).items()}

def shape(self) -> Dict[str, Any]:
    """Get `shape` directly from values in `dict` returned by `get_sample`"""
    # Supported: int, float, dict, PIL.Image, None and obj.shape (if available)
    # This function can be overwritten for specific cases in child class
    sample = self.get_sample(0)

    if not isinstance(sample, dict):
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Expected `sample` to be a `dict`, got "
            f"{type(sample).__name__}"
        )

    output = {}
    for key, val in sample.items():
        if hasattr(val, "shape"):
            output[key] = val.shape
        elif isinstance(val, (int, float)):
            output[key] = 1
        elif isinstance(val, dict):
            output[key] = len(val)
        elif isinstance(val, Image):
            output[key] = {"size": val.size, "mode": val.mode}
        elif val is None:
            output[key] = None
        else:
            raise FedbiomedError(
                f"{ErrorNumbers.FB632.value}: Not possible to get shape for value "
                f"of key: {key} that is type {type(val).__name__}"
            )
    return output

def validate(self) -> None:
    """Validates coherence of controller

    Raises:
        FedbiomedError: if coherence issue is found
    """
    return None

def get_sample(self, index):
    """Shouldn't be called as the custom controller does not need to get a sample.

    It is handled by get_item in CustomDataset.
    """
    return None

def get_types(self):
    """Controller does not know how to get samples."""
    pass

def shape(self):
    """Len and shape should never be used, as the dataset implements it specifically.

    It is impossible for the custom controller to know the shape of the dataset.
    """
    return None

def get_sample(self, index: int) -> Dict[str, Any]:
    """Retrieve a data sample without applying transforms"""
    try:
        path, target = self._samples[index]
        data = self._loader(path)
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to retrieve item at index {index}"
        ) from e
    return {"data": data, "target": target}

def get_sample(self, index: int) -> Dict[str, Any]:
    """Retrieve a data sample without applying transforms"""
    try:
        data, target = self._dataset[index]
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to retrieve item at index {index}"
        ) from e
    return {"data": data, "target": target}

def available_subjects(
    self,
    subjects_from_index: Union[list, pd.Series],
    subjects_from_folder: list = None,
) -> Dict[str, str]:
    """Checks missing subject folders and missing entries in demographics

    Args:
        subjects_from_index: Given subject folder names in demographics
        subjects_from_folder: list of subject folder names

    Returns:
        Dict with next keys:
        - missing_folders: subjects in demographics absent in folder structure
        - missing_entries: subjects in folder structure absent in demographics
        - intersection: subjects present in folder structure and demographics
    """
    # Select all subject folders if it is not given
    if subjects_from_folder is None:
        subjects_from_folder = self.subjects

    return {
        # Missing subject that will cause warnings
        "missing_folders": list(
            set(subjects_from_index).difference(subjects_from_folder)
        ),
        # Missing entries that will cause errors
        "missing_entries": list(
            set(subjects_from_folder).difference(subjects_from_index)
        ),
        # Intersection
        "intersection": list(
            set(subjects_from_index).intersection(subjects_from_folder)
        ),
    }

def demographics_column_names(self, path: Union[str, Path]):
    return self.read_demographics(path).columns.values

@staticmethod
def get_dataset_type() -> DatasetTypes:
    return DatasetTypes.MEDICAL_FOLDER

def get_sample(
    self, index: int
) -> Dict[str, NiftiReader.data_type | Dict[str, Any]]:
    """Retrieve a data sample without applying transforms"""
    if self._validated is False:
        self.validate()

    sample = self._samples[index]
    try:
        data = {
            modality: NiftiReader.read(sample[modality])
            for modality in self.modalities
        }
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to retrieve item at index {index}"
        ) from e
    if "demographics" in sample:
        data["demographics"] = sample["demographics"]
    return data

@staticmethod
def read_demographics(
    tabular_file: Union[str, Path],
    index_col: Optional[Union[int, str]] = None,
) -> pd.DataFrame:
    """Read demographics tabular file

    Args:
        tabular_file: path to demographics file
        index_col: Index column that matches <subject>. Defaults to None.

    Raises:
        FedbiomedError: if the file can not be loaded

    Returns:
        Demographics in DataFrame format
    """
    tabular_file = MedicalFolderController._validate_tabular_file(tabular_file)

    try:
        demographics = CsvReader(tabular_file).data.to_pandas()
        if index_col is not None:
            if isinstance(index_col, int):
                if index_col < 0 or index_col >= len(demographics.columns):
                    raise FedbiomedError(
                        f"{ErrorNumbers.FB613.value}: "
                        f"Index column {index_col} is out of bounds"
                    )
                index_col = demographics.columns[index_col]
            demographics = demographics.set_index(index_col)

    except FedbiomedError as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB613.value}: :"
            f"Can not load demographics tabular file. Error message is: {e}"
        ) from e

    length = len(demographics)
    logger.info(f"Number of rows in demographics file: {length}")

    # Keep the first one in duplicated subjects
    demographics = demographics.loc[~demographics.index.duplicated(keep="first")]
    if length != len(demographics):
        logger.info(f"Length of demographics for unique index {len(demographics)}")

    return demographics

def set_dlp(self, dlp):
    """Ensures validation of controller object after set_dlp is executed"""
    DataLoadingPlanMixin.set_dlp(self, dlp)
    self.validate()

def subject_modality_status(self, index: Union[list, pd.Series] = None) -> Dict:
    """Scans subjects and checks which modalities exist for each subject

    Args:
        index: Array-like index that comes from reference csv file.
            It represents subject folder names. Defaults to None.
    Returns:
        Modality status that indicates which modalities are available per subject
    """
    # Pivot into wide format with boolean indicators
    df_ = (
        self.df_dir.assign(val=True)
        .pivot_table(
            index="subject",
            columns="modality",
            values="val",
            fill_value=0,
        )
        .astype(bool)
    )

    if index is not None:
        df_["in_folder"] = True
        # Merge with pivot (outer join)
        df_ = pd.merge(
            df_,
            pd.DataFrame(True, index=index, columns=["in_index"]),
            left_index=True,
            right_index=True,
            how="outer",
        )
        # Fill missing values with False
        df_ = df_.fillna(False)

    return {
        "columns": df_.columns.tolist(),
        "data": df_.values.tolist(),
        "index": df_.index.tolist(),
    }

def validate(self) -> None:
    # Filter subjects to contain all modalities
    self._modalities, df_dir = self._prepare_df_dir_for_use(self.df_dir, self._dlp)
    # Generate list of samples: dict, with demographics and path to modalities
    self._subjects, self._samples = self._make_dataset(self.demographics, df_dir)
    # Change flag to be able to recover a sample
    self._validated = True
    # Check if is possible to use `reader` to recover a valid item
    _ = self.get_sample(0)

    self._controller_kwargs = {
        "root": str(self.root),
        "tabular_file": (
            None if self.tabular_file is None else str(self.tabular_file)
        ),
        "index_col": self.index_col,
        "dlp": self._dlp,
    }

def get_sample(self, index: int) -> Dict[str, Any]:
    """Retrieve a data sample without applying transforms"""
    try:
        data, target = self._dataset[index]
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to retrieve item at index {index}"
        ) from e
    return {"data": data, "target": target}

def get_sample(self, index: int) -> pl.DataFrame:
    """Retrieve a data sample without applying transforms"""
    if index >= self.__len__():
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to retrieve item at index {index}"
        )
    return self._reader.get(index)

def get_types(self):
    """Get dtypes of the columns in the Tabular dataset"""
    return {
        col: dtype.__class__.__name__
        for col, dtype in self._reader.data.schema.items()
    }

def normalize_columns(self, columns: Union[Iterable, int, str]) -> list[str]:
    """Validate and normalize `columns` to a list of column names

    Args:
        columns: Columns to normalize

    Raises:
        FedbiomedError: if `columns` is not valid

    Returns:
        List of column names
    """
    return self._reader.normalize_columns(columns=columns)

def shape(self) -> Dict:
    return self._reader.shape()