import logging
from typing import Any, Dict, Tuple, Union

import numpy as np
import pandas as pd
from sklearn.preprocessing import (
    MaxAbsScaler,
    MinMaxScaler,
    RobustScaler,
    StandardScaler,
)

from ..utils import (
    detect_numeric_columns,
    pack_pipeline_output,
    resolve_columns,
    unpack_pipeline_input,
)
from .base import BaseApplier, BaseCalculator

logger = logging.getLogger(__name__)

# --- Standard Scaler ---


class StandardScalerCalculator(BaseCalculator):
    def fit(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        config: Dict[str, Any],
    ) -> Dict[str, Any]:
        X, _, _ = unpack_pipeline_input(df)

        # Config: {'with_mean': True, 'with_std': True, 'columns': [...]}
        with_mean = config.get("with_mean", True)
        with_std = config.get("with_std", True)

        cols = resolve_columns(X, config, detect_numeric_columns)

        if not cols:
            return {}

        scaler = StandardScaler(with_mean=with_mean, with_std=with_std)
        scaler.fit(X[cols])

        return {
            "type": "standard_scaler",
            "mean": scaler.mean_.tolist() if scaler.mean_ is not None else None,
            "scale": scaler.scale_.tolist() if scaler.scale_ is not None else None,
            "var": scaler.var_.tolist() if scaler.var_ is not None else None,
            "with_mean": with_mean,
            "with_std": with_std,
            "columns": cols,
        }


class StandardScalerApplier(BaseApplier):
    def apply(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        params: Dict[str, Any],
    ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]]:
        X, y, is_tuple = unpack_pipeline_input(df)

        cols = params.get("columns", [])
        mean = params.get("mean")
        scale = params.get("scale")

        valid_cols = [c for c in cols if c in X.columns]
        if not valid_cols or mean is None or scale is None:
            return pack_pipeline_output(X, y, is_tuple)

        X_out = X.copy()
        mean_arr = np.array(mean)
        scale_arr = np.array(scale)
        col_indices = [cols.index(c) for c in valid_cols]

        vals = X_out[valid_cols].values
        if params.get("with_mean", True):
            vals = vals - mean_arr[col_indices]
        if params.get("with_std", True):
            safe_scale = scale_arr[col_indices]
            safe_scale[safe_scale == 0] = 1.0
            vals = vals / safe_scale

        X_out[valid_cols] = vals
        return pack_pipeline_output(X_out, y, is_tuple)


# --- MinMax Scaler ---


class MinMaxScalerCalculator(BaseCalculator):
    def fit(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        config: Dict[str, Any],
    ) -> Dict[str, Any]:
        X, _, _ = unpack_pipeline_input(df)

        # Config: {'feature_range': (0, 1), 'columns': [...]}
        feature_range = config.get("feature_range", (0, 1))

        cols = resolve_columns(X, config, detect_numeric_columns)

        if not cols:
            return {}

        scaler = MinMaxScaler(feature_range=feature_range)
        scaler.fit(X[cols])

        return {
            "type": "minmax_scaler",
            "min": scaler.min_.tolist(),
            "scale": scaler.scale_.tolist(),
            "data_min": scaler.data_min_.tolist(),
            "data_max": scaler.data_max_.tolist(),
            "feature_range": feature_range,
            "columns": cols,
        }


class MinMaxScalerApplier(BaseApplier):
    def apply(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        params: Dict[str, Any],
    ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]]:
        X, y, is_tuple = unpack_pipeline_input(df)

        cols = params.get("columns", [])
        min_val = params.get("min")
        scale = params.get("scale")

        valid_cols = [c for c in cols if c in X.columns]
        if not valid_cols or min_val is None or scale is None:
            return pack_pipeline_output(X, y, is_tuple)

        X_out = X.copy()
        min_arr = np.array(min_val)
        scale_arr = np.array(scale)
        col_indices = [cols.index(c) for c in valid_cols]

        vals = X_out[valid_cols].values
        vals = vals * scale_arr[col_indices] + min_arr[col_indices]
        X_out[valid_cols] = vals
        return pack_pipeline_output(X_out, y, is_tuple)


