from typing import Dict, Any, Union, Tuple
import logging
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler, MaxAbsScaler

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

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)
