"""
models/train.py
Classe Trainer pour PyTorch.
Fonctionnalités : early stopping, ReduceLROnPlateau,
                  sauvegarde du meilleur modèle, courbes train/val.
"""

import torch
import torch.nn as nn
from tqdm import tqdm
import matplotlib.pyplot as plt


class Trainer:
    def __init__(self, model, train_dataloader, test_dataloader,
                 lr=1e-3, epochs=30, device="cpu", patience=5):
        self.model            = model
        self.train_dataloader = train_dataloader
        self.test_dataloader  = test_dataloader
        self.epochs           = epochs
        self.patience         = patience
        self.device           = device
        self.criterion        = nn.CrossEntropyLoss()
        self.optimizer        = torch.optim.Adam(model.parameters(), lr=lr)
        self.scheduler        = torch.optim.lr_scheduler.ReduceLROnPlateau(
                                    self.optimizer, mode="min",
                                    factor=0.5, patience=3)

    # ── Entraînement complet ───────────────────────────────────────────────
    def train(self, save_path=None, plot=False):
        self.train_loss, self.train_acc = [], []
        self.val_loss,   self.val_acc   = [], []

        best_val_loss     = float("inf")
        epochs_no_improve = 0
        best_state        = None

        for epoch in range(self.epochs):
            tr_loss, tr_acc = self._train_one_epoch(epoch)
            v_loss,  v_acc  = self._validate()

            self.train_loss.append(tr_loss)
            self.train_acc.append(tr_acc)
            self.val_loss.append(v_loss)
            self.val_acc.append(v_acc)

            self.scheduler.step(v_loss)
            lr = self.optimizer.param_groups[0]["lr"]

            print(f"Epoch {epoch+1:02d}/{self.epochs} "
                  f"| Train  loss={tr_loss:.4f}  acc={tr_acc:.2f}% "
                  f"| Val    loss={v_loss:.4f}  acc={v_acc:.2f}% "
                  f"| LR={lr:.2e}")

            # Early stopping
            if v_loss < best_val_loss:
                best_val_loss     = v_loss
                epochs_no_improve = 0
                best_state = {k: v.clone() for k, v in self.model.state_dict().items()}
                if save_path:
                    torch.save(best_state, save_path)
                    print(f"  ✓ Best model saved  (val_loss={v_loss:.4f})")
            else:
                epochs_no_improve += 1
                print(f"  ⚠ No improvement {epochs_no_improve}/{self.patience}")
                if epochs_no_improve >= self.patience:
                    print(f"\n⛔ Early stopping at epoch {epoch+1}")
                    break

        if best_state:
            self.model.load_state_dict(best_state)
        if plot:
            self.plot_history()

    # ── Une epoch de train ─────────────────────────────────────────────────
    def _train_one_epoch(self, epoch):
        self.model.train()
        total_loss, total_correct, total_samples = 0, 0, 0
        pbar = tqdm(self.train_dataloader,
                    desc=f"Epoch {epoch+1}/{self.epochs} [train]", leave=False)

        for imgs, labels in pbar:
            imgs, labels = imgs.to(self.device), labels.to(self.device)
            self.optimizer.zero_grad()
            out  = self.model(imgs)
            loss = self.criterion(out, labels)
            loss.backward()
            self.optimizer.step()

            _, preds       = out.max(1)
            correct        = (preds == labels).sum().item()
            total          = labels.size(0)
            total_correct += correct
            total_samples += total
            total_loss    += loss.item()

            pbar.set_postfix({
                "Batch Acc": f"{100.*correct/total:.1f}%",
                "Avg Acc":   f"{100.*total_correct/total_samples:.1f}%",
                "Loss":      f"{total_loss/total_samples:.4f}",
            })

        return total_loss / total_samples, 100. * total_correct / total_samples

    # ── Validation ────────────────────────────────────────────────────────
    @torch.no_grad()
    def _validate(self):
        self.model.eval()
        total_loss, total_correct, total_samples = 0, 0, 0
        for imgs, labels in self.test_dataloader:
            imgs, labels  = imgs.to(self.device), labels.to(self.device)
            out           = self.model(imgs)
            loss          = self.criterion(out, labels)
            _, preds      = out.max(1)
            total_correct += (preds == labels).sum().item()
            total_samples += labels.size(0)
            total_loss    += loss.item() * labels.size(0)
        return total_loss / total_samples, 100. * total_correct / total_samples

    # ── Évaluation finale (public) ────────────────────────────────────────
    @torch.no_grad()
    def evaluate(self):
        loss, acc = self._validate()
        print(f"\nTest Accuracy : {acc:.2f}%  |  Test Loss : {loss:.4f}")
        return acc, loss

    # ── Courbes ───────────────────────────────────────────────────────────
    def plot_history(self, save_path="/kaggle/working/history_pytorch.png"):
        epochs = range(1, len(self.train_loss) + 1)
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

        ax1.plot(epochs, self.train_loss, label="Train", color="tab:blue")
        ax1.plot(epochs, self.val_loss,   label="Val",   color="tab:orange")
        ax1.set_title("Loss"); ax1.set_xlabel("Epoch")
        ax1.legend(); ax1.grid(alpha=.3)

        ax2.plot(epochs, self.train_acc, label="Train", color="tab:blue")
        ax2.plot(epochs, self.val_acc,   label="Val",   color="tab:orange")
        ax2.set_title("Accuracy (%)"); ax2.set_xlabel("Epoch")
        ax2.legend(); ax2.grid(alpha=.3)

        fig.suptitle("Training History — PyTorch", fontsize=13)
        fig.tight_layout()
        plt.savefig(save_path, dpi=120)
        plt.show()
        print(f"✓ Courbes sauvegardées → {save_path}")