train_utils.py

import os
import json
import time
from datetime import datetime
from typing import List, Tuple

import torch
import torch.nn as nn
import torch.optim as optim

from config import MODEL_DIR, META_DIR, DATASET_DISPLAY_NAME, CLASSICAL_MODEL_TYPES, IMAGE_SIZE
from data_utils import make_loaders
from metrics_utils import compute_classification_metrics, save_confusion_matrix_figure, save_loss_curve_figure
from model import SimpleCNN, BackboneWithFC, MLP


# ---------------------------------------------------------------------------
# Path helpers
# ---------------------------------------------------------------------------

def model_weight_path(model_name: str) -> str:
    return os.path.join(MODEL_DIR, f"{model_name}.pt")


def classifier_weight_path(model_name: str) -> str:
    return os.path.join(MODEL_DIR, f"{model_name}.joblib")


def model_meta_path(model_name: str) -> str:
    return os.path.join(META_DIR, f"{model_name}.json")


def list_saved_models() -> List[str]:
    names = []
    for fn in os.listdir(META_DIR):
        if fn.endswith(".json"):
            names.append(fn[:-5])
    return sorted(names, reverse=True)


def get_runtime_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ---------------------------------------------------------------------------
# Save / load
# ---------------------------------------------------------------------------

def save_model(model: nn.Module, model_name: str, config: dict, training_summary: dict):
    if config["model_type"] == "fc_head":
        state_dict = {k: v.detach().cpu() for k, v in model.classifier.state_dict().items()}
    else:
        state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}

    torch.save(state_dict, model_weight_path(model_name))

    with open(model_meta_path(model_name), "w", encoding="utf-8") as f:
        json.dump(
            {
                "model_name": model_name,
                "config": config,
                "training_summary": training_summary,
                "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            },
            f,
            indent=2,
            ensure_ascii=False,
        )


def _load_meta(model_name: str) -> dict:
    path = model_meta_path(model_name)
    if not os.path.exists(path):
        raise FileNotFoundError(f"Métadonnées introuvables : {model_name}")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_model(model_name: str, device: torch.device) -> Tuple[nn.Module, dict]:
    meta = _load_meta(model_name)
    cfg = meta["config"]
    model_type = cfg.get("model_type", "cnn")

    if model_type == "fc_head":
        from backbone_utils import load_backbone
        backbone = load_backbone(device)
        model = BackboneWithFC(backbone, cfg["num_classes"], cfg.get("dropout", 0.4), cfg.get("fc_dim", 256))
        model.classifier.load_state_dict(
            torch.load(model_weight_path(model_name), map_location="cpu")
        )

    elif model_type == "cnn":
        model = SimpleCNN(
            num_classes=cfg["num_classes"],
            num_conv_blocks=cfg.get("num_conv_blocks", 3),
            base_filters=cfg.get("base_filters", 32),
            kernel_size=cfg.get("kernel_size", 3),
            use_batchnorm=cfg.get("use_batchnorm", True),
            dropout=cfg.get("dropout", 0.4),
            fc_dim=cfg.get("fc_dim", 256),
        )
        model.load_state_dict(torch.load(model_weight_path(model_name), map_location="cpu"))

    elif model_type == "mlp":
        model = MLP(
            num_classes=cfg["num_classes"],
            input_size=cfg.get("input_size", 3 * IMAGE_SIZE * IMAGE_SIZE),
            num_layers=cfg.get("num_layers", 2),
            hidden_dim=cfg.get("hidden_dim", 256),
            dropout=cfg.get("dropout", 0.4),
        )
        model.load_state_dict(torch.load(model_weight_path(model_name), map_location="cpu"))

    else:
        raise ValueError(f"load_model n'accepte pas le type '{model_type}'. Utilisez load_classical_pipeline pour les modèles ML classiques.")

    model.to(device)
    model.eval()
    return model, meta


# ---------------------------------------------------------------------------
# Training helpers
# ---------------------------------------------------------------------------

def evaluate_loss_acc(model, loader, criterion, device):
    model.eval()
    total_loss, total, correct = 0.0, 0, 0

    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)
            correct += (outputs.argmax(1) == labels).sum().item()
            total += labels.size(0)

    return (total_loss / total if total else 0.0), (correct / total if total else 0.0)


