"""Train BagOfWords on AG News (12k headlines, 4 topics); about 1 s per epoch on a laptop CPU.

    python train.py --epochs 8

Lines marked `# forge` are the patch the forge setup adds; the rest is a plain training script.
"""

from __future__ import annotations

import argparse
import json

import torch
import torch.nn.functional as F

import forge  # forge
from data import CLASSES, HERE, get_loaders
from model import BagOfWords

CHECKPOINT = HERE / "checkpoints" / "best.pt"
BEST = CHECKPOINT.with_suffix(".json")  # val acc of best.pt, so a weaker run never replaces it


def saved_best_acc() -> float:
    try:
        return float(json.loads(BEST.read_text())["val_acc"])
    except (OSError, ValueError, KeyError, TypeError):
        return 0.0


def run_epoch(model, loader, opt) -> tuple[float, float]:
    training = opt is not None
    model.train(training)
    total_loss, correct, n = 0.0, 0, 0
    with torch.set_grad_enabled(training):
        for x, y in loader:
            out = model(x)
            loss = F.cross_entropy(out, y)
            if training:
                opt.zero_grad()
                loss.backward()
                opt.step()
            total_loss += loss.item() * len(y)
            correct += (out.argmax(1) == y).sum().item()
            n += len(y)
    return total_loss / n, correct / n


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=8)
    parser.add_argument("--data", default=str(HERE / "data"))
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=3e-3)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    train_loader, val_loader, vocab = get_loaders(args.data, args.batch_size)
    model = BagOfWords(len(vocab), len(CLASSES))  # CPU: the model is small enough that a GPU wouldn't help
    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
    forge.watch(model, val_loader.dataset, classes=CLASSES, epochs=args.epochs)  # forge

    best_acc = run_best = 0.0
    kept_acc = saved_best_acc()
    for epoch in range(1, args.epochs + 1):
        train_loss, train_acc = run_epoch(model, train_loader, opt)
        val_loss, val_acc = run_epoch(model, val_loader, None)
        print(f"epoch {epoch}/{args.epochs}  train {train_loss:.3f} / {train_acc:.3f}  val {val_loss:.3f} / {val_acc:.3f}")
        forge.log(epoch=epoch, train_loss=train_loss, train_acc=train_acc, val_loss=val_loss,  # forge
                  val_acc=val_acc, lr=opt.param_groups[0]["lr"], model=model)  # forge
        run_best = max(run_best, val_acc)
        if val_acc > max(best_acc, kept_acc):
            best_acc = val_acc
            CHECKPOINT.parent.mkdir(exist_ok=True)
            torch.save(model.state_dict(), CHECKPOINT)
            BEST.write_text(json.dumps({"val_acc": val_acc, "epoch": epoch, "epochs": args.epochs}))

    forge.evaluate(model, val_loader)  # forge
    if best_acc:
        print(f"best val acc {best_acc:.3f}, saved to {CHECKPOINT}")
    else:
        print(f"best val acc {run_best:.3f}; kept {CHECKPOINT} ({kept_acc:.3f} from an earlier run)")


if __name__ == "__main__":
    main()