"""Phase 3 — boucle d'entraînement du transformer texte (from scratch).

Tout est écrit à la main : planificateur de LR warmup+cosine, découpage des lots,
boucle d'entraînement/validation, gradient clipping, autocast, checkpoints, logs.

Aucun `Trainer` de `transformers` : seulement PyTorch.

Usage :

    python -m src.train.train_text --data-dir data/processed/wikipedia_fr \\
        --out-dir checkpoints/text_50m --max-steps 20000 --batch-size 16 \\
        --grad-accum 4 --ctx-len 512 --device auto
"""

from __future__ import annotations

import argparse
import json
import math
import time
from contextlib import nullcontext
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import torch

from src.data.prepare_text import DTYPE
from src.model.text_model import CortexTextLM, TextConfig

ROOT = Path(__file__).resolve().parents[2]


# ---------------------------------------------------------------------------
# Hyperparamètres / planificateur
# ---------------------------------------------------------------------------


@dataclass
class TrainConfig:
    data_dir: str = "data/processed/wikipedia_fr"
    out_dir: str = "checkpoints/text"
    max_steps: int = 2000
    batch_size: int = 16
    grad_accum: int = 4
    ctx_len: int = 512
    lr: float = 3e-4
    min_lr_ratio: float = 0.1
    warmup_steps: int = 200
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0
    eval_interval: int = 100
    eval_batches: int = 20
    log_interval: int = 10
    seed: int = 1337
    device: str = "auto"
    dtype: str = "auto"  # auto | float32 | bfloat16 | float16
    n_layer: int = 6
    n_head: int = 8
    n_embd: int = 512
    dropout: float = 0.0
    tie_embeddings: bool = False
    compile: bool = False


def get_lr(step: int, cfg: TrainConfig) -> float:
    """LR warmup linéaire puis décroissance cosinus jusqu'à ``lr * min_lr_ratio``."""
    if step < cfg.warmup_steps:
        return cfg.lr * (step + 1) / max(1, cfg.warmup_steps)
    if step >= cfg.max_steps:
        return cfg.lr * cfg.min_lr_ratio
    ratio = (step - cfg.warmup_steps) / max(1, cfg.max_steps - cfg.warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * ratio))
    min_lr = cfg.lr * cfg.min_lr_ratio
    return min_lr + coeff * (cfg.lr - min_lr)


def pick_device(requested: str) -> torch.device:
    if requested != "auto":
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def pick_autocast(device: torch.device, dtype: str):
    """Renvoie (context manager, dtype) pour l'autocast, ou (nullcontext, ...)."""
    if dtype == "float32":
        return nullcontext(), torch.float32
    if dtype == "auto":
        if device.type == "cuda":
            dtype = "bfloat16" if torch.cuda.is_bf16_supported() else "float16"
        else:
            dtype = "float32"
    if dtype == "float32":
        return nullcontext(), torch.float32
    amp_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
    if device.type == "cuda":
        return torch.autocast("cuda", dtype=amp_dtype), amp_dtype
    # CPU : bfloat16 supporté par autocast sur les CPU récents, sinon pas d'autocast.
    if dtype == "bfloat16":
        try:
            return torch.autocast("cpu", dtype=torch.bfloat16), torch.bfloat16
        except Exception:
            return nullcontext(), torch.float32
    return nullcontext(), torch.float32


# ---------------------------------------------------------------------------
# Données
# ---------------------------------------------------------------------------


def load_split(data_dir: Path, split: str) -> np.memmap:
    path = Path(data_dir) / f"{split}.bin"
    if not path.is_file():
        raise FileNotFoundError(
            f"{path} introuvable. Lance d'abord src/data/prepare_text.py."
        )
    return np.memmap(path, dtype=DTYPE, mode="r")


