#  Copyright © 2025 Emmi AI GmbH. All rights reserved.


import numpy as np

from noether.core.schemas.dataset import SubsetWrapperConfig
from noether.data.base import Dataset, Subset
from noether.data.base.wrapper import DatasetWrapper


class SubsetWrapper(Subset):
    """Wraps the dataset with a noether.data.Subset using indices generated by the properties from the constructor."""

    def __init__(
        self,
        dataset: Dataset | DatasetWrapper,
        config: SubsetWrapperConfig,
    ):
        """
        Args:
            dataset: The dataset to wrap.
            config: The configuration to use.

        Raises:
            ValueError: If the input parameters are invalid.
            RuntimeError: If no valid indices are provided.
        """
        if config.indices is None:
            if config.start_index is not None or config.end_index is not None:
                if config.start_percent is not None or config.end_percent is not None:
                    raise ValueError("Cannot specify both start_index/end_index and start_percent/end_percent.")
                if not isinstance(config.start_index, int | None) or not isinstance(config.end_index, int | None):
                    raise ValueError("start_index and end_index must be of type int or None.")
                end_index = config.end_index or len(dataset)
                end_index = min(end_index, len(dataset))
                start_index = config.start_index or 0
                if not (start_index <= end_index):
                    raise ValueError(
                        f"start_index ({start_index}) must be less than or equal to end_index ({end_index})."
                    )
                indices = np.arange(start_index, end_index, dtype=int)  # type: ignore
            elif config.start_percent is not None or config.end_percent is not None:
                # create indices from start/end percent
                if not (
                    config.start_percent is None
                    or (isinstance(config.start_percent, float) and 0.0 <= config.start_percent <= 1.0)
                ):
                    raise ValueError("start_percent must be of type float and between 0.0 and 1.0.")
                if not (
                    config.end_percent is None
                    or (isinstance(config.end_percent, float) and 0.0 <= config.end_percent <= 1.0)
                ):
                    raise ValueError("end_percent must be of type float and between 0.0 and 1.0.")
                start_percent = config.start_percent or 0.0
                end_percent = config.end_percent or 1.0
                if end_percent <= start_percent:
                    raise ValueError(
                        f"end_percent ({end_percent}) must be larger than start_percent ({start_percent})."
                    )
                start_index = int(start_percent * len(dataset))
                end_index = int(end_percent * len(dataset))
                indices = np.arange(start_index, end_index, dtype=int)  # type: ignore
            else:
                raise RuntimeError(
                    "Either indices or start_index/end_index or start_percent/end_percent must be provided."
                )
        else:
            if not (
                config.start_index is None
                and config.end_index is None
                and config.start_percent is None
                and config.end_percent is None
            ):
                raise ValueError(
                    "Cannot specify indices together with start_index/end_index or start_percent/end_percent."
                )
            indices = np.array(config.indices, dtype=int)  # type: ignore
        if indices is None:
            raise RuntimeError("No valid indices were provided to create the SubsetWrapper.")
        super().__init__(dataset=dataset, indices=indices)  # type: ignore

    def __str__(self) -> str:
        dataset_str = (
            str(self.dataset.__class__.__name__) if self.dataset.__str__ is object.__str__ else str(self.dataset)
        )
        if len(self.indices) == 0:
            return f"{dataset_str} (length=0)"
        return f"{dataset_str} (length={len(self.indices)} indices={self.indices[0]}-{self.indices[-1]})"
