backend/app/core/anexos/model_loader.py

"""
Carregador de modelos de avaliação.
Carrega arquivos .dai via joblib.
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Optional

import joblib
import pandas as pd

from .model_data import ModelData


def _get_first(container: dict[str, Any], *keys: str, default: Any = None) -> Any:
    """Retorna o primeiro valor não-nulo encontrado nas chaves informadas."""
    for key in keys:
        if key in container and container[key] is not None:
            return container[key]
    return default


def _as_dict(value: Any) -> dict[str, Any]:
    """Garante retorno como dicionário."""
    return value if isinstance(value, dict) else {}


def _as_dataframe(value: Any) -> pd.DataFrame:
    """Converte entrada em DataFrame quando possível."""
    if isinstance(value, pd.DataFrame):
        return value
    if isinstance(value, pd.Series):
        return value.to_frame()
    return pd.DataFrame()


def _as_series(value: Any) -> pd.Series:
    """Converte entrada em Series quando possível."""
    if isinstance(value, pd.Series):
        return value
    if isinstance(value, pd.DataFrame) and value.shape[1] == 1:
        return value.iloc[:, 0]
    return pd.Series(dtype=float)


def _as_string_list(value: Any) -> list[str]:
    """Converte listas/tuplas/índices/séries em lista de strings."""
    if value is None:
        return []
    if isinstance(value, pd.Series):
        value = value.tolist()
    elif isinstance(value, pd.Index):
        value = value.tolist()
    elif not isinstance(value, (list, tuple, set)):
        return []
    return [str(item) for item in value]


def _extract_diagnostic_summary(diag: dict[str, Any]) -> dict[str, Any]:
    """
    Normaliza estatísticas do bloco de diagnósticos para o formato interno.

    O formato dos .dai pode variar entre versões; esta função tolera aliases
    comuns sem quebrar o carregamento do modelo.
    """
    gerais = _as_dict(_get_first(diag, "gerais", "estatisticas_gerais", default={}))
    teste_f = _as_dict(_get_first(diag, "teste_f", "f_test", default={}))
    teste_ks = _as_dict(_get_first(diag, "teste_ks", "ks_test", default={}))
    teste_normalidade = _as_dict(
        _get_first(diag, "teste_normalidade", "normalidade", default={})
    )
    teste_dw = _as_dict(_get_first(diag, "teste_dw", "durbin_watson", default={}))
    teste_bp = _as_dict(_get_first(diag, "teste_bp", "breusch_pagan", default={}))

    def stat_from(block: dict[str, Any]) -> Any:
        return _get_first(block, "estatistica", "stat", "value")

    def pvalue_from(block: dict[str, Any]) -> Any:
        return _get_first(block, "p_valor", "pvalue", "p_val", "p")

    def interp_from(block: dict[str, Any]) -> Any:
        return _get_first(block, "interpretacao", "interpretation", "texto")

    return {
        **gerais,
        "Fc": stat_from(teste_f),
        "p_valor_F": pvalue_from(teste_f),
        "Interpretacao_F": interp_from(teste_f),
        "ks_stat": stat_from(teste_ks),
        "ks_p": pvalue_from(teste_ks),
        "Interpretacao_KS": interp_from(teste_ks),
        "perc_resid": _get_first(
            teste_normalidade, "percentuais", "percentis", "percentages"
        ),
        "dw": stat_from(teste_dw),
        "Interpretacao_DW": interp_from(teste_dw),
        "bp_lm": stat_from(teste_bp),
        "bp_p": pvalue_from(teste_bp),
        "Interpretacao_BP": interp_from(teste_bp),
        "equacao": _get_first(diag, "equacao", "equation", "formula"),
    }


def _build_xy_preview(
    dados: dict[str, Any], transformacoes: dict[str, Any]
) -> pd.DataFrame:
    """
    Monta DataFrame de preview para o Anexo I.

    Em versões novas, `dados['df']` pode conter o dataset completo (centenas de
    colunas), o que degrada a geração/abertura no Word. Quando disponível, usa
    `transformacoes['X'] + y` como preview estável da modelagem.
    """
    df_preview = _as_dataframe(
        _get_first(
            dados,
            "xy_preview",
            "df_preview",
            "df_modelagem",
            "preview",
            "df",
            default=pd.DataFrame(),
        )
    )

    x_df = _as_dataframe(_get_first(transformacoes, "X", "x", default=pd.DataFrame()))
    y_series = _as_series(_get_first(transformacoes, "y", "Y", default=pd.Series(dtype=float)))

    if not x_df.empty:
        modelagem_df = x_df.copy()

        if not y_series.empty:
            y_name = str(y_series.name) if y_series.name else "y"
            y_alinhado = y_series.reset_index(drop=True)
            modelagem_df = modelagem_df.reset_index(drop=True)
            tamanho = min(len(modelagem_df), len(y_alinhado))
            if tamanho > 0:
                modelagem_df = modelagem_df.iloc[:tamanho].copy()
                modelagem_df[y_name] = y_alinhado.iloc[:tamanho].values

        # Se o preview original veio "explodido" (df completo), prioriza modelagem.
        if df_preview.empty or df_preview.shape[1] > max(40, modelagem_df.shape[1] + 10):
            return modelagem_df

    return df_preview


def model_data_from_package(
    data: Any,
    dai_path: Path,
    nome: str | None = None,
) -> Optional[ModelData]:
    """Converte um pacote .dai já carregado para os dados usados nos anexos."""
    try:
        data_dict = _as_dict(data)
        dados = _as_dict(_get_first(data_dict, "dados", "data", default={}))
        transformacoes = _as_dict(
            _get_first(data_dict, "transformacoes", "transformations", default={})
        )
        modelo_dict = _as_dict(_get_first(data_dict, "modelo", "model", default={}))
        diag = _as_dict(_get_first(modelo_dict, "diagnosticos", "diagnostics", default={}))

        modelos_resumos = _extract_diagnostic_summary(diag)
        xy_preview = _build_xy_preview(dados, transformacoes)
        top_x_esc = _as_dataframe(
            _get_first(transformacoes, "X", "x", default=pd.DataFrame())
        )
        top_y_esc = _as_series(
            _get_first(transformacoes, "y", "Y", default=pd.Series(dtype=float))
        )
        estatisticas = _as_dataframe(
            _get_first(dados, "estatisticas", "statistics", default=pd.DataFrame())
        )
        tabelas_coef = _as_dataframe(
            _get_first(
                modelo_dict,
                "coeficientes",
                "coeficients",
                "coefficients",
                default=pd.DataFrame(),
            )
        )
        for coluna_variavel in ("Variável", "Variavel", "variable"):
            if coluna_variavel in tabelas_coef.columns:
                tabelas_coef = tabelas_coef.set_index(coluna_variavel)
                break

        tabelas_obs_calc = _as_dataframe(
            _get_first(
                modelo_dict,
                "obs_calc",
                "observado_calculado",
                "observed_vs_calculated",
                default=pd.DataFrame(),
            )
        )
        modelos_sm = _get_first(
            modelo_dict, "sm", "statsmodels", "statsmodels_result", default=None
        )

        return ModelData(
            nome=str(nome or dai_path.stem),
            path=dai_path,
            xy_preview=xy_preview,
            top_x_esc=top_x_esc,
            top_y_esc=top_y_esc,
            estatisticas=estatisticas,
            tabelas_coef=tabelas_coef,
            tabelas_obs_calc=tabelas_obs_calc,
            modelos_resumos=modelos_resumos,
            modelos_sm=modelos_sm,
            formatted_top_transformation_info=transformacoes.get("info", []),
            variaveis_dicotomicas=_as_string_list(
                _get_first(
                    transformacoes,
                    "dicotomicas",
                    "dicotômicas",
                    "dummies",
                    "dummy",
                    "binarias",
                    "binárias",
                    default=[],
                )
            ),
            graf_model="",
        )
    except (KeyError, TypeError, ValueError):
        return None


def load_model(dai_path: Path) -> Optional[ModelData]:
    """
    Carrega um modelo de avaliação a partir de um arquivo .dai.

    Args:
        dai_path: Caminho para o arquivo .dai

    Returns:
        ModelData com os dados carregados, ou None se falhar
    """
    if not dai_path.exists():
        return None

    try:
        data = joblib.load(dai_path)
    except Exception as e:
        print(f"Erro ao carregar {dai_path}: {e}")
        return None

    return model_data_from_package(data, dai_path)


def load_model_by_name(name: str, models_dir: Path) -> Optional[ModelData]:
    """
    Carrega um modelo pelo nome.

    Args:
        name: Nome do modelo (nome do arquivo .dai sem extensão)
        models_dir: Diretório dos modelos

    Returns:
        ModelData com os dados carregados, ou None se não encontrar
    """
    dai_path = models_dir / f"{name}.dai"
    return load_model(dai_path)