"""
finetune_doctr.py
------------------
Fine-tuning de crnn_vgg16_bn (doctr) sur les valeurs LCD
de pompes à carburant.

Copie de référence pour la reproductibilité (documentation technique).
Ce script attend le dataset généré par `annotator/` à la racine du dépôt
d'origine (../annotator/dataset_doctr/) — il n'est pas autonome dans ce
dossier livrable, qui ne contient volontairement pas les ~150 images
sources annotées. Voir docs/LIMITATIONS.md pour le contexte du dataset.

Dataset attendu (généré par prepare_doctr_dataset.py) :
  annotator/dataset_doctr/
    train/images/ + labels.json
    val/images/   + labels.json

Usage (depuis le dépôt d'origine, dossier train/) :
  python finetune_doctr.py
  python finetune_doctr.py --epochs 50 --lr 3e-5
  python finetune_doctr.py --from-scratch   # si pas de connexion internet
"""

import os, json, argparse, time
import torch
from torch.utils.data import DataLoader
from doctr.datasets import RecognitionDataset, VOCABS
from doctr.models import crnn_vgg16_bn

# ─── Chemins (depuis train/) ──────────────────────────────────────────────────
DATASET_DIR = os.path.join("..", "annotator", "dataset_doctr")
OUTPUT_DIR  = os.path.join("models", "v2")  # ne pas écraser models/crnn_fuel_pump_best.pt (baseline)
VOCAB       = VOCABS["french"]   # contient 0-9, virgule, point


# ─── Dataloader ───────────────────────────────────────────────────────────────

IMG_H = 32

IMG_W = 256

def _resize_preserve_aspect(img, target_h: int, target_w: int):
    """Redimensionne en conservant le ratio d'aspect (hauteur fixe, largeur au
    prorata) puis complète par du padding noir jusqu'à target_w, au lieu
    d'étirer l'image. C'est ce que fait déjà `preprocess_crop` à l'inférence
    (src/fine_tuned_inference.py) : entraîner avec un étirement forcé alors
    que l'inférence préserve le ratio créait un décalage systématique entre
    les chiffres vus à l'entraînement (déformés) et en production (non
    déformés), ce qui pénalisait la précision du modèle.
    """
    import torch.nn.functional as F

    c, h, w = img.shape
    new_w = max(1, min(target_w, round(target_h * w / max(h, 1))))
    resized = F.interpolate(img.unsqueeze(0), size=(target_h, new_w),
                            mode="bilinear", align_corners=False).squeeze(0)
    if new_w < target_w:
        pad = torch.zeros(c, target_h, target_w - new_w, dtype=resized.dtype)
        resized = torch.cat([resized, pad], dim=2)
    elif new_w > target_w:
        resized = resized[:, :, :target_w]
    return resized


def collate_fn(batch):
    """Redimensionne toutes les images à taille fixe pour le CRNN, en
    préservant le ratio d'aspect (voir _resize_preserve_aspect)."""

    imgs, targets = zip(*batch)

    resized = [_resize_preserve_aspect(img, IMG_H, IMG_W) for img in imgs]

    return torch.stack(resized, 0), list(targets)

def build_loaders(batch_size: int):
    paths = {
        "train_img": os.path.join(DATASET_DIR, "train", "images"),
        "train_lbl": os.path.join(DATASET_DIR, "train", "labels.json"),
        "val_img":   os.path.join(DATASET_DIR, "val",   "images"),
        "val_lbl":   os.path.join(DATASET_DIR, "val",   "labels.json"),
    }
    for k, p in paths.items():
        if not os.path.exists(p):
            raise FileNotFoundError(
                f"❌ Introuvable : {p}\n"
                f"   Lance d'abord : python prepare_doctr_dataset.py")

    train_ds = RecognitionDataset(
        img_folder=paths["train_img"], labels_path=paths["train_lbl"])
    val_ds = RecognitionDataset(
        img_folder=paths["val_img"], labels_path=paths["val_lbl"])

    print(f"📦 Dataset : {len(train_ds)} train | {len(val_ds)} val")

    trn = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                     num_workers=0, collate_fn=collate_fn, drop_last=False)
    val = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
                     num_workers=0, collate_fn=collate_fn)
    return trn, val, len(train_ds), len(val_ds)


# ─── Modèle ───────────────────────────────────────────────────────────────────

def build_model(from_scratch: bool):
    if from_scratch:
        print("⚠️  Mode from-scratch (poids aléatoires).")
        return crnn_vgg16_bn(
            pretrained=False, pretrained_backbone=False, vocab=VOCAB)
    try:
        print("⬇️  Chargement des poids pré-entraînés Mindee...")
        m = crnn_vgg16_bn(pretrained=True, vocab=VOCAB)
        print("✅ Poids pré-entraînés chargés.")
        return m
    except Exception as e:
        print(f"❌ Téléchargement échoué : {e}")
        print("   Relance avec --from-scratch si pas de connexion internet.")
        raise SystemExit(1)