def collect_predictions(model, loader, device):
    model.eval()
    y_true, y_pred = [], []

    with torch.no_grad():
        for images, labels in loader:
            outputs = model(images.to(device))
            y_pred.extend(outputs.argmax(1).detach().cpu().tolist())
            y_true.extend(labels.tolist())

    return y_true, y_pred


def _training_loop(model, train_loader, val_loader, criterion, optimizer, scheduler, epochs, device):
    history = []
    logs = []
    best_val_loss = float("inf")
    best_state = None

    for epoch in range(1, epochs + 1):
        model.train()
        running_loss, total, correct = 0.0, 0, 0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            running_loss += loss.item() * images.size(0)
            correct += (outputs.argmax(1) == labels).sum().item()
            total += labels.size(0)

        train_loss = running_loss / total if total else 0.0
        train_acc = correct / total if total else 0.0
        val_loss, val_acc = evaluate_loss_acc(model, val_loader, criterion, device)
        scheduler.step(val_loss)
        current_lr = optimizer.param_groups[0]["lr"]

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}

        history.append({
            "epoch": epoch,
            "train_loss": round(train_loss, 4),
            "train_acc": round(train_acc, 4),
            "val_loss": round(val_loss, 4),
            "val_acc": round(val_acc, 4),
        })
        logs.append(
            f"Époque {epoch}/{epochs} | "
            f"perte train={train_loss:.4f} acc train={train_acc:.4f} | "
            f"perte val={val_loss:.4f} acc val={val_acc:.4f} | "
            f"lr={current_lr:.2e}"
        )

    return history, logs, best_state, best_val_loss


# ---------------------------------------------------------------------------
# Train FC head on frozen backbone
# ---------------------------------------------------------------------------

def train_fc_head(
    dropout: float = 0.4,
    fc_dim: int = 256,
    learning_rate: float = 1e-4,
    weight_decay: float = 1e-4,
    batch_size: int = 16,
    epochs: int = 20,
    model_tag: str = "",
):
    from backbone_utils import load_backbone

    device = get_runtime_device()
    train_loader, val_loader, test_loader, class_names = make_loaders(batch_size)
    num_classes = len(class_names)

    backbone = load_backbone(device)

    model = BackboneWithFC(backbone, num_classes, dropout, fc_dim).to(device)

    trainable_params = sum(p.numel() for p in model.classifier.parameters())
    total_params = sum(p.numel() for p in model.parameters())

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.classifier.parameters(), lr=learning_rate, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=5, min_lr=learning_rate * 0.1
    )

    t0 = time.time()
    history, logs, best_state, best_val_loss = _training_loop(
        model, train_loader, val_loader, criterion, optimizer, scheduler, epochs, device
    )

    model.load_state_dict(best_state)
    test_loss, test_acc = evaluate_loss_acc(model, test_loader, criterion, device)
    y_true, y_pred = collect_predictions(model, test_loader, device)
    metrics = compute_classification_metrics(y_true, y_pred, class_names)
    elapsed = time.time() - t0

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_tag = model_tag.strip().replace(" ", "_") if model_tag.strip() else "fc_head"
    model_name = f"{safe_tag}_{timestamp}"

    cm_path = save_confusion_matrix_figure(metrics["confusion_matrix"], model_name)

    config = {
        "dataset_name": DATASET_DISPLAY_NAME,
        "model_type": "fc_head",
        "architecture": f"ResNet18 backbone (gelé) + FC({fc_dim})",
        "num_classes": num_classes,
        "class_names": class_names,
        "dropout": dropout,
        "fc_dim": fc_dim,
        "learning_rate": learning_rate,
        "weight_decay": weight_decay,
        "batch_size": batch_size,
        "epochs": epochs,
    }

    training_summary = {
        "final_train_loss": history[-1]["train_loss"] if history else None,
        "final_train_acc": history[-1]["train_acc"] if history else None,
        "best_val_loss": round(best_val_loss, 4),
        "final_val_loss": history[-1]["val_loss"] if history else None,
        "final_val_acc": history[-1]["val_acc"] if history else None,
        "test_cross_entropy_loss": round(test_loss, 4),
        "test_accuracy": round(test_acc, 4),
        "test_f1_macro": metrics["f1_macro"],
        "test_f1_weighted": metrics["f1_weighted"],
        "elapsed_seconds": round(elapsed, 2),
        "device": str(device),
        "total_params": total_params,
        "trainable_params": trainable_params,
    }

    save_model(model, model_name, config, training_summary)

    logs += [
        "",
        "Entraînement terminé.",
        f"Modèle sauvegardé : {model_name}",
        f"Architecture : {config['architecture']}",
        f"Paramètres entraînables : {trainable_params} / {total_params}",
        f"Perte test : {test_loss:.4f}  |  Accuracy test : {test_acc:.4f}",
        f"F1 macro : {metrics['f1_macro']:.4f}  |  F1 pondéré : {metrics['f1_weighted']:.4f}",
        f"Temps : {elapsed:.1f}s  |  Appareil : {device}",
    ]

    return {
        "logs": "\n".join(logs),
        "history": history,
        "summary": training_summary,
        "model_name": model_name,
        "classification_report": metrics["classification_report"],
        "confusion_matrix": metrics["confusion_matrix"],
        "confusion_matrix_path": cm_path,
    }


