"""
Implementação concreta (PyTorch) 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

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# 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


# ---------------------------------------------------------------------------
# Arquitetura MLP
# ---------------------------------------------------------------------------
class _MLP(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden_dims: List[int],
        num_classes: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        layers: List[nn.Module] = []
        prev = input_size
        for h in hidden_dims:
            layers.append(nn.Linear(prev, h))
            layers.append(nn.ReLU())
            if dropout and dropout > 0:
                layers.append(nn.Dropout(dropout))
            prev = h
        layers.append(nn.Linear(prev, num_classes))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # noqa: D401
        return self.net(x)


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

    def build_model(self) -> nn.Module:
        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)))

        # device
        self.device = torch.device(
            "cuda"
            if torch.cuda.is_available() and self.config.get("use_cuda", True)
            else "cpu"
        )

        model = _MLP(
            input_size=self.input_size,
            hidden_dims=list(hidden_dims),
            num_classes=self.num_classes,
            dropout=dropout,
        ).to(self.device)

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

    # ---------------- treinamento ----------------
    def train_model(
        self,
        train_loader: DataLoader,
        val_loader: Optional[DataLoader],
        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()

        optimizer = self._build_optimizer(optimizer_name, lr, weight_decay)

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

        n_batches = max(len(train_loader), 1)
        total_train_steps = max(epochs * n_batches, 1)
        best_val_acc = -float("inf")
        best_epoch = -1
        for epoch in range(1, epochs + 1):
            self.model.train()
            run_loss = 0.0
            correct = 0
            total = 0
            for batch_idx, (X, y) in enumerate(train_loader, start=1):
                X = X.to(self.device)
                y = y.to(self.device)

                optimizer.zero_grad()
                logits = self.model(X)
                loss = self.criterion(logits, y)
                loss.backward()
                optimizer.step()

                run_loss += float(loss.item()) * X.size(0)
                preds = logits.argmax(dim=1)
                correct += int((preds == y).sum().item())
                total += int(y.size(0))

            global_step = (epoch - 1) * n_batches
            inner_pct = int(global_step * 100 / total_train_steps)
            emit_progress(inner_pct, total_train_steps)

            train_loss = run_loss / max(total, 1)
            train_acc = correct / max(total, 1)
            history["train_loss"].append(round(train_loss, 6))
            history["train_acc"].append(round(train_acc, 6))

            val_loss, val_acc = (float("nan"), float("nan"))
            if val_loader is not None:
                val_metrics = self.evaluate(val_loader)
                val_loss = val_metrics["loss"]
                val_acc = val_metrics["accuracy"]
                if val_acc > best_val_acc:
                    best_val_acc = val_acc
                    best_epoch = epoch
            history["val_loss"].append(round(val_loss, 6))
            history["val_acc"].append(round(val_acc, 6))

            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: DataLoader) -> Dict[str, float]:
        self.model.eval()
        loss_sum = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for X, y in data_loader:
                X = X.to(self.device)
                y = y.to(self.device)
                logits = self.model(X)
                loss = self.criterion(logits, y)
                loss_sum += float(loss.item()) * X.size(0)
                preds = logits.argmax(dim=1)
                correct += int((preds == y).sum().item())
                total += int(y.size(0))
        return {
            "loss": loss_sum / max(total, 1),
            "accuracy": correct / max(total, 1),
            "n": total,
        }

    # ---------------- inferência ----------------
    def predict(self, inputs: Any) -> np.ndarray:
        self.model.eval()
        with torch.no_grad():
            if isinstance(inputs, DataLoader):
                outs: List[np.ndarray] = []
                for batch in inputs:
                    X = batch[0] if isinstance(batch, (tuple, list)) else batch
                    X = X.to(self.device)
                    outs.append(self.model(X).argmax(dim=1).cpu().numpy())
                return np.concatenate(outs, axis=0)
            if isinstance(inputs, np.ndarray):
                X = torch.from_numpy(inputs.astype(np.float32)).to(self.device)
            elif isinstance(inputs, torch.Tensor):
                X = inputs.to(self.device)
            else:
                X = torch.tensor(inputs, dtype=torch.float32).to(self.device)
            if X.ndim == 1:
                X = X.unsqueeze(0)
            return self.model(X).argmax(dim=1).cpu().numpy()

    # ---------------- persistência ----------------
    def save_model(self, filename: str) -> str:
        path = os.path.join(self.model_dir, filename)
        torch.save(
            {
                "state_dict": self.model.state_dict(),
                "input_size": self.input_size,
                "num_classes": self.num_classes,
                "hidden_dims": self._hidden_dims,
                "dropout": self._dropout,
                "project_name": self.project_name,
            },
            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
        )
        ckpt = torch.load(path, map_location=self.device)
        self._hidden_dims = ckpt.get("hidden_dims", self._hidden_dims)
        self._dropout = ckpt.get("dropout", self._dropout)
        self.model = _MLP(
            input_size=ckpt["input_size"],
            hidden_dims=self._hidden_dims,
            num_classes=ckpt["num_classes"],
            dropout=self._dropout,
        ).to(self.device)
        self.model.load_state_dict(ckpt["state_dict"])

    # ---------------- helpers ----------------
    def _build_optimizer(
        self, name: str, lr: float, weight_decay: float
    ) -> optim.Optimizer:
        params = self.model.parameters()
        if name == "sgd":
            return optim.SGD(params, lr=lr, momentum=0.9, weight_decay=weight_decay)
        if name == "rmsprop":
            return optim.RMSprop(params, lr=lr, weight_decay=weight_decay)
        # default: adam
        return optim.Adam(params, lr=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="pytorch",
        config=config or {},
    )