import numpy as np

from .transformation import Transformation


class Normalization(Transformation):
    def __init__(self, norm_tech: str) -> None:
        """
        Initialize the normalization technique

        Parameters
        ----------
        norm_tech : str
            Normalization technique to be applied. The following techniques are supported:
            - "arctan"
            - "boxcox"
            - "decimal"
            - "logistic"
            - "mad"
            - "max"
            - "maxabs"
            - "maxmin"
            - "mean"
            - "median"
            - "min"
            - "minabs"
            - "minmax"
            - "power"
            - "quantile"
            - "robust"
            - "softmax"
            - "std"
            - "tanh"
            - "var"
            - "yeojohnson"
            - "zscore"
        """
        self.norm_tech: str = norm_tech

    def __call__(self, x: np.ndarray) -> np.ndarray:
        """
        Apply normalization technique to the input data

        Parameters
        ----------
        x : np.ndarray
            Input data to be normalized

        Returns
        -------
        np.ndarray
            Normalized data
        """
        eps = 1e-12

        def safe_div(numerator: np.ndarray, denominator: float | np.ndarray) -> np.ndarray:
            denominator = np.asarray(denominator, dtype=float)
            denominator = np.where(np.abs(denominator) < eps, eps, denominator)
            return numerator / denominator

        match self.norm_tech:
            case "arctan":
                return np.arctan(x)
            case "boxcox":
                return np.sign(x) * np.power(np.abs(x), 1 / 2)
            case "decimal":
                max_abs = np.max(np.abs(x))
                scale = 10 ** np.ceil(np.log10(max(max_abs, eps)))
                return safe_div(x, scale)
            case "logistic":
                return 1 / (1 + np.exp(-x))
            case "mad":
                return safe_div(x, np.median(np.abs(x - np.median(x))))
            case "max":
                return safe_div(x, np.max(x))
            case "maxabs":
                return safe_div(x, np.max(np.abs(x)))
            case "maxmin":
                return safe_div(x, (np.max(x) - np.min(x)))
            case "mean":
                return safe_div(x, np.mean(x))
            case "median":
                return safe_div(x, np.median(x))
            case "min":
                return safe_div(x, np.min(x))
            case "minabs":
                return safe_div(x, np.min(np.abs(x)))
            case "minmax":
                return safe_div(x - np.min(x), (np.max(x) - np.min(x)))
            case "power":
                return np.sign(x) * np.power(np.abs(x), 1 / 2)
            case "quantile":
                return safe_div(x, np.quantile(x, 0.95))
            case "robust":
                return safe_div(x, (np.percentile(x, 75) - np.percentile(x, 25)))
            case "softmax":
                shifted = x - np.max(x, axis=-1, keepdims=True)
                exp_x = np.exp(shifted)
                return safe_div(exp_x, np.sum(exp_x, axis=-1, keepdims=True))
            case "std":
                return safe_div(x, np.std(x))
            case "tanh":
                return np.tanh(x)
            case "var":
                return safe_div(x, np.var(x))
            case "yeojohnson":
                return np.sign(x) * np.power(np.abs(x), 1 / 2)
            case "zscore":
                return safe_div(x - np.mean(x), np.std(x))
            case _:
                raise ValueError("Unknown normalization technique")