# ---------------------------------------------------------------------------
# Train SimpleCNN from scratch
# ---------------------------------------------------------------------------

def train_cnn(
    num_conv_blocks: int = 3,
    base_filters: int = 32,
    kernel_size: int = 3,
    use_batchnorm: bool = True,
    dropout: float = 0.4,
    fc_dim: int = 256,
    learning_rate: float = 1e-3,
    weight_decay: float = 1e-4,
    batch_size: int = 16,
    epochs: int = 30,
    model_tag: str = "",
):
    device = get_runtime_device()
    train_loader, val_loader, test_loader, class_names = make_loaders(batch_size)
    num_classes = len(class_names)

    model = SimpleCNN(
        num_classes=num_classes,
        num_conv_blocks=num_conv_blocks,
        base_filters=base_filters,
        kernel_size=kernel_size,
        use_batchnorm=use_batchnorm,
        dropout=dropout,
        fc_dim=fc_dim,
    ).to(device)

    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=8, min_lr=learning_rate * 0.2
    )

    t0 = time.time()
    history, logs, best_state, best_val_loss = _training_loop(
        model, train_loader, val_loader, criterion, optimizer, scheduler, epochs, device
    )

    model.load_state_dict(best_state)
    test_loss, test_acc = evaluate_loss_acc(model, test_loader, criterion, device)
    y_true, y_pred = collect_predictions(model, test_loader, device)
    metrics = compute_classification_metrics(y_true, y_pred, class_names)
    elapsed = time.time() - t0

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_tag = model_tag.strip().replace(" ", "_") if model_tag.strip() else "cnn"
    model_name = f"{safe_tag}_{timestamp}"

    cm_path = save_confusion_matrix_figure(metrics["confusion_matrix"], model_name)
    loss_curve_path = save_loss_curve_figure(history, model_name)

    architecture = f"CNN simple ({num_conv_blocks} blocs, filtres={base_filters}, noyau={kernel_size}×{kernel_size})"

    config = {
        "dataset_name": DATASET_DISPLAY_NAME,
        "model_type": "cnn",
        "architecture": architecture,
        "num_classes": num_classes,
        "class_names": class_names,
        "num_conv_blocks": num_conv_blocks,
        "base_filters": base_filters,
        "kernel_size": kernel_size,
        "use_batchnorm": use_batchnorm,
        "dropout": dropout,
        "fc_dim": fc_dim,
        "learning_rate": learning_rate,
        "weight_decay": weight_decay,
        "batch_size": batch_size,
        "epochs": epochs,
    }

    training_summary = {
        "final_train_loss": history[-1]["train_loss"] if history else None,
        "final_train_acc": history[-1]["train_acc"] if history else None,
        "best_val_loss": round(best_val_loss, 4),
        "final_val_loss": history[-1]["val_loss"] if history else None,
        "final_val_acc": history[-1]["val_acc"] if history else None,
        "test_cross_entropy_loss": round(test_loss, 4),
        "test_accuracy": round(test_acc, 4),
        "test_f1_macro": metrics["f1_macro"],
        "test_f1_weighted": metrics["f1_weighted"],
        "elapsed_seconds": round(elapsed, 2),
        "device": str(device),
        "total_params": total_params,
        "trainable_params": trainable_params,
    }

    save_model(model, model_name, config, training_summary)

    logs += [
        "",
        "Entraînement terminé.",
        f"Modèle sauvegardé : {model_name}",
        f"Architecture : {architecture}",
        f"Paramètres : {total_params}",
        f"Perte test : {test_loss:.4f}  |  Accuracy test : {test_acc:.4f}",
        f"F1 macro : {metrics['f1_macro']:.4f}  |  F1 pondéré : {metrics['f1_weighted']:.4f}",
        f"Temps : {elapsed:.1f}s  |  Appareil : {device}",
    ]

    return {
        "logs": "\n".join(logs),
        "history": history,
        "summary": training_summary,
        "model_name": model_name,
        "classification_report": metrics["classification_report"],
        "confusion_matrix": metrics["confusion_matrix"],
        "confusion_matrix_path": cm_path,
        "loss_curve_path": loss_curve_path,
    }


