from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer, KNNImputer
from sklearn.preprocessing import (
    RobustScaler,
    OneHotEncoder,
    TargetEncoder,
    PowerTransformer,
    StandardScaler,
)
from .scout import Scout


class Forge:
    """
    Forge v2.0: Constrói pipelines de transformação dinâmicos.
    Lê o 'Feature Catalog' do AuditTrail para decidir a melhor estratégia de Scaling.
    """

    def __init__(self, scout: Scout):
        self.scout = scout
        self.audit = scout.audit

    def build(self, model_type="tree"):
        """
        Constrói o pré-processador.

        Args:
            model_type: 'tree' (XGBoost/LGBM) ou 'linear'.
                        Modelos de árvore são menos sensíveis a escala, mas 'RobustScaler'
                        ajuda na convergência se houver outliers extremos.
        """
        transformers = []

        # --- 1. LÓGICA NUMÉRICA DINÂMICA (A Grande Mudança v2.0) ---
        # Tenta ler as estatísticas calculadas pelo novo Scout
        stats_df = self.audit.feature_catalog

        # Listas para agrupar colunas por estratégia
        robust_cols = []
        power_cols = []
        standard_cols = []

        # Se temos colunas numéricas para processar
        if self.scout.numeric_cols:

            # Se o Scout correu o profile estatístico (v2.0), usamos essa inteligência
            if not stats_df.empty:
                for col in self.scout.numeric_cols:
                    # Busca a linha correspondente no DataFrame de estatísticas
                    row = stats_df[stats_df["column"] == col]

                    if not row.empty:
                        rec = row.iloc[0]["recommendation"]

                        # Decisão baseada na recomendação e no tipo de modelo
                        if rec == "RobustScaler":
                            robust_cols.append(col)
                        elif rec == "PowerTransformer" and model_type != "tree":
                            # Só vale a pena gaussianizar (caro computacionalmente) se não for árvore
                            power_cols.append(col)
                        else:
                            standard_cols.append(col)
                    else:
                        # Se não houver stats para esta coluna específica, fallback para Standard
                        standard_cols.append(col)

            else:
                # Fallback V1: Se não houver stats (Scout antigo), tudo é Standard
                standard_cols = self.scout.numeric_cols

            # --- CRIAÇÃO DOS PIPELINES NUMÉRICOS ---

            # A) Pipeline para Outliers (RobustScaler)
            if robust_cols:
                self.audit.log_pipeline_decision(
                    feature=robust_cols,
                    action="RobustScaler",
                    reason="Outliers excessivos (>5%) detetados via IQR.",
                )
                transformers.append(
                    (
                        "num_robust",
                        Pipeline(
                            [
                                ("imputer", KNNImputer(n_neighbors=5)),
                                ("scaler", RobustScaler()),
                            ]
                        ),
                        robust_cols,
                    )
                )

            # B) Pipeline para Distribuições Enviesadas (Yeo-Johnson)
            if power_cols:
                self.audit.log_pipeline_decision(
                    feature=power_cols,
                    action="PowerTransformer",
                    reason="Alta Assimetria (Skewness > 1.0). Gaussianização necessária.",
                )
                transformers.append(
                    (
                        "num_power",
                        Pipeline(
                            [
                                ("imputer", KNNImputer(n_neighbors=5)),
                                ("scaler", PowerTransformer()),
                            ]
                        ),
                        power_cols,
                    )
                )

            # C) Pipeline Padrão (StandardScaler)
            if standard_cols:
                # Não logamos decisão aqui para não poluir o relatório (é o default)
                transformers.append(
                    (
                        "num_std",
                        Pipeline(
                            [
                                ("imputer", KNNImputer(n_neighbors=5)),
                                ("scaler", StandardScaler()),
                            ]
                        ),
                        standard_cols,
                    )
                )

        # --- 2. LÓGICA CATEGÓRICA (Baseada em Cardinalidade) ---

        # Divide as colunas baseando-se no mapa de cardinalidade criado pelo Scout
        low_card = [
            c
            for c in self.scout.categorical_cols
            if self.scout.cardinality_map.get(c, 0) <= 10
        ]
        high_card = [
            c
            for c in self.scout.categorical_cols
            if self.scout.cardinality_map.get(c, 0) > 10
        ]

        # A) Baixa Cardinalidade -> OneHot (Dummies)
        if low_card:
            self.audit.log_pipeline_decision(
                feature=low_card,
                action="OneHotEncoder",
                reason="Baixa Cardinalidade (<= 10 classes únicas).",
            )
            transformers.append(
                (
                    "cat_low",
                    Pipeline(
                        [
                            ("imputer", SimpleImputer(strategy="most_frequent")),
                            (
                                "encoder",
                                OneHotEncoder(
                                    handle_unknown="ignore", sparse_output=False
                                ),
                            ),
                        ]
                    ),
                    low_card,
                )
            )

        # B) Alta Cardinalidade -> TargetEncoder (Melhor que OneHot para muitas categorias)
        if high_card:
            self.audit.log_pipeline_decision(
                feature=high_card,
                action="TargetEncoder",
                reason="Alta Cardinalidade (> 10 classes). Evita explosão dimensional.",
            )
            transformers.append(
                (
                    "cat_high",
                    Pipeline(
                        [
                            ("imputer", SimpleImputer(strategy="most_frequent")),
                            ("encoder", TargetEncoder()),
                        ]
                    ),
                    high_card,
                )
            )

        # --- 3. MONTAGEM FINAL ---

        # ColumnTransformer aplica os pipelines em paralelo e junta tudo no final
        # verbose_feature_names_out=False mantem os nomes das colunas limpos (sem prefixos feios)
        preprocessor = ColumnTransformer(
            transformers=transformers, remainder="drop", verbose_feature_names_out=False
        ).set_output(
            transform="pandas"
        )  # Garante que o output seja DataFrame pandas, vital para rastreabilidade

        return preprocessor
