#!/usr/bin/env python3
"""Train the loan grade classifier in either LoRA or full fine-tuning mode.

    python train.py --mode lora --rank 8 --epochs 4
    python train.py --mode full --epochs 4

Written as an explicit PyTorch loop rather than ``transformers.Trainer``, because
the mechanics are the thing worth showing.

Writes ``results/{mode}_metrics.json`` and ``checkpoints/{mode}.pt``.

COLAB
-----
Runtime -> Change runtime type -> T4 GPU, then::

    !git clone <your-repo> && cd RiscAutious
    !pip install -q -r requirements.txt
    !python data/download.py
    !python train.py --mode lora
    !python train.py --mode full

It falls back to CPU automatically, but full fine-tuning on CPU is slow enough
that you will notice.
"""

from __future__ import annotations

import argparse
import json
import logging
import random
import sys
import time
from pathlib import Path
from typing import Sequence

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from data.dataset import build_dataloaders, load_dataframe, load_labels
from models.classifier import TextClassifier

log = logging.getLogger("train")

#: Default learning rates, per mode. These differ **on purpose** and using one
#: value for both would make the comparison meaningless:
#:
#:   LoRA (1e-3) — the adapters start at exactly zero and have to travel a long
#:     way. At 2e-5 they barely move and LoRA looks far worse than it is.
#:   Full (2e-5) — the pretrained weights are already close to useful. At 1e-3
#:     the first few steps overwrite what pretraining learned ("catastrophic
#:     forgetting") and accuracy collapses.
#:
#: Each mode gets the learning rate that is standard practice for it. That is
#: the fair comparison, not an identical number.
DEFAULT_LR: dict[str, float] = {"lora": 1e-3, "full": 2e-5}