# ---------------------------------------------------------------------------
# Train MLP from scratch (baseline avant le CNN)
# ---------------------------------------------------------------------------

def train_mlp(
    num_layers: int = 2,
    hidden_dim: int = 256,
    dropout: float = 0.4,
    learning_rate: float = 1e-3,
    weight_decay: float = 1e-4,
    batch_size: int = 16,
    epochs: int = 30,
    model_tag: str = "",
):
    device = get_runtime_device()
    train_loader, val_loader, test_loader, class_names = make_loaders(batch_size)
    num_classes = len(class_names)
    input_size = 3 * IMAGE_SIZE * IMAGE_SIZE

    model = MLP(
        num_classes=num_classes,
        input_size=input_size,
        num_layers=num_layers,
        hidden_dim=hidden_dim,
        dropout=dropout,
    ).to(device)

    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=8, min_lr=learning_rate * 0.2
    )

    t0 = time.time()
    history, logs, best_state, best_val_loss = _training_loop(
        model, train_loader, val_loader, criterion, optimizer, scheduler, epochs, device
    )

    model.load_state_dict(best_state)
    test_loss, test_acc = evaluate_loss_acc(model, test_loader, criterion, device)
    y_true, y_pred = collect_predictions(model, test_loader, device)
    metrics = compute_classification_metrics(y_true, y_pred, class_names)
    elapsed = time.time() - t0

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_tag = model_tag.strip().replace(" ", "_") if model_tag.strip() else "mlp"
    model_name = f"{safe_tag}_{timestamp}"

    cm_path = save_confusion_matrix_figure(metrics["confusion_matrix"], model_name)
    loss_curve_path = save_loss_curve_figure(history, model_name)

    architecture = f"MLP ({num_layers} couches cachées de {hidden_dim} neurones)"

    config = {
        "dataset_name": DATASET_DISPLAY_NAME,
        "model_type": "mlp",
        "architecture": architecture,
        "num_classes": num_classes,
        "class_names": class_names,
        "input_size": input_size,
        "num_layers": num_layers,
        "hidden_dim": hidden_dim,
        "dropout": dropout,
        "learning_rate": learning_rate,
        "weight_decay": weight_decay,
        "batch_size": batch_size,
        "epochs": epochs,
    }

    training_summary = {
        "final_train_loss": history[-1]["train_loss"] if history else None,
        "final_train_acc": history[-1]["train_acc"] if history else None,
        "best_val_loss": round(best_val_loss, 4),
        "final_val_loss": history[-1]["val_loss"] if history else None,
        "final_val_acc": history[-1]["val_acc"] if history else None,
        "test_cross_entropy_loss": round(test_loss, 4),
        "test_accuracy": round(test_acc, 4),
        "test_f1_macro": metrics["f1_macro"],
        "test_f1_weighted": metrics["f1_weighted"],
        "elapsed_seconds": round(elapsed, 2),
        "device": str(device),
        "total_params": total_params,
        "trainable_params": trainable_params,
    }

    save_model(model, model_name, config, training_summary)

    logs += [
        "",
        "Entraînement terminé.",
        f"Modèle sauvegardé : {model_name}",
        f"Architecture : {architecture}",
        f"Paramètres : {total_params}",
        f"Perte test : {test_loss:.4f}  |  Accuracy test : {test_acc:.4f}",
        f"F1 macro : {metrics['f1_macro']:.4f}  |  F1 pondéré : {metrics['f1_weighted']:.4f}",
        f"Temps : {elapsed:.1f}s  |  Appareil : {device}",
    ]

    return {
        "logs": "\n".join(logs),
        "history": history,
        "summary": training_summary,
        "model_name": model_name,
        "classification_report": metrics["classification_report"],
        "confusion_matrix": metrics["confusion_matrix"],
        "confusion_matrix_path": cm_path,
        "loss_curve_path": loss_curve_path,
    }


