import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.model_selection import KFold
from rapidfuzz import process, fuzz
from typing import List, Dict, Union, Optional
import logging

logger = logging.getLogger("scoutml")


class DateFeaturizer(ScoutEstimator):
    """Detecta colunas de data e extrai features ciclicas/temporais."""

    def __init__(self, date_format: str = None):
        super().__init__()
        self.date_format = date_format
        self.date_cols_ = []
        self.generated_features_ = []

    def fit(self, X: pd.DataFrame, y=None):
        self.history = []
        self.date_cols_ = []

        # Heurística simples: converte pra datetime e vê se não dá erro massivo
        # Ou verifica se o dtype já é datetime
        for col in X.columns:
            if pd.api.types.is_datetime64_any_dtype(X[col]):
                self.date_cols_.append(col)
            elif X[col].dtype == "object":
                # Tenta inferir se é data (amostra para performance)
                try:
                    pd.to_datetime(X[col].dropna().iloc[:100], format=self.date_format)
                    self.date_cols_.append(col)
                except (ValueError, TypeError):
                    continue

        if self.date_cols_:
            self.history.append({"action": "detect_dates", "cols": self.date_cols_})

        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        X = X.copy()
        for col in self.date_cols_:
            # Garante conversão
            if not pd.api.types.is_datetime64_any_dtype(X[col]):
                X[col] = pd.to_datetime(
                    X[col], format=self.date_format, errors="coerce"
                )

            # Extração de Features
            pfx = col
            X[f"{pfx}_year"] = X[col].dt.year
            X[f"{pfx}_month"] = X[col].dt.month
            X[f"{pfx}_day"] = X[col].dt.day
            X[f"{pfx}_dayofweek"] = X[col].dt.dayofweek
            X[f"{pfx}_is_weekend"] = (X[col].dt.dayofweek >= 5).astype(int)

            # Remove a original para não quebrar o Encoder depois
            X.drop(columns=[col], inplace=True)

        return X


class ScoutEstimator(BaseEstimator, TransformerMixin):
    """Classe base para manter o histórico de decisões."""

    def __init__(self):
        self.history = []


class AutoCleaner(ScoutEstimator):
    """Remove colunas constantes, IDs e detecta vazamento de dados."""

    def __init__(self, leakage_threshold: float = 0.95):
        super().__init__()
        self.leakage_threshold = leakage_threshold
        self.cols_to_drop_ = []
        self.leakage_warnings_ = []

    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
        self.cols_to_drop_ = []
        self.history = []
        n_rows = len(X)

        for col in X.columns:
            # 1. ID Detection
            if X[col].nunique() >= (n_rows * 0.99):
                self.cols_to_drop_.append(col)
                self.history.append(
                    {"action": "drop", "col": col, "reason": "High Cardinality (ID)"}
                )
                continue

            # 2. Constant Columns
            if X[col].nunique() <= 1:
                self.cols_to_drop_.append(col)
                self.history.append(
                    {"action": "drop", "col": col, "reason": "Zero Variance"}
                )
                continue

            # 3. Leakage Hunter
            if y is not None and pd.api.types.is_numeric_dtype(X[col]):
                try:
                    corr = X[col].corr(y, method="spearman")
                    if abs(corr) > self.leakage_threshold:
                        self.leakage_warnings_.append((col, corr))
                        self.history.append(
                            {
                                "action": "warning",
                                "col": col,
                                "reason": f"Leakage ({corr:.2f})",
                            }
                        )
                except:
                    pass
        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        if not self.cols_to_drop_:
            return X
        return X.drop(columns=self.cols_to_drop_, errors="ignore")


