Dataset

@abstractmethod
def get_item(self, index):
    """Return the sample for the given index.

    May return either ``data`` alone, or a ``(data, target)`` tuple. When
    only ``data`` is returned, the target is treated as ``None``.

    Args:
        index (int): Index of the sample to retrieve.
    """
    pass

def load(
    self,
    root: Union[str, Path],
    to_format: DataReturnFormat,
) -> None:
    """Finalize initialization of object to be able to recover items.

    Args:
        root: path to the dataset (must not be ``None``).
        to_format: expected format of data returned by ``__getitem__``.
    """

    if root is None:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Custom Dataset ERROR: 'root' must be provided to specify dataset location."
        )
    self.__path = root
    self._to_format = to_format

    # Call user defined read function to read the dataset
    try:
        self.read()
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to read "
            f"from dataset using read method. Please see error: {e}"
        ) from e

    if len(self) == 0:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Custom Dataset ERROR: dataset is empty (len == 0)."
        )

    try:
        sample = self.get_item(0)
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to retrieve item "
            f"from dataset using get_item method. Please see error: {e}"
        ) from e

    data, target = self._split_sample(sample)
    target = self._normalize_target(target)
    self._composed: dict[str, Union[bool, None]] = {
        "data": None,
        "target": None,
    }
    self._check_type(data, "data")
    if target is not None:
        self._check_type(target, "target")

@abstractmethod
def read(self) -> None:
    """Reads the dataset from the specified path.

    This method should be implemented by subclasses to load the dataset
    from the given path and prepare it for use.
    """
    pass

def apply_transforms(self, sample: Dict[str, Any]) -> Dict[str, Any]:
    """Apply transforms to sample in place

    Args:
        sample: sample returned by `self._controller.get_sample`

    Raises:
        FedbiomedError: if there is a problem applying `transform` or `target_transform`
    """
    try:
        sample["data"] = self._transform(
            self._get_default_types_callable()(
                self._get_format_conversion_callable()(sample["data"])
            )
        )
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to apply `transform` to `data` "
            f"in sample in {self._to_format.value} format."
        ) from e

    try:
        sample["data"] = self._get_default_types_callable()(sample["data"])
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to apply default training plan types to `data` "
            f"in sample in {self._to_format.value} format."
        ) from e

    if sample.get("target") is not None:
        try:
            sample["target"] = self._target_transform(
                self._get_default_types_callable()(
                    self._get_format_conversion_callable()(sample["target"])
                )
            )
        except Exception as e:
            raise FedbiomedError(
                f"{ErrorNumbers.FB632.value}: Failed to apply `target_transform` to "
                f"`target` in sample in {self._to_format.value} format."
            ) from e

        try:
            sample["target"] = self._get_default_types_callable()(sample["target"])
        except Exception as e:
            raise FedbiomedError(
                f"{ErrorNumbers.FB632.value}: Failed to apply default training plan types to `target` "
                f"in sample in {self._to_format.value} format."
            ) from e

    return sample

def complete_initialization(
    self, controller_kwargs: Dict[str, Any], to_format: DataReturnFormat
) -> None:
    """Deprecated alias for `load`.

    Preserves the legacy signature where ``root`` was passed inside the
    ``controller_kwargs`` dict; it is unpacked so it binds to the ``root``
    parameter of `load`.

    .. deprecated::
        Use `load` instead. This method will be removed in a future Fed-BioMed release.
    """
    logger.warning(
        "`complete_initialization` is deprecated and will be removed in future "
        "Fed-BioMed releases; use `load` instead."
    )
    return self.load(to_format=to_format, **controller_kwargs)

def compute_stats(
    self,
    dataset_schema: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    stats: Optional[List[str]] = None,
    stats_args: Optional[Dict[str, Any]] = None,
) -> Any:
    """Computes statistics over the dataset using the AnalyticsOrchestrator.

    Args:
        schema_args: Selection arguments to filter the schema (e.g. subset of columns/keys).
        stats: List of statistics names to compute (e.g. ['mean', 'std']).
               If None or empty, default statistics are chosen based on data type.
        stats_args: Specific arguments for statistics, structured matching the schema.

    Returns:
        Computed statistics structure.

    Raises:
        FedbiomedError: If the dataset does not support analytics (missing get_schema_for_analytics).
    """
    orchestrator = AnalyticsOrchestrator()
    return orchestrator.compute_stats(
        self,
        dataset_schema=dataset_schema,
        stats=stats,
        stats_args=stats_args,
    )

