Dataset reader

def get(
    self,
    indexes: int | Iterable,
    columns: Optional[Iterable | int | str] = None,
) -> pl.DataFrame:
    """Gets the specified rows and columns in the dataset.

    Args:
        indexes: Row indexes to retrieve.
        columns: (Optional) list of columns to retrieve.
    Returns:
        Polars DataFrame: The specified dataframe.
    """
    return self._get_entry(indexes=indexes, columns=columns)

def len(self) -> int:
    """Get number of samples"""
    return self._len

def normalize_columns(self, columns: Iterable | int | str) -> list[str]:
    """Validates `columns` and returns them in type `list`

    Args:
        columns: Columns to extract, can be an iterable of column names, a single
            column name, or an integer index.
    Returns:
        A list of column names.

    Raises:
        FedbiomedUserInputError: If the input does not match the types expected
    """

    if isinstance(columns, str) or isinstance(columns, int):
        columns = [columns]

    # if columns is a list of int, convert it to a list of column names
    # (auto-generated by polars as column_0, column_1, etc. if there is no header)
    if all(isinstance(item, int) for item in columns):
        n_cols = len(self.columns)
        _inter = list(filter(lambda x: 0 > x or x > n_cols, columns))
        if any(_inter):
            raise FedbiomedUserInputError(
                f"Column index(es) {_inter} is out of range (0 to {n_cols - 1})"
            )
        columns = [self.columns[i] for i in columns]
    _faulty_col = list(filter(lambda x: x not in self.columns, columns))
    if any(_faulty_col):
        msg = f"Cannot read columns {_faulty_col}: file does not contain those columns specified"
        raise FedbiomedUserInputError(msg)

    return columns

def shape(self):
    """Returns the shape of the csv dataset.

    Computed before applying transforms or conversion to other format.

    Returns:
        Dictionary with the shape and other necessary info for the dataset
    """
    return {"csv": self._shape}

def to_numpy(self):
    """Returns the data as a Numpy ndarray."""
    return self.data.to_numpy()

def to_pandas(self) -> pd.DataFrame:
    """Returns the data as a Pandas Dataframe."""
    return self.data.to_pandas()

def unsafe_to_torch(self):
    """This is an unsafe method that returns the data as a Torch Tensor.

    Warning: This method requires that columns have homogeneous data types. Havinng
    mixed types will raise an error.
    """
    self.data.to_torch()

def validate(self) -> None:
    """Validate the path of the CSV file.

    Raises:
        FedbiomedError: If the path is invalid
    """
    if not os.path.isfile(self._path):
        raise FedbiomedError(f"error: cannot find csv file {self._path}")

@classmethod
def read(cls, path: Union[str, Path]) -> nib.Nifti1Image:
    """Reads the NIfTI file and returns it as a tensor, optionally transformed.

    Args:
        path (Union[str, Path]): Path to the NIfTI file (.nii or .nii.gz)

    Returns:
        Union[torch.Tensor, np.ndarray, nib.Nifti1Image]: The image data in the specified format.
    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the file is not a valid NIfTI file.
    """
    # IMPORTANT: Read function does not return header. It may bne useful in the future.
    if isinstance(path, str):
        path = Path(path)
    elif not isinstance(path, Path):
        raise TypeError(f"Expected path to be a string or Path, got {type(path)}")

    path = Path(path)

    cls.validate(path)

    img = nib.load(str(path))

    return img

@staticmethod
def validate(path: Path) -> None:
    """Validate the file path and extension.

    Args:
        path (Path): Path to the NIfTI file.
    """
    if not path.exists():
        raise FileNotFoundError(f"NIfTI file does not exist: {path}")
    if not path.is_file():
        raise ValueError(f"Provided path is not a file: {path}")
    if path.suffix not in {".nii", ".gz"} and not path.name.endswith(".nii.gz"):
        raise ValueError(f"File must be .nii or .nii.gz: {path.name}")