class SmartImputer(ScoutEstimator):
    """Imputação Híbrida: Mediana para numéricos, Moda para categóricos."""

    def __init__(
        self, null_threshold_drop: float = 0.60, structural_threshold: float = 0.05
    ):
        super().__init__()
        self.null_threshold_drop = null_threshold_drop
        self.structural_threshold = structural_threshold
        self.imputation_map_ = {}
        self.missing_indicators_ = []
        self.cols_to_drop_ = []

    def fit(self, X: pd.DataFrame, y=None):
        self.history = []
        null_percent = X.isnull().mean()

        for col in X.columns:
            pct = null_percent[col]
            if pct == 0:
                continue

            if pct > self.null_threshold_drop:
                self.cols_to_drop_.append(col)
                self.history.append(
                    {"action": "drop", "col": col, "reason": f"Nulls {pct:.1%}"}
                )
                continue

            if pct > self.structural_threshold:
                self.missing_indicators_.append(col)
                self.history.append({"action": "add_indicator", "col": col})

            # Estratégia de Preenchimento Robusta
            if pd.api.types.is_numeric_dtype(X[col]):
                fill_val = X[col].median()
                if pd.isna(fill_val):
                    fill_val = 0
                self.imputation_map_[col] = float(fill_val)
                self.history.append(
                    {"action": "impute_num", "col": col, "val": fill_val}
                )
            else:
                # Usa a Moda (valor mais frequente) ou "MISSING" se vazio
                modes = X[col].mode()
                fill_val = modes[0] if not modes.empty else "MISSING"
                self.imputation_map_[col] = str(fill_val)
                self.history.append(
                    {"action": "impute_cat", "col": col, "val": fill_val}
                )

        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        X = X.copy()
        if self.cols_to_drop_:
            X.drop(columns=self.cols_to_drop_, inplace=True, errors="ignore")

        for col in self.missing_indicators_:
            if col in X.columns:
                X[f"is_missing_{col}"] = X[col].isnull().astype(int)

        X.fillna(value=self.imputation_map_, inplace=True)
        return X