def set_seed(seed: int) -> None:
    """Seed every RNG that affects training, so runs are reproducible.

    Three separate generators matter here: Python's ``random`` (used by the
    sampler), NumPy's, and PyTorch's (weight init, dropout masks). Seeding only
    ``torch`` is a common half-measure that leaves runs non-reproducible.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def pick_device(requested: str | None = None) -> torch.device:
    """Choose the compute device: explicit request, else CUDA > MPS > CPU."""
    if requested:
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():  # Apple Silicon
        return torch.device("mps")
    return torch.device("cpu")


def synchronize(device: torch.device) -> None:
    """Block until queued GPU work finishes.

    GPU kernels launch asynchronously: ``time.time()`` right after a forward pass
    records when the work was *queued*, not when it completed. Without this, GPU
    timings come out absurdly fast and the training-time comparison is fiction.
    """
    if device.type == "cuda":
        torch.cuda.synchronize()
    elif device.type == "mps":
        torch.mps.synchronize()


@torch.no_grad()
def evaluate_split(
    model: nn.Module, loader: DataLoader, device: torch.device
) -> tuple[float, float, list[int], list[int]]:
    """Run the model over a split without training on it.

    ``model.eval()`` matters: it switches dropout off and makes the forward pass
    deterministic. Forgetting it means your validation accuracy is measured on a
    randomly-perturbed model and jitters between runs.

    Returns:
        ``(mean_loss, accuracy, predictions, true_labels)``.
    """
    model.eval()
    total_loss, correct, seen = 0.0, 0, 0
    predictions: list[int] = []
    truths: list[int] = []

    for batch in loader:
        batch = {k: v.to(device) for k, v in batch.items()}
        out = model(batch["input_ids"], batch["attention_mask"], batch["labels"])

        # Weight by batch size: the last batch is usually smaller, so a plain
        # mean over batches would over-weight it.
        total_loss += out["loss"].item() * len(batch["labels"])
        preds = out["logits"].argmax(dim=-1)
        correct += (preds == batch["labels"]).sum().item()
        seen += len(batch["labels"])

        predictions.extend(preds.cpu().tolist())
        truths.extend(batch["labels"].cpu().tolist())

    return total_loss / seen, correct / seen, predictions, truths


def train_one_epoch(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    device: torch.device,
    max_grad_norm: float = 1.0,
    log_every: int = 50,
) -> tuple[float, float]:
    """One pass over the training data. Returns ``(mean_loss, accuracy)``."""
    model.train()  # enables dropout
    total_loss, correct, seen = 0.0, 0, 0

    for step, batch in enumerate(loader):
        batch = {k: v.to(device) for k, v in batch.items()}

        # --- the four lines that are the whole of gradient descent ---
        out = model(batch["input_ids"], batch["attention_mask"], batch["labels"])
        loss = out["loss"]
        loss.backward()  # accumulates d(loss)/d(param) into every param.grad

        # Rescale gradients if their combined norm exceeds the threshold. Cheap
        # insurance against one bad batch producing a huge step that wrecks the
        # weights. Standard practice for transformer fine-tuning.
        torch.nn.utils.clip_grad_norm_(
            [p for p in model.parameters() if p.requires_grad], max_grad_norm
        )

        optimizer.step()   # apply the update
        scheduler.step()   # advance the learning rate schedule
        optimizer.zero_grad(set_to_none=True)
        # zero_grad is NOT optional: PyTorch *accumulates* into .grad rather than
        # overwriting, so skipping it silently sums every batch's gradients.
        # set_to_none=True frees the tensors instead of filling them with zeros.
        # -------------------------------------------------------------

        total_loss += loss.item() * len(batch["labels"])
        correct += (out["logits"].argmax(dim=-1) == batch["labels"]).sum().item()
        seen += len(batch["labels"])

        if log_every and step % log_every == 0:
            log.info(
                "    step %4d/%d  loss %.4f  lr %.2e",
                step, len(loader), loss.item(), scheduler.get_last_lr()[0],
            )

    return total_loss / seen, correct / seen


def build_scheduler(
    optimizer: torch.optim.Optimizer, total_steps: int, warmup_ratio: float = 0.1
) -> torch.optim.lr_scheduler.LRScheduler:
    """Linear warmup then linear decay to zero.

    Warmup: the first steps use a tiny learning rate while Adam's running moment
    estimates are still based on almost no data and are therefore unreliable.
    Taking full-size steps on bad estimates destabilizes early training.

    Decay: large steps early to explore, small steps late to settle.
    """
    warmup_steps = max(1, int(total_steps * warmup_ratio))

    def lr_lambda(step: int) -> float:
        if step < warmup_steps:
            return step / warmup_steps
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return max(0.0, 1.0 - progress)

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)


def run(args: argparse.Namespace) -> dict:
    """Train one model end to end and return its metrics dict."""
    set_seed(args.seed)
    device = pick_device(args.device)
    log.info("Device: %s | mode: %s", device, args.mode)

    train_loader, val_loader, test_loader, _, labels = build_dataloaders(
        data_path=args.data,
        batch_size=args.batch_size,
        max_length=args.max_length,
        seed=args.seed,
    )

    log.info("%d classes", len(labels))

    weights = None
    if args.class_weights:
        from data.dataset import class_weights as compute_weights
        weights = compute_weights(load_dataframe(args.data, labels), labels).to(device)
        log.info("Using class-weighted loss")

    model = TextClassifier(
        model_name=args.model,
        num_labels=len(labels),
        label_names=labels,
        mode=args.mode,
        lora_r=args.rank,
        lora_alpha=args.alpha,
        lora_dropout=args.lora_dropout,
        class_weights=weights,
    ).to(device)

    report = model.trainable_parameter_report()
    log.info(
        "Trainable: %s / %s (%.3f%%)",
        f"{report['trainable_params']:,}",
        f"{report['total_params']:,}",
        report["trainable_pct"],
    )

    # Only hand the optimizer parameters that actually need updating. Passing
    # frozen ones would allocate Adam moment buffers for all 66M of them and
    # throw away most of LoRA's memory advantage.
    trainable = [p for p in model.parameters() if p.requires_grad]
    lr = args.lr if args.lr is not None else DEFAULT_LR[args.mode]
    optimizer = torch.optim.AdamW(trainable, lr=lr, weight_decay=args.weight_decay)
    scheduler = build_scheduler(optimizer, total_steps=len(train_loader) * args.epochs)
    log.info("Optimizer: AdamW lr=%.2e over %d trainable tensors", lr, len(trainable))

    history: list[dict] = []
    best_val_acc = -1.0
    checkpoint_path = Path(args.checkpoint_dir) / f"{args.mode}.pt"

    synchronize(device)
    start = time.perf_counter()

    for epoch in range(1, args.epochs + 1):
        log.info("Epoch %d/%d", epoch, args.epochs)
        train_loss, train_acc = train_one_epoch(
            model, train_loader, optimizer, scheduler, device, args.max_grad_norm
        )
        val_loss, val_acc, _, _ = evaluate_split(model, val_loader, device)

        history.append({
            "epoch": epoch,
            "train_loss": train_loss, "train_acc": train_acc,
            "val_loss": val_loss, "val_acc": val_acc,
        })
        log.info(
            "  train loss %.4f acc %.4f | val loss %.4f acc %.4f",
            train_loss, train_acc, val_loss, val_acc,
        )

        # Keep the epoch that generalized best, not the last one. Later epochs
        # usually have lower *training* loss while overfitting.
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            model.save(checkpoint_path)
            log.info("  new best val acc %.4f -> saved", val_acc)

    synchronize(device)
    train_seconds = time.perf_counter() - start
    log.info("Training finished in %.1fs", train_seconds)

    metrics = {
        "mode": args.mode,
        "history": history,
        "best_val_acc": best_val_acc,
        "final_val_acc": history[-1]["val_acc"],
        "train_seconds": train_seconds,
        "seconds_per_epoch": train_seconds / args.epochs,
        "device": str(device),
        "learning_rate": lr,
        "checkpoint": str(checkpoint_path),
        "checkpoint_kb": checkpoint_path.stat().st_size / 1024,
        "n_train": len(train_loader.dataset),
        "n_val": len(val_loader.dataset),
        "n_test": len(test_loader.dataset),
        "num_labels": len(labels),
        **report,
        "args": {k: str(v) for k, v in vars(args).items()},
    }

    out_path = Path(args.results_dir) / f"{args.mode}_metrics.json"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(metrics, indent=2))
    log.info("Wrote %s", out_path)

    print(f"\n  {args.mode.upper()}  best val acc {best_val_acc:.4f}  "
          f"| {report['trainable_params']:,} trainable ({report['trainable_pct']:.3f}%)  "
          f"| {train_seconds:.1f}s  | checkpoint {metrics['checkpoint_kb']:.0f} KB\n")
    return metrics


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    """Define and parse the command-line interface."""
    p = argparse.ArgumentParser(
        description="Fine-tune DistilBERT for loan grade prediction.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--mode", choices=("lora", "full"), default="lora")
    p.add_argument("--rank", type=int, default=8, help="LoRA rank r.")
    p.add_argument("--alpha", type=int, default=16, help="LoRA scaling numerator.")
    p.add_argument("--lora-dropout", type=float, default=0.05)
    p.add_argument("--epochs", type=int, default=4)
    p.add_argument("--batch-size", type=int, default=32, help="Drop to 16 on CUDA OOM.")
    p.add_argument("--lr", type=float, default=None,
                   help="Overrides the per-mode default (lora 1e-3, full 2e-5).")
    p.add_argument("--weight-decay", type=float, default=0.01)
    p.add_argument("--max-grad-norm", type=float, default=1.0)
    p.add_argument("--max-length", type=int, default=128)
    p.add_argument("--class-weights", action="store_true",
                   help="Weight the loss by inverse class frequency.")
    p.add_argument("--seed", type=int, default=42,
                   help="Same seed across modes = same split = fair comparison.")
    p.add_argument("--data", type=Path, default=Path("data/processed/dataset.csv"))
    p.add_argument("--model", default="distilbert-base-uncased")
    p.add_argument("--results-dir", type=Path, default=Path("results"))
    p.add_argument("--checkpoint-dir", type=Path, default=Path("checkpoints"))
    p.add_argument("--device", default=None, help="cuda / mps / cpu. Auto-detected.")
    return p.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    """Entry point. Returns a process exit code."""
    logging.basicConfig(level=logging.INFO, format="%(levelname)-7s %(message)s")
    for noisy in ("httpx", "urllib3", "filelock", "huggingface_hub"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    import transformers
    transformers.logging.set_verbosity_error()  # hides the unused-MLM-head report

    args = parse_args(argv)
    try:
        run(args)
    except FileNotFoundError as exc:
        log.error("%s", exc)
        log.error("Run: python data/download.py --source synthetic --rows 8000")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())