Dataset

@classmethod
def from_dict(cls, data: dict):
    field_names = {f.name for f in fields(cls)}
    filtered = {k: v for k, v in data.items() if k in field_names}
    return cls(**filtered)

def to_dict(self) -> dict:
    """Convert entry to dictionary - removes None values"""
    return {k: v for k, v in asdict(self).items() if v is not None}

def complete_initialization(
    self, controller_kwargs: Dict[str, Any], to_format: DataReturnFormat
) -> None:
    """Finalize initialization of object to be able to recover items.

    Args:
        controller_kwargs: must contain a ``"root"`` key with the path to the dataset.
        to_format: expected format of data returned by ``__getitem__``.
    """

    self.path = controller_kwargs.get("root", None)
    if self.path is None:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Custom Dataset ERROR: 'root' must be provided in controller_kwargs to specify dataset location."
        )
    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

    if not isinstance(sample, tuple) or len(sample) != 2:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: get_item method must return a tuple of two elements"
            f" (data, target), but got {type(sample).__name__} with"
            f" length {len(sample) if isinstance(sample, (list, tuple)) else 'N/A'}"
        )

    data, target = sample
    self._check_type(data, "data")
    self._check_type(target, "target")

@abstractmethod
def get_item(self, index):
    """Return a (data, target) tuple for the given index.

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

@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

@abstractmethod
def complete_initialization(
    self,
    controller_kwargs: Dict[str, Any],
    to_format: DataReturnFormat,
) -> None:
    """Finalize initialization of object to be able to recover items"""
    # Recover sample and validate consistency of transforms
    pass

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,
    )

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 complete_initialization(
    self,
    controller_kwargs: Dict[str, Any],
    to_format: DataReturnFormat,
) -> None:
    """Finalize initialization of object to be able to recover items

    Args:
        controller_kwargs: arguments to create controller
        to_format: format associated to expected return format
    """
    self.to_format = to_format
    self._init_controller(controller_kwargs=controller_kwargs)

    # 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 complete_initialization(
    self,
    controller_kwargs: Dict[str, Any],
    to_format: DataReturnFormat,
) -> None:
    """Select data and target, and check if they can be converted to requested format.

    Args:
        controller_kwargs: keyword arguments for controller (not used here).
        to_format: format associated to expected return format.
    Raises:
        FedbiomedError: if there is a problem converting dataset items to requested format.
    """

    self._to_format = to_format
    self._converter = self._get_format_conversion_callable()

    if self._is_supervised:
        data, target = self._dataset[0]
    elif self._target is not None:
        data = self._dataset[0]
        target = self._target[0]
    else:
        data = self._dataset[0]
        target = None

    try:
        self._validate_format_conversion(data)
    except FedbiomedError as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB632.value}: Failed to convert dataset items to "
            f"requested format {to_format}. Details: {e}"
        ) from e

    if target is not None:
        try:
            self._validate_format_conversion(target)
        except FedbiomedError as e:
            raise FedbiomedError(
                f"{ErrorNumbers.FB632.value}: Failed to convert dataset items to "
                f"requested format {to_format}. Details: {e}"
            ) from e

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

def complete_initialization(
    self,
    controller_kwargs: Dict[str, Any],
    to_format: DataReturnFormat,
) -> None:
    """Finalize initialization of object to be able to recover items

    Args:
        controller_kwargs: arguments to create controller
        to_format: format associated to expected return format
    """
    self.to_format = to_format

    self._init_controller(controller_kwargs=controller_kwargs)

    # 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._transform,
        )