class AdaptiveEncoder(ScoutEstimator):
    """
    Pipeline Completo: Fuzzy Match -> Rare Label Grouping -> Encoding (Target/OHE).
    """

    def __init__(
        self,
        card_threshold: int = 10,
        smoothing: float = 10.0,
        fuzzy_threshold: int = 90,
        rare_threshold: float = 0.01,
    ):
        super().__init__()
        self.card_threshold = card_threshold
        self.smoothing = smoothing
        self.fuzzy_threshold = fuzzy_threshold
        self.rare_threshold = rare_threshold

        # Estado aprendido
        self.fuzzy_maps_ = {}  # {col: {bad_val: good_val}}
        self.rare_groups_ = {}  # {col: [list_of_kept_values]}
        self.encoding_map_ = {}  # {col: {val: target_mean}}
        self.global_means_ = {}  # {col: global_mean}
        self.ohe_cols_ = []

    def _learn_fuzzy_map(self, series: pd.Series) -> Dict[str, str]:
        """Agrupa strings muito parecidas (ex: 'SP', 'sp', 'S. Paulo')."""
        unique_vals = series.dropna().unique().astype(str)
        if len(unique_vals) < 3:
            return {}

        mapping = {}
        # Ordena por tamanho para processar os mais limpos primeiro (heurística simples)
        unique_vals = sorted(unique_vals, key=len)

        covered = set()
        for val in unique_vals:
            if val in covered:
                continue

            # Encontra similares
            matches = process.extract(
                val, unique_vals, limit=None, scorer=fuzz.token_sort_ratio
            )
            # Filtra por threshold
            group = [m[0] for m in matches if m[1] >= self.fuzzy_threshold]

            # O "Líder" do grupo é o valor mais curto ou mais frequente (simplificado aqui como o primeiro)
            leader = val
            for member in group:
                if member != leader:
                    mapping[member] = leader
                covered.add(member)
        return mapping

    def _learn_rare_groups(self, series: pd.Series) -> List[str]:
        """Retorna lista de categorias para MANTER. O resto vira 'OTHER'."""
        freqs = series.value_counts(normalize=True)
        kept = freqs[freqs >= self.rare_threshold].index.tolist()
        return kept

    def _get_smoothed_mean(self, df, col, y_name):
        """Target Encoding com Suavização (Smoothing)."""
        global_mean = df[y_name].mean()
        agg = df.groupby(col)[y_name].agg(["count", "mean"])
        counts = agg["count"]
        means = agg["mean"]
        smooth = (counts * means + self.smoothing * global_mean) / (
            counts + self.smoothing
        )
        return smooth, global_mean

    def fit(self, X: pd.DataFrame, y: pd.Series):
        self.history = []
        X = X.copy()
        y_name = "target_internal"
        X[y_name] = y.values

        cat_cols = X.select_dtypes(include=["object", "category"]).columns

        for col in cat_cols:
            # 1. Fuzzy Matching Learning
            f_map = self._learn_fuzzy_map(X[col])
            if f_map:
                self.fuzzy_maps_[col] = f_map
                self.history.append(
                    {"action": "fuzzy_map", "col": col, "count": len(f_map)}
                )
                X[col] = X[col].replace(
                    f_map
                )  # Aplica localmente para continuar o treino

            # 2. Rare Label Learning
            kept_vals = self._learn_rare_groups(X[col])
            self.rare_groups_[col] = kept_vals
            # Aplica localmente: tudo que não está em kept_vals vira "OTHER"
            X.loc[~X[col].isin(kept_vals), col] = "OTHER"
            self.history.append(
                {"action": "rare_group", "col": col, "kept": len(kept_vals)}
            )

            # 3. Decision: OHE vs Target Encoding
            cardinality = X[col].nunique()
            if cardinality <= self.card_threshold:
                self.ohe_cols_.append(col)
                self.history.append({"action": "ohe", "col": col})
            else:
                smooth_map, global_m = self._get_smoothed_mean(X, col, y_name)
                self.encoding_map_[col] = smooth_map.to_dict()
                self.global_means_[col] = global_m
                self.history.append({"action": "target_encode", "col": col})

        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        X = X.copy()

        # 1. Aplica Fuzzy Maps
        for col, f_map in self.fuzzy_maps_.items():
            if col in X.columns:
                X[col] = X[col].replace(f_map)

        # 2. Aplica Rare Groups
        for col, kept_vals in self.rare_groups_.items():
            if col in X.columns:
                # Se não estiver na lista de mantidos, vira "OTHER"
                X[col] = np.where(X[col].isin(kept_vals), X[col], "OTHER")

        # 3. Aplica Target Encoding
        for col, mapping in self.encoding_map_.items():
            if col in X.columns:
                fallback = self.global_means_[col]
                X[col] = X[col].map(mapping).fillna(fallback)

        # 4. Aplica OHE
        if self.ohe_cols_:
            X = pd.get_dummies(X, columns=self.ohe_cols_, drop_first=True)

        return X

    def fit_transform(
        self, X: pd.DataFrame, y: pd.Series = None, **fit_params
    ) -> pd.DataFrame:
        """K-Fold Target Encoding para evitar Data Leakage."""
        self.fit(X, y)  # Aprende os mapas globais (Fuzzy, Rare, Encoding)

        X_out = X.copy()

        # Primeiro, precisamos aplicar Fuzzy e Rare em todo o X_out para garantir consistência
        # antes de fazer o K-Fold do Target Encoding
        for col, f_map in self.fuzzy_maps_.items():
            X_out[col] = X_out[col].replace(f_map)

        for col, kept_vals in self.rare_groups_.items():
            X_out[col] = np.where(X_out[col].isin(kept_vals), X_out[col], "OTHER")

        # Se não houver target encoding, só retorna transformado
        if not self.encoding_map_:
            return self.transform(X)

        # K-Fold logic apenas para as colunas de Target Encoding
        kf = KFold(n_splits=5, shuffle=True, random_state=42)
        y_series = y if isinstance(y, pd.Series) else pd.Series(y)
        y_name = "target_tmp"

        for col in self.encoding_map_.keys():
            # Cria coluna temporária
            X_out[f"{col}_TE"] = np.nan

        for train_idx, val_idx in kf.split(X_out, y_series):
            # Dados de Treino do Fold (já limpos de fuzzy/rare)
            X_fold_train = X_out.iloc[train_idx].copy()
            y_fold_train = y_series.iloc[train_idx]

            # Dados de Validação
            X_fold_val = X_out.iloc[val_idx].copy()

            df_train = X_fold_train.copy()
            df_train[y_name] = y_fold_train.values

            for col in self.encoding_map_.keys():
                smooth_map, global_m = self._get_smoothed_mean(df_train, col, y_name)
                # Aplica map aprendido no fold de treino sobre o fold de validação
                X_out.loc[val_idx, col] = (
                    X_fold_val[col].map(smooth_map).fillna(global_m)
                )

        # Preencher buracos eventuais com média global aprendida no fit() total
        for col in self.encoding_map_.keys():
            X_out[col] = X_out[col].fillna(self.global_means_[col])

        # Finalmente, OHE (pós tratamento)
        if self.ohe_cols_:
            X_out = pd.get_dummies(X_out, columns=self.ohe_cols_, drop_first=True)

        return X_out