# ─── Évaluation ───────────────────────────────────────────────────────────────

@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    total_loss = n_ok = n_tot = 0
    samples = []

    for imgs, targets in loader:
        imgs = imgs.to(device)
        out  = model(imgs, target=targets, return_preds=True)
        total_loss += out["loss"].item()
        for (pred, _), gt in zip(out["preds"], targets):
            n_tot += 1
            ok = pred.strip() == gt.strip()
            if ok: n_ok += 1
            if len(samples) < 8:
                samples.append((gt, pred, "✅" if ok else "❌"))

    return total_loss / max(len(loader), 1), n_ok / max(n_tot, 1), samples


# ─── Entraînement ─────────────────────────────────────────────────────────────

def train(epochs, lr, batch_size, from_scratch, patience):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"🖥️  Device : {device}")

    trn_loader, val_loader, n_trn, n_val = build_loaders(batch_size)
    model = build_model(from_scratch)
    model.to(device)

    # LR faible car fine-tuning — on préserve les poids pré-entraînés
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=3)

    os.makedirs(OUTPUT_DIR, exist_ok=True)
    best_loss = float("inf")
    best_acc  = 0.0
    no_improv = 0
    history   = []

    print(f"\n🚀 Fine-tuning : {epochs} epochs | lr={lr} | batch={batch_size}")
    print("─" * 65)

    for epoch in range(1, epochs + 1):
        model.train()
        t0 = time.time()
        tr_loss = 0.0
        nb = 0

        for imgs, targets in trn_loader:
            imgs = imgs.to(device)
            optimizer.zero_grad()
            out  = model(imgs, target=targets, return_preds=False)
            loss = out["loss"]
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()
            tr_loss += loss.item()
            nb += 1

        avg_tr = tr_loss / max(nb, 1)
        vl, acc, samples = evaluate(model, val_loader, device)
        scheduler.step(vl)
        cur_lr = optimizer.param_groups[0]["lr"]

        print(f"Ep {epoch:3d}/{epochs} | "
              f"train={avg_tr:.4f} | val={vl:.4f} | "
              f"acc={acc:.1%} | lr={cur_lr:.1e} | {time.time()-t0:.0f}s")

        # Afficher des exemples de prédictions toutes les 5 epochs
        if epoch % 5 == 0 or epoch == 1:
            print("  Exemples val :")
            for gt, pred, status in samples:
                print(f"    {status}  GT='{gt}'  →  PRED='{pred}'")

        history.append({"epoch": epoch, "train_loss": avg_tr,
                        "val_loss": vl, "val_acc": acc, "lr": cur_lr})

        # Sauvegarder le meilleur modèle. Écrire dans un fichier temporaire
        # puis renommer (os.replace, atomique) évite un crash Windows déjà
        # observé (RuntimeError code 1224, fichier verrouillé pendant
        # l'écriture directe — probablement un scan antivirus).
        if vl < best_loss:
            best_loss = vl
            best_acc  = acc
            no_improv = 0
            best_path = os.path.join(OUTPUT_DIR, "crnn_fuel_pump_best.pt")
            tmp_path = best_path + ".tmp"
            torch.save(model.state_dict(), tmp_path)
            os.replace(tmp_path, best_path)
            print(f"  💾 Meilleur modèle → {best_path}")
        else:
            no_improv += 1

        if no_improv >= patience:
            print(f"\n⏹  Early stopping ({patience} epochs sans amélioration)")
            break

    # Sauvegarde finale + historique
    final = os.path.join(OUTPUT_DIR, "crnn_fuel_pump_final.pt")
    torch.save(model.state_dict(), final)
    with open(os.path.join(OUTPUT_DIR, "history.json"), "w") as f:
        json.dump({"history": history, "best_val_loss": best_loss,
                   "best_val_acc": best_acc}, f, indent=2)

    print("\n" + "=" * 65)
    print(f"  ✅ Entraînement terminé")
    print(f"  Meilleure val_acc  : {best_acc:.1%}")
    print(f"  Meilleur modèle    : {best_path}")
    print(f"  Modèle final       : {final}")
    print("=" * 65)
    print("\n👉 Prochaine étape : python test_model.py")


# ─── Entrée ───────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs",       type=int,   default=80)
    parser.add_argument("--lr",           type=float, default=5e-5)
    parser.add_argument("--batch-size",   type=int,   default=8)
    parser.add_argument("--patience",     type=int,   default=8)
    parser.add_argument("--from-scratch", action="store_true")
    args = parser.parse_args()

    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    train(args.epochs, args.lr, args.batch_size,
          args.from_scratch, args.patience)