# --- Robust Scaler ---


class RobustScalerCalculator(BaseCalculator):
    def fit(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        config: Dict[str, Any],
    ) -> Dict[str, Any]:
        X, _, _ = unpack_pipeline_input(df)

        # Config: {'quantile_range': (25.0, 75.0), 'with_centering': True, 'with_scaling': True, 'columns': [...]}
        quantile_range = config.get("quantile_range", (25.0, 75.0))
        with_centering = config.get("with_centering", True)
        with_scaling = config.get("with_scaling", True)

        cols = resolve_columns(X, config, detect_numeric_columns)

        if not cols:
            return {}

        scaler = RobustScaler(
            quantile_range=quantile_range,
            with_centering=with_centering,
            with_scaling=with_scaling,
        )
        scaler.fit(X[cols])

        return {
            "type": "robust_scaler",
            "center": scaler.center_.tolist() if scaler.center_ is not None else None,
            "scale": scaler.scale_.tolist() if scaler.scale_ is not None else None,
            "quantile_range": quantile_range,
            "with_centering": with_centering,
            "with_scaling": with_scaling,
            "columns": cols,
        }


class RobustScalerApplier(BaseApplier):
    def apply(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        params: Dict[str, Any],
    ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]]:
        X, y, is_tuple = unpack_pipeline_input(df)

        cols = params.get("columns", [])
        center = params.get("center")
        scale = params.get("scale")

        valid_cols = [c for c in cols if c in X.columns]
        if not valid_cols:
            return pack_pipeline_output(X, y, is_tuple)

        X_out = X.copy()
        col_indices = [cols.index(c) for c in valid_cols]
        vals = X_out[valid_cols].values

        if params.get("with_centering", True) and center is not None:
            center_arr = np.array(center)
            vals = vals - center_arr[col_indices]

        if params.get("with_scaling", True) and scale is not None:
            scale_arr = np.array(scale)
            # Avoid division by zero
            safe_scale = scale_arr[col_indices]
            safe_scale[safe_scale == 0] = 1.0
            vals = vals / safe_scale

        X_out[valid_cols] = vals
        return pack_pipeline_output(X_out, y, is_tuple)


# --- MaxAbs Scaler ---


class MaxAbsScalerCalculator(BaseCalculator):
    def fit(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        config: Dict[str, Any],
    ) -> Dict[str, Any]:
        X, _, _ = unpack_pipeline_input(df)

        cols = resolve_columns(X, config, detect_numeric_columns)

        if not cols:
            return {}

        scaler = MaxAbsScaler()
        scaler.fit(X[cols])

        return {
            "type": "maxabs_scaler",
            "scale": scaler.scale_.tolist() if scaler.scale_ is not None else None,
            "max_abs": (
                scaler.max_abs_.tolist() if scaler.max_abs_ is not None else None
            ),
            "columns": cols,
        }


class MaxAbsScalerApplier(BaseApplier):
    def apply(
        self,
        df: Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]],
        params: Dict[str, Any],
    ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, pd.Series]]:
        X, y, is_tuple = unpack_pipeline_input(df)

        cols = params.get("columns", [])
        scale = params.get("scale")

        valid_cols = [c for c in cols if c in X.columns]
        if not valid_cols or scale is None:
            return pack_pipeline_output(X, y, is_tuple)

        X_out = X.copy()
        scale_arr = np.array(scale)
        col_indices = [cols.index(c) for c in valid_cols]

        vals = X_out[valid_cols].values
        # MaxAbsScaler just divides by max_abs (which is stored in scale_)
        safe_scale = scale_arr[col_indices]
        safe_scale[safe_scale == 0] = 1.0
        vals = vals / safe_scale

        X_out[valid_cols] = vals
        return pack_pipeline_output(X_out, y, is_tuple)
