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: int):
    """Not meaningful for custom datasets — data access is owned by CustomDataset."""
    logger.warning(
        "CustomController.get_sample() should not be called; data access is handled by CustomDataset."
    )
    return None

def get_types(self):
    """Not meaningful for custom datasets — type info is owned by CustomDataset."""
    logger.warning(
        "CustomController.get_types() should not be called; type information is defined by the CustomDataset."
    )
    return None

def shape(self):
    """Not meaningful for custom datasets — shape is owned by CustomDataset."""
    logger.warning(
        "CustomController.shape() should not be called; shape is defined by the CustomDataset."
    )
    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: Optional[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, PathLike, 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:
        # pandas is used for flexibility in handling delimiters and index columns
        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:  # type: ignore[override]
    """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()