"""Phase 5 — entraînement de l'alignement image-texte (CLIP maison).

Boucle écrite à la main : on encode les images et les légendes d'un lot, on calcule
la perte contrastive symétrique, on met à jour les deux encodeurs + les projecteurs.

Usage :

    python -m src.train.train_clip --data-dir data/processed/toy_clip \\
        --out-dir checkpoints/clip_toy --epochs 30 --batch-size 64
"""

from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import torch

from src.data.prepare_image_text import load_pairs
from src.model.clip import (
    ClipConfig,
    CortexClip,
    class_retrieval_accuracy,
    clip_loss,
)
from src.tokenizer import load_tokenizer
from src.tokenizer.special_tokens import BOS_ID, PAD_ID
from src.train.train_text import pick_device


@dataclass
class ClipTrainConfig:
    data_dir: str = "data/processed/toy_clip"
    out_dir: str = "checkpoints/clip_toy"
    epochs: int = 30
    batch_size: int = 64
    lr: float = 1e-3
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    max_len: int = 16
    embed_dim: int = 64
    text_layers: int = 2
    text_embd: int = 128
    vision_layers: int = 2
    vision_embd: int = 128
    patch_size: int = 8
    log_interval: int = 10
    seed: int = 1337
    device: str = "auto"


def encode_captions(tokenizer, captions: list[str], max_len: int, device):
    """Tokenise + tronque/pad à ``max_len``. Retourne ``(ids, mask)``.

    ``mask`` vaut 1 sur les vrais tokens (BOS inclus) et 0 sur le padding, afin que
    le pooling du texte **ignore les PAD** (sinon toutes les légendes se ressemblent).
    """
    rows, masks = [], []
    for cap in captions:
        ids = [BOS_ID, *tokenizer.encode(cap)][:max_len]
        n_real = len(ids)
        ids += [PAD_ID] * (max_len - n_real)
        rows.append(ids)
        masks.append([1] * n_real + [0] * (max_len - n_real))
    return (torch.tensor(rows, dtype=torch.long, device=device),
            torch.tensor(masks, dtype=torch.long, device=device))


def build_model(cfg: ClipTrainConfig, tokenizer, image_size: int = 32) -> CortexClip:
    from src.model.text_model import TextConfig
    from src.model.vit import ViTConfig

    return CortexClip(ClipConfig(
        embed_dim=cfg.embed_dim,
        text=TextConfig(vocab_size=tokenizer.vocab_size, n_layer=cfg.text_layers,
                        n_head=4, n_embd=cfg.text_embd, ctx_len=cfg.max_len,
                        tie_embeddings=True),
        vision=ViTConfig(image_size=image_size, patch_size=cfg.patch_size,
                         n_layer=cfg.vision_layers, n_head=4, n_embd=cfg.vision_embd),
    ))


def _build_batches(captions: list[str], batch_size: int, rng: torch.Generator):
    """Construit des lots où **chaque légende apparaît une seule fois**.

    Indispensable : avec des légendes dupliquées, l'appariement image↔texte n'est
    plus 1-1 et la perte contrastive symétrique devient contradictoire (collapse).
    On regroupe donc les images par légende et on tire au plus une image par groupe
    et par lot.
    """
    groups: dict[str, list[int]] = {}
    for i, cap in enumerate(captions):
        groups.setdefault(cap, []).append(i)
    group_keys = sorted(groups)
    b = max(2, min(batch_size, len(group_keys)))

    order = torch.randperm(len(group_keys), generator=rng).tolist()
    for start in range(0, len(group_keys), b):
        chunk = order[start:start + b]
        if len(chunk) < 2:
            continue
        idx = []
        for g in chunk:
            choices = groups[group_keys[g]]
            pick = torch.randint(len(choices), (1,), generator=rng).item()
            idx.append(choices[pick])
        yield idx


def train(cfg: ClipTrainConfig) -> dict:
    torch.manual_seed(cfg.seed)
    device = pick_device(cfg.device)
    tokenizer = load_tokenizer()
    out_dir = Path(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    tr_imgs, tr_caps = load_pairs(cfg.data_dir, "train")
    va_imgs, va_caps = load_pairs(cfg.data_dir, "val")
    image_size = int(tr_imgs.shape[-1])
    model = build_model(cfg, tokenizer, image_size=image_size).to(device)

    # Banque de légendes **uniques** (les captions se répètent par classe) : c'est la
    # référence correcte pour mesurer le retrieval image->texte.
    bank = sorted(set(tr_caps))
    bank_index = {cap: i for i, cap in enumerate(bank)}
    bank_ids, bank_mask = encode_captions(tokenizer, bank, cfg.max_len, device)

    decay, no_decay = [], []
    for _n, p in model.named_parameters():
        if p.requires_grad:
            (decay if p.dim() >= 2 else no_decay).append(p)
    optimizer = torch.optim.AdamW(
        [{"params": decay, "weight_decay": cfg.weight_decay},
         {"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr)

    history, t0 = [], time.time()
    best_acc = -1.0
    rng = torch.Generator().manual_seed(cfg.seed)
    for epoch in range(cfg.epochs):
        epoch_loss, steps = 0.0, 0
        for idx in _build_batches(tr_caps, cfg.batch_size, rng):
            imgs = torch.from_numpy(tr_imgs[idx]).to(device)
            ids, mask = encode_captions(tokenizer, [tr_caps[i] for i in idx], cfg.max_len, device)
            loss = clip_loss(model, imgs, ids, mask)
            optimizer.zero_grad()
            loss.backward()
            if cfg.grad_clip > 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
            optimizer.step()
            epoch_loss += loss.item()
            steps += 1

        # Évaluation retrieval au niveau des classes sur la val.
        with torch.no_grad():
            v_imgs = torch.from_numpy(va_imgs).to(device)
            v_labels = torch.tensor([bank_index[c] for c in va_caps], dtype=torch.long,
                                    device=device)
            acc = class_retrieval_accuracy(model, v_imgs, bank_ids, bank_mask, v_labels)
        rec = {"epoch": epoch, "train_loss": round(epoch_loss / max(1, steps), 4),
               "val_retrieval_acc": round(acc, 4)}
        history.append(rec)
        if (epoch + 1) % cfg.log_interval == 0 or epoch == cfg.epochs - 1:
            print(f"epoch {epoch:3d} | loss {rec['train_loss']:.4f} | "
                  f"retrieval acc {acc:.3f} | {time.time() - t0:.0f}s", flush=True)
        if acc > best_acc:
            best_acc = acc
            torch.save({"model": model.state_dict(), "clip_config": _cfg_dict(cfg),
                        "val_retrieval_acc": acc}, out_dir / "best.pt")

    summary = {"best_val_retrieval_acc": best_acc, "epochs": cfg.epochs,
               "elapsed_s": round(time.time() - t0, 1), "history": history}
    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2),
                                         encoding="utf-8")
    print("Résumé :", json.dumps({k: v for k, v in summary.items() if k != "history"}), flush=True)
    return summary


def _cfg_dict(cfg: ClipTrainConfig) -> dict:
    return asdict(cfg)


def parse_args() -> ClipTrainConfig:
    parser = argparse.ArgumentParser(description="Entraîner le CLIP maison (Phase 5).")
    for field, value in asdict(ClipTrainConfig()).items():
        parser.add_argument(f"--{field.replace('_', '-')}", type=type(value), default=value)
    ns = parser.parse_args()
    return ClipTrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()})


if __name__ == "__main__":
    train(parse_args())