"""Feature normalisation transform."""

from __future__ import annotations

import numpy as np

from .base_transform import FeatureTransform


class Normalise(FeatureTransform):
    """Standardise features to zero mean and unit variance."""

    def __init__(self, name: str | None = None) -> None:
        super().__init__(name=name)
        self.mean_: np.ndarray | None = None
        self.std_: np.ndarray | None = None

    def fit(self, X: np.ndarray) -> Normalise:
        """Estimate normalisation statistics.

        Parameters
        ----------
        X : np.ndarray
            Input feature matrix.

        Returns
        -------
        Normalise
            The fitted transform.
        """

        X_array = self._as_2d(X)
        self.mean_ = np.mean(X_array, axis=0)
        std = np.std(X_array, axis=0)
        self.std_ = np.where(std == 0.0, 1.0, std)
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        """Standardise an input feature matrix.

        Parameters
        ----------
        X : np.ndarray
            Input feature matrix.

        Returns
        -------
        np.ndarray
            Standardised feature matrix.
        """

        X_array = self._as_2d(X)
        if self.mean_ is None or self.std_ is None:
            self.fit(X_array)
        input_names = self.input_feature_names or [
            f"feature_{index}" for index in range(X_array.shape[1])
        ]
        self._feature_names = list(input_names)
        return (X_array - self.mean_) / self.std_