def get_batch(
    data: np.memmap,
    batch_size: int,
    ctx_len: int,
    device: torch.device,
    generator: torch.Generator | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Échantillonne un lot (x, y) où y est x décalé d'un token."""
    if len(data) <= ctx_len + 1:
        raise ValueError(f"corpus trop petit ({len(data)} tokens) pour ctx_len={ctx_len}")
    hi = len(data) - ctx_len - 1
    idx = torch.randint(hi, (batch_size,), generator=generator)
    x = torch.stack([torch.from_numpy(data[i : i + ctx_len].astype(np.int64)) for i in idx.tolist()])
    y = torch.stack([torch.from_numpy(data[i + 1 : i + 1 + ctx_len].astype(np.int64)) for i in idx.tolist()])
    return x.to(device), y.to(device)


# ---------------------------------------------------------------------------
# Évaluation
# ---------------------------------------------------------------------------


@torch.no_grad()
def estimate_loss(
    model: CortexTextLM,
    splits: dict[str, np.memmap],
    cfg: TrainConfig,
    device: torch.device,
    autocast_ctx,
) -> dict[str, float]:
    model.eval()
    out = {}
    for name, data in splits.items():
        losses = torch.zeros(cfg.eval_batches)
        for i in range(cfg.eval_batches):
            x, y = get_batch(data, cfg.batch_size, cfg.ctx_len, device)
            with autocast_ctx:
                _, loss = model(x, targets=y)
            losses[i] = loss.item()
        out[name] = losses.mean().item()
    model.train()
    return out


# ---------------------------------------------------------------------------
# Entraînement
# ---------------------------------------------------------------------------


def train(cfg: TrainConfig) -> dict:
    torch.manual_seed(cfg.seed)
    device = pick_device(cfg.device)
    autocast_ctx, _ = pick_autocast(device, cfg.dtype)
    out_dir = Path(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    from src.tokenizer.special_tokens import VOCAB_SIZE

    text_cfg = TextConfig(
        vocab_size=VOCAB_SIZE,
        n_layer=cfg.n_layer,
        n_head=cfg.n_head,
        n_embd=cfg.n_embd,
        ctx_len=cfg.ctx_len,
        dropout=cfg.dropout,
        tie_embeddings=cfg.tie_embeddings,
    )
    model = CortexTextLM(text_cfg).to(device)
    if cfg.compile and hasattr(torch, "compile"):
        model = torch.compile(model)  # type: ignore[assignment]

    # AdamW avec weight decay uniquement sur les matrices (pas sur normes/biais).
    decay, no_decay = [], []
    for _name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        (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, betas=(cfg.beta1, cfg.beta2),
    )

    train_data = load_split(cfg.data_dir, "train")
    splits = {"train": train_data, "val": load_split(cfg.data_dir, "val")}

    gen = torch.Generator().manual_seed(cfg.seed)
    log_path = out_dir / "train_log.jsonl"
    best_val = float("inf")
    history: list[dict] = []
    model.train()
    t0 = time.time()

    for step in range(cfg.max_steps):
        lr = get_lr(step, cfg)
        for group in optimizer.param_groups:
            group["lr"] = lr

        optimizer.zero_grad(set_to_none=True)
        accum_loss = 0.0
        for _ in range(cfg.grad_accum):
            x, y = get_batch(train_data, cfg.batch_size, cfg.ctx_len, device, gen)
            with autocast_ctx:
                _, loss = model(x, targets=y)
                loss = loss / cfg.grad_accum
            loss.backward()
            accum_loss += loss.item()
        if cfg.grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
        optimizer.step()

        if step % cfg.log_interval == 0:
            elapsed = time.time() - t0
            record = {"step": step, "loss": round(accum_loss, 4), "lr": round(lr, 6),
                      "elapsed_s": round(elapsed, 1)}
            history.append(record)
            print(f"step {step:6d} | loss {accum_loss:.4f} | lr {lr:.2e} | {elapsed:.0f}s", flush=True)
            with log_path.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(record) + "\n")

        if (step + 1) % cfg.eval_interval == 0 or step == cfg.max_steps - 1:
            metrics = estimate_loss(model, splits, cfg, device, autocast_ctx)
            ppl = math.exp(min(metrics["val"], 20))
            print(f"  >> eval step {step}: train {metrics['train']:.4f} | "
                  f"val {metrics['val']:.4f} | ppl {ppl:.2f}", flush=True)
            eval_rec = {"step": step, **metrics, "val_ppl": round(ppl, 3)}
            history.append(eval_rec)
            with log_path.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps(eval_rec) + "\n")

            ckpt = {
                "step": step,
                "model": {k: v for k, v in model.state_dict().items()},
                "optimizer": optimizer.state_dict(),
                "text_config": asdict(text_cfg),
                "train_config": asdict(cfg),
                "val_loss": metrics["val"],
            }
            torch.save(ckpt, out_dir / "last.pt")
            if metrics["val"] < best_val:
                best_val = metrics["val"]
                torch.save(ckpt, out_dir / "best.pt")

    summary = {"best_val_loss": best_val, "best_val_ppl": math.exp(min(best_val, 20)),
               "steps": cfg.max_steps, "elapsed_s": round(time.time() - t0, 1),
               "n_params": model.count_parameters()}
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print("Résumé :", json.dumps(summary), flush=True)
    return summary


def parse_args() -> TrainConfig:
    parser = argparse.ArgumentParser(description="Entraîner le transformer texte Cortex-0.")
    defaults = TrainConfig()
    for field, value in asdict(defaults).items():
        arg = f"--{field.replace('_', '-')}"
        if isinstance(value, bool):
            parser.add_argument(arg, type=lambda s: s.lower() in ("1", "true", "yes"),
                                default=value)
        else:
            parser.add_argument(arg, type=type(value), default=value)
    ns = parser.parse_args()
    return TrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()})


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