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

class Forge:
    def __init__(self, scout):
        self.scout = scout
        self.audit = scout.audit

    def build(self, model_type="tree"):
        transformers = []
        stats = self.audit.feature_catalog
        
        # --- LÓGICA DE CONJUNTOS (Evita sobreposição de colunas) ---
        # 1. Tudo o que não foi dropado
        all_cols = set(self.scout.numeric_cols + self.scout.categorical_cols) - set(self.scout.dropped_cols)
        
        # 2. Define numéricos validos
        valid_num = set(self.scout.numeric_cols) & all_cols
        
        # 3. Define categóricos validos (EXCLUINDO o que já é numérico para evitar o erro)
        valid_cat = (set(self.scout.categorical_cols) & all_cols) - valid_num
        
        # Converte para listas ordenadas para consistência
        valid_num = sorted(list(valid_num))
        valid_cat = sorted(list(valid_cat))
        # -----------------------------------------------------------

        rob, pwr, std = [], [], []
        if not stats.empty:
            for col in valid_num:
                row = stats[stats['column'] == col]
                rec = row.iloc[0]['recommendation'] if not row.empty else "StandardScaler"
                
                if rec == 'RobustScaler': rob.append(col)
                elif rec == 'PowerTransformer' and model_type != 'tree': pwr.append(col)
                else: std.append(col)
        else: std = valid_num

        if rob: transformers.append(("rob", Pipeline([("imp", KNNImputer()), ("scl", RobustScaler())]), rob))
        if pwr: transformers.append(("pwr", Pipeline([("imp", KNNImputer()), ("scl", PowerTransformer())]), pwr))
        if std: transformers.append(("std", Pipeline([("imp", KNNImputer()), ("scl", StandardScaler())]), std))

        low = [c for c in valid_cat if self.scout.cardinality_map.get(c,0) <= 10]
        high = [c for c in valid_cat if self.scout.cardinality_map.get(c,0) > 10]

        if low: transformers.append(("ohe", Pipeline([("imp", SimpleImputer(strategy="most_frequent")), ("enc", OneHotEncoder(handle_unknown="ignore", sparse_output=False))]), low))
        if high: transformers.append(("te", Pipeline([("imp", SimpleImputer(strategy="most_frequent")), ("enc", TargetEncoder())]), high))

        return ColumnTransformer(transformers, remainder="drop", verbose_feature_names_out=False).set_output(transform="pandas")
