src/utils.py

"""
Módulo de funções auxiliares para carregar dados, preparar a base
e treinar modelos de classificação para previsão de reclamações.
"""

from typing import Tuple, List, Dict

import logging

import numpy as np
import pandas as pd
from imblearn.over_sampling import SMOTE
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from lightgbm import LGBMClassifier
from sklearn.metrics import (
    roc_auc_score,
    precision_score,
    recall_score,
    f1_score,
    confusion_matrix,
    roc_curve,
)

logger = logging.getLogger(__name__)


def carregar_dados(caminho: str) -> pd.DataFrame:
    """
    Carrega a base de dados de marketing.

    Parâmetros:
        caminho: caminho do arquivo CSV.

    Retorno:
        DataFrame com os dados carregados.

    Observação:
        A função lança ValueError se o arquivo estiver vazio.
    """
    df = pd.read_csv(caminho, sep=None, engine="python")
    if df.empty:
        raise ValueError("Erro: arquivo de dados vazio.")
    logger.info("Base carregada com %d linhas e %d colunas.", df.shape[0], df.shape[1])
    return df


def preparar_base(
    df: pd.DataFrame, aplicar_smote: bool = True, teste: float = 0.3, random_state: int = 42
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[str], StandardScaler]:
    """
    Limpa a base, cria variáveis, aplica dummies, faz split e (opcionalmente) SMOTE.

    Parâmetros:
        df: DataFrame original.
        aplicar_smote: se True, aplica SMOTE na base de treino.
        teste: proporção da base para teste.
        random_state: semente aleatória.

    Retorno:
        X_train, X_test, y_train, y_test, nomes_features, scaler
    """
    df = df.copy()

    # Remover colunas que não ajudam no modelo
    colunas_remover = ["ID", "Z_CostContact", "Z_Revenue"]
    for col in colunas_remover:
        if col in df.columns:
            df = df.drop(columns=[col])

    # Tratar valores nulos em Income
    if "Income" in df.columns:
        df["Income"] = df["Income"].fillna(df["Income"].median())

    # Garantir que a variável-alvo existe
    if "Complain" not in df.columns:
        raise ValueError("Erro: coluna 'Complain' não encontrada na base.")

    # Criar variável numérica a partir da data
    if "Dt_Customer" in df.columns:
        df["Dt_Customer"] = pd.to_datetime(df["Dt_Customer"], dayfirst=True, errors="coerce")
        max_data = df["Dt_Customer"].max()
        df["Customer_Since_Days"] = (max_data - df["Dt_Customer"]).dt.days
        df = df.drop(columns=["Dt_Customer"])

    # Separar X e y
    y = df["Complain"].astype(int)
    X = df.drop(columns=["Complain"])

    # One-hot encoding para variáveis categóricas
    X = pd.get_dummies(X, drop_first=True)

    # Treino e teste
    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=teste,
        random_state=random_state,
        stratify=y,
    )

    # Padronização
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # Aplicar SMOTE apenas no treino
    if aplicar_smote:
        smote = SMOTE(random_state=random_state)
        X_train_scaled, y_train = smote.fit_resample(X_train_scaled, y_train)
        logger.info(
            "SMOTE aplicado. Base de treino balanceada: %s",
            y_train.value_counts().to_dict(),
        )

    feature_names = list(X.columns)
    return X_train_scaled, X_test_scaled, y_train, y_test, feature_names, scaler


def treinar_modelo(
    nome_modelo: str,
    X_train: np.ndarray,
    y_train: np.ndarray,
    random_state: int = 42,
) -> object:
    """
    Treina um modelo de classificação de acordo com o nome informado.

    Parâmetros:
        nome_modelo: nome do algoritmo ("Regressão Logística", "Random Forest", "LightGBM").
        X_train: matriz de treino.
        y_train: vetor alvo de treino.
        random_state: semente aleatória.

    Retorno:
        Modelo treinado.

    Observação:
        Lança ValueError se o nome do modelo não for reconhecido.
    """
    if nome_modelo == "Regressão Logística":
        modelo = LogisticRegression(max_iter=1000)
    elif nome_modelo == "Random Forest":
        modelo = RandomForestClassifier(n_estimators=300, random_state=random_state)
    elif nome_modelo == "LightGBM":
        modelo = LGBMClassifier(random_state=random_state)
    else:
        raise ValueError("Erro: modelo escolhido inválido.")

    modelo.fit(X_train, y_train)
    logger.info("Modelo '%s' treinado com sucesso.", nome_modelo)
    return modelo


def avaliar_modelo(
    modelo: object,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> Dict[str, float]:
    """
    Calcula as métricas AUC, precisão, recall e F1-score.

    Parâmetros:
        modelo: modelo já treinado.
        X_test: matriz de teste.
        y_test: alvo de teste.

    Retorno:
        Dicionário com métricas.
    """
    y_pred = modelo.predict(X_test)
    if hasattr(modelo, "predict_proba"):
        y_proba = modelo.predict_proba(X_test)[:, 1]
    else:
        # fallback para modelos sem predict_proba (não é o caso aqui)
        y_proba = y_pred

    return {
        "AUC": float(roc_auc_score(y_test, y_proba)),
        "Precisão": float(precision_score(y_test, y_pred)),
        "Recall": float(recall_score(y_test, y_pred)),
        "F1-score": float(f1_score(y_test, y_pred)),
    }


def calcular_curva_roc(
    modelo: object, X_test: np.ndarray, y_test: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, float]:
    """
    Calcula pontos da curva ROC e o valor de AUC.

    Parâmetros:
        modelo: modelo treinado.
        X_test: dados de teste.
        y_test: alvo de teste.

    Retorno:
        fpr, tpr, auc
    """
    if hasattr(modelo, "predict_proba"):
        y_proba = modelo.predict_proba(X_test)[:, 1]
    else:
        y_proba = modelo.predict(X_test)

    fpr, tpr, _ = roc_curve(y_test, y_proba)
    auc = roc_auc_score(y_test, y_proba)
    return fpr, tpr, float(auc)


def obter_importancias(
    modelo: object, feature_names: List[str]
) -> pd.DataFrame:
    """
    Retorna um DataFrame com importâncias de variáveis, se disponível.

    Parâmetros:
        modelo: modelo treinado.
        feature_names: lista de nomes das variáveis.

    Retorno:
        DataFrame com colunas 'Variavel' e 'Importancia'.
    """
    if hasattr(modelo, "feature_importances_"):
        importancias = modelo.feature_importances_
        df_imp = pd.DataFrame(
            {"Variavel": feature_names, "Importancia": importancias}
        ).sort_values(by="Importancia", ascending=False)
        return df_imp

    if hasattr(modelo, "coef_"):
        coef = modelo.coef_[0]
        df_imp = pd.DataFrame(
            {"Variavel": feature_names, "Importancia": np.abs(coef)}
        ).sort_values(by="Importancia", ascending=False)
        return df_imp

    # Caso não exista importância
    return pd.DataFrame(
        {"Variavel": feature_names, "Importancia": [0.0] * len(feature_names)}
    )