model/mpl_pequeno_keras/model.py

"""
Implementação concreta (Keras/TensorFlow) do modelo treinado pelo pipeline.

Arquitetura: MLP simples para dados tabulares. A configuração de
camadas, dropout, otimizador, etc. vem de ``config`` (lido do
``config.json`` do treino) ou de defaults razoáveis.

A classe ``TabularMLP`` herda de ``BaseModel`` (em ``docker/base_model.py``)
para manter a API uniforme entre frameworks. Uma função ``create_model``
fábrica é exposta no fim do arquivo para que ``main.py`` continue
funcionando sem mudanças.
"""

from __future__ import annotations

import os
import sys
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Permite importar base_model.py de ../
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from base_model import BaseModel  # noqa: E402
from utils import emit_progress  # noqa: E402


# ---------------------------------------------------------------------------
# Wrapper que implementa a API de BaseModel para Keras
# ---------------------------------------------------------------------------
class TabularMLP(BaseModel):
    """MLP tabular em Keras/TensorFlow, plugável no pipeline."""

    def build_model(self) -> keras.Model:
        hp = self.config.get("hyperparameters", {})
        hidden_dims = self.config.get("hidden_dims") or hp.get("hidden_dims") or [
            self.config.get("hidden_dim1", 32),
            self.config.get("hidden_dim2", 16),
        ]
        dropout = float(hp.get("dropout", self.config.get("dropout", 0.0)))

        # Constrói o modelo Sequential
        model = keras.Sequential()
        model.add(layers.InputLayer(input_shape=(self.input_size,)))

        # Adiciona camadas ocultas
        for h in hidden_dims:
            model.add(layers.Dense(h, activation="relu"))
            if dropout and dropout > 0:
                model.add(layers.Dropout(dropout))

        # Camada de saída
        model.add(layers.Dense(self.num_classes, activation="softmax"))

        self._hidden_dims = list(hidden_dims)
        self._dropout = dropout
        return model

    # ---------------- treinamento ----------------
    def train_model(
        self,
        train_loader: Any,
        val_loader: Optional[Any],
        epochs: int,
        lr: float,
        **kwargs: Any,
    ) -> Dict[str, list]:
        hp = self.config.get("hyperparameters", {})
        epochs = int(hp.get("epochs", epochs))
        lr = float(hp.get("learning_rate", lr))
        weight_decay = float(hp.get("weight_decay", 0.0))
        optimizer_name = str(hp.get("optimizer", "adam")).lower()

        # Desempacota tuplas (X, y) do dataset
        if isinstance(train_loader, tuple) and len(train_loader) == 2:
            X_train, y_train = train_loader
        else:
            raise ValueError(f"train_loader deve ser tupla (X, y), recebeu {type(train_loader)}")
        
        if isinstance(val_loader, tuple) and len(val_loader) == 2:
            X_val, y_val = val_loader
        else:
            raise ValueError(f"val_loader deve ser tupla (X, y), recebeu {type(val_loader)}")

        # Compila o modelo
        optimizer = self._build_optimizer(optimizer_name, lr, weight_decay)
        self.model.compile(
            optimizer=optimizer,
            loss="sparse_categorical_crossentropy",
            metrics=["accuracy"],
        )

        history = {
            "train_loss": [],
            "train_acc": [],
            "val_loss": [],
            "val_acc": [],
        }

        n_batches = max((len(X_train) + 31) // 32, 1)  # Batches por epoch
        total_train_steps = max(epochs * n_batches, 1)
        batch_size = int(hp.get("batch_size", 32))
        best_val_acc = -float("inf")
        best_epoch = -1

        for epoch in range(1, epochs + 1):
            # Treina 1 época
            hist = self.model.fit(
                X_train,
                y_train,
                batch_size=batch_size,
                epochs=1,
                validation_data=(X_val, y_val),
                verbose=0,
            )

            # Extrai histórico
            train_loss = float(hist.history["loss"][0])
            train_acc = float(hist.history["accuracy"][0])
            val_loss = float(hist.history["val_loss"][0])
            val_acc = float(hist.history["val_accuracy"][0])

            history["train_loss"].append(round(train_loss, 6))
            history["train_acc"].append(round(train_acc, 6))
            history["val_loss"].append(round(val_loss, 6))
            history["val_acc"].append(round(val_acc, 6))

            # Emite progresso por batch (simula progresso contínuo por época)
            for batch_idx in range(1, n_batches + 1):
                global_step = (epoch - 1) * n_batches + batch_idx
                inner_pct = int(global_step * 100 / total_train_steps)
                emit_progress(inner_pct, 1)

            # Rastreia melhor validação
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                best_epoch = epoch

            print(
                f"Epoch [{epoch:>3}/{epochs}] "
                f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} "
                f"val_loss={val_loss:.4f} val_acc={val_acc:.4f}"
            )

        history["best_epoch"] = best_epoch
        history["best_val_acc"] = round(best_val_acc, 6) if best_epoch > 0 else None
        self.history = history
        return history

    # ---------------- avaliação ----------------
    def evaluate(self, data_loader: Any) -> Dict[str, float]:
        if isinstance(data_loader, tuple) and len(data_loader) == 2:
            X_val, y_val = data_loader
        else:
            X_val, y_val = data_loader[2], data_loader[3]
        
        loss, accuracy = self.model.evaluate(X_val, y_val, verbose=0)
        return {
            "loss": float(loss),
            "accuracy": float(accuracy),
            "n": len(X_val),
        }

    # ---------------- inferência ----------------
    def predict(self, inputs: Any) -> np.ndarray:
        if isinstance(inputs, tuple):
            # Se for tuple de (X_train, y_train, X_val, y_val), usa X_val
            X = inputs[2].astype(np.float32)
        elif isinstance(inputs, np.ndarray):
            X = inputs.astype(np.float32)
        else:
            X = np.array(inputs, dtype=np.float32)

        if X.ndim == 1:
            X = X.reshape(1, -1)

        # Predições como classe
        logits = self.model.predict(X, verbose=0)
        return np.argmax(logits, axis=1)

    # ---------------- persistência ----------------
    def save_model(self, filename: str) -> str:
        path = os.path.join(self.model_dir, filename)
        # Remove extensão se for .pt (compatibilidade)
        if path.endswith(".pt"):
            path = path[:-3] + ".keras"

        self.model.save(path)
        print(f"✓ Model saved to {path}")
        return path

    def load_model(self, filename: str) -> None:
        path = filename if os.path.isabs(filename) else os.path.join(
            self.model_dir, filename
        )
        # Converte .pt para .keras
        if path.endswith(".pt"):
            path = path[:-3] + ".keras"

        self.model = keras.models.load_model(path)

    # ---------------- helpers ----------------
    def _build_optimizer(
        self, name: str, lr: float, weight_decay: float
    ) -> keras.optimizers.Optimizer:
        if name == "sgd":
            return keras.optimizers.SGD(learning_rate=lr, weight_decay=weight_decay)
        if name == "rmsprop":
            return keras.optimizers.RMSprop(learning_rate=lr, weight_decay=weight_decay)
        # default: adam
        return keras.optimizers.Adam(learning_rate=lr, weight_decay=weight_decay)


# ---------------------------------------------------------------------------
# Fábrica usada pelo main.py
# ---------------------------------------------------------------------------
def create_model(
    input_size: int,
    num_classes: int,
    project_name: str,
    base_path: str,
    config: Optional[Dict[str, Any]] = None,
) -> TabularMLP:
    """Cria um ``TabularMLP`` pronto pra uso pelo main.py."""
    return TabularMLP(
        input_size=input_size,
        num_classes=num_classes,
        project_name=project_name,
        base_path=base_path,
        framework="keras",
        config=config or {},
    )