@abstractmethod
def load(
    self,
    root: Union[str, Path],
    to_format: DataReturnFormat,
    **kwargs: Any,
) -> None:
    """Finalize initialization of object to be able to recover items"""
    # Recover sample and validate consistency of transforms
    pass

def analytics_schema(self):
    """Return schema associated with federated analytics."""
    if self._controller is None:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Dataset object has not completed "
            "initialization. It is not ready to use yet."
        )
    schema = {}

    # Add demographics schema if available
    if self.demographics_columns is not None:
        schema["demographics"] = RowSpec(columns=self.demographics_columns)

    # Add image schema for all other modalities
    schema.update(
        {
            modality: ImageSpec()
            for modality in self._data_modalities
            if modality != "demographics"
        }
    )
    return schema, None

def load(
    self,
    root: Union[str, Path],
    to_format: DataReturnFormat,
    tabular_file: Optional[str] = None,
    index_col: Optional[str] = None,
    dlp: Optional[DataLoadingPlan] = None,
) -> None:
    """Finalize initialization of object to be able to recover items

    Args:
        root: path to the dataset root
        to_format: format associated to expected return format
        tabular_file: path to the CSV file holding the demographic information
        index_col: column in the tabular file holding the subject names
        dlp: data loading plan to apply
    """
    self.to_format = to_format
    self._init_controller(
        root=root,
        tabular_file=tabular_file,
        index_col=index_col,
        dlp=dlp,
    )

    # Recover sample and validate consistency of transforms
    sample = self._controller.get_sample(0)
    self._validate_format_and_transformations(
        {modality: sample[modality] for modality in self._data_modalities},
        transform=self._transform,
    )
    if self._target_modalities is not None:
        self._validate_format_and_transformations(
            {modality: sample[modality] for modality in self._target_modalities},
            transform=self._target_transform,
            is_target=True,
        )

def load(
    self,
    root: Union[str, Path],
    to_format: DataReturnFormat,
    train: bool = True,
    download: bool = True,
) -> None:
    """Finalize initialization of object to be able to recover items

    Args:
        root: path to the dataset root
        to_format: format associated to expected return format
        train: if true then train files are used
        download: if true then downloads and extracts the files if they do not exist
    """
    self.to_format = to_format
    self._init_controller(root=root, train=train, download=download)
    self._validate_initial_sample()

def analytics_schema(self):
    """Return schema for federated analytics"""
    return RowSpec(columns=self._input_columns), None

def load(
    self,
    root: Union[str, Path],
    to_format: DataReturnFormat,
) -> None:
    """Finalize initialization of object to be able to recover items

    Args:
        root: path to the dataset root
        to_format: format associated to expected return format
    """
    self.to_format = to_format

    self._init_controller(root=root)

    # Normalize columns using controller (implies validation)
    self._input_columns = self._controller.normalize_columns(self._input_columns)
    if self._target_columns is not None:
        self._target_columns = self._controller.normalize_columns(
            self._target_columns
        )
        # Check for overlap between input_columns and target_columns
        _intersection_cols = list(
            set(self._input_columns) & set(self._target_columns)
        )
        if _intersection_cols:
            logger.warning(
                f"Columns {_intersection_cols} are present in both input_columns and target_columns."
            )

    sample = self._controller.get_sample(0)  # type: ignore

    n_rows, _ = sample.shape
    if n_rows > 1:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: TabularDataset currently only supports "
            "row-wise samples. Sample obtained from controller has multiple rows."
        )

    self._validate_format_and_transformations(
        self._get_item_from_sample(sample, self._input_columns),
        transform=self._transform,
    )
    if self._target_columns is not None:
        self._validate_format_and_transformations(
            self._get_item_from_sample(sample, self._target_columns),
            transform=self._target_transform,
        )