"""Remove Constant Features Step."""

from __future__ import annotations

from typing_extensions import override

import numpy as np
import torch

from tabpfn.preprocessors.preprocessing_helpers import (
    FeaturePreprocessingTransformerStep,
)


class RemoveConstantFeaturesStep(FeaturePreprocessingTransformerStep):
    """Remove features that are constant in the training data."""

    def __init__(self) -> None:
        super().__init__()
        self.sel_: list[bool] | None = None

    @override
    def _fit(  # type: ignore
        self, X: np.ndarray | torch.Tensor, categorical_features: list[int]
    ) -> list[int]:
        if isinstance(X, torch.Tensor):
            sel_ = torch.max(X[0:1, :] != X, dim=0)[0].cpu()
        else:
            sel_ = ((X[0:1, :] == X).mean(axis=0) < 1.0).tolist()

        if not any(sel_):
            raise ValueError(
                "All features are constant and would have been removed!"
                " Unable to predict using TabPFN.",
            )
        self.sel_ = sel_

        return [
            new_idx
            for new_idx, idx in enumerate(np.where(sel_)[0])
            if idx in categorical_features
        ]

    @override
    def _transform(
        self, X: np.ndarray | torch.Tensor, *, is_test: bool = False
    ) -> np.ndarray:
        assert self.sel_ is not None, "You must call fit first"
        return X[:, self.sel_]