# ---------------------------------------------------------------------------
# Evaluate any saved model
# ---------------------------------------------------------------------------

def evaluate_saved_model(model_name: str):
    if not model_name:
        raise ValueError("Aucun modèle sélectionné.")

    meta = _load_meta(model_name)
    model_type = meta["config"].get("model_type", "cnn")

    if model_type in CLASSICAL_MODEL_TYPES:
        return _evaluate_classical(model_name, meta)
    else:
        return _evaluate_neural(model_name, meta)


def _evaluate_neural(model_name: str, meta: dict):
    device = get_runtime_device()
    model, meta = load_model(model_name, device)

    batch_size = int(meta["config"].get("batch_size", 16))
    _, _, test_loader, class_names = make_loaders(batch_size)

    criterion = nn.CrossEntropyLoss()
    test_loss, test_acc = evaluate_loss_acc(model, test_loader, criterion, device)
    y_true, y_pred = collect_predictions(model, test_loader, device)

    metrics = compute_classification_metrics(y_true, y_pred, class_names)
    cm_path = save_confusion_matrix_figure(metrics["confusion_matrix"], model_name)

    return (
        {
            "test_cross_entropy_loss": round(test_loss, 4),
            "test_accuracy": round(test_acc, 4),
            "test_f1_macro": metrics["f1_macro"],
            "test_f1_weighted": metrics["f1_weighted"],
            "device": str(device),
        },
        metrics["classification_report"],
        metrics["confusion_matrix"],
        cm_path,
    )


def _evaluate_classical(model_name: str, meta: dict):
    from backbone_utils import get_cached_features, extract_all_features
    from classical_ml_utils import load_classical_pipeline

    features_cache = get_cached_features()
    if features_cache is None:
        features_cache, _, _ = extract_all_features()

    class_names = meta["config"]["class_names"]
    pipeline = load_classical_pipeline(model_name)

    X_test = features_cache["test"]["X"]
    y_test = features_cache["test"]["y"]
    y_pred = pipeline.predict(X_test)

    metrics = compute_classification_metrics(y_test.tolist(), y_pred.tolist(), class_names)
    cm_path = save_confusion_matrix_figure(metrics["confusion_matrix"], model_name)

    return (
        {
            "test_accuracy": metrics["accuracy"],
            "test_f1_macro": metrics["f1_macro"],
            "test_f1_weighted": metrics["f1_weighted"],
        },
        metrics["classification_report"],
        metrics["confusion_matrix"],
        cm_path,
    )