"""Train one of the five models on the graphics card and write its test scores (2026-10-29).

    python models/train.py <name>      e.g. business-pii-detector
Reads models/<name>/data/{train,test}.jsonl (models/build.py), fine-tunes distilroberta-base, saves the model to
models/<name>/out/ (not committed; uploaded to Hugging Face) and the scores to models/<name>/scores.json.
Span models (rows with "entities") are scored by exact entity match; classifiers by accuracy and macro F1.
"""

import json
import os
import sys
from collections import Counter
from pathlib import Path

import torch
from torch.utils.data import DataLoader
from transformers import (AutoModelForSequenceClassification, AutoModelForTokenClassification, AutoTokenizer,
                          get_linear_schedule_with_warmup)

BASE = "distilroberta-base"
ROOT = Path(__file__).resolve().parent
name = sys.argv[1]
folder = ROOT / name
read = lambda s: [json.loads(l) for l in open(folder / "data" / f"{s}.jsonl", encoding="utf-8")]
train_rows, test_rows = read("train"), read("test")
spans = "entities" in train_rows[0]
dev = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(1029)
tok = AutoTokenizer.from_pretrained(BASE, add_prefix_space=True)

if spans:
    types = sorted({e["label"] for r in train_rows + test_rows for e in r["entities"]})
    labels = ["O"] + [f"{p}-{t}" for t in types for p in ("B", "I")]
else:
    labels = sorted({r["label"] for r in train_rows + test_rows})
lid = {l: i for i, l in enumerate(labels)}
MAX = 256 if not spans else 160


def encode_spans(r):
    enc = tok(r["text"], truncation=True, max_length=MAX, return_offsets_mapping=True)
    tags = []
    for s, e in enc["offset_mapping"]:
        if s == e:
            tags.append(-100)
            continue
        tag = "O"
        for ent in r["entities"]:
            if s >= ent["start"] and e <= ent["end"] + 1 and s < ent["end"]:
                tag = ("B-" if s <= ent["start"] else "I-") + ent["label"]
                break
        tags.append(lid[tag])
    enc.pop("offset_mapping")
    enc["labels"] = tags
    return enc


def encode_cls(r):
    enc = tok(r["text"], truncation=True, max_length=MAX)
    enc["labels"] = lid[r["label"]]
    return enc


def collate(batch):
    n = max(len(b["input_ids"]) for b in batch)
    out = {"input_ids": [], "attention_mask": [], "labels": []}
    for b in batch:
        pad = n - len(b["input_ids"])
        out["input_ids"].append(b["input_ids"] + [tok.pad_token_id] * pad)
        out["attention_mask"].append(b["attention_mask"] + [0] * pad)
        out["labels"].append(b["labels"] + [-100] * pad if spans else b["labels"])
    return {k: torch.tensor(v) for k, v in out.items()}


encode = encode_spans if spans else encode_cls
train = [encode(r) for r in train_rows]
Model = AutoModelForTokenClassification if spans else AutoModelForSequenceClassification
model = Model.from_pretrained(BASE, num_labels=len(labels), id2label=dict(enumerate(labels)), label2id=lid).to(dev)
EPOCHS = int(os.environ.get("EPOCHS", 8 if spans else 6))  # EPOCHS=20 for many-class models with few rows each
loader = DataLoader(train, batch_size=16, shuffle=True, collate_fn=collate)
opt = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
sched = get_linear_schedule_with_warmup(opt, len(loader) * EPOCHS // 10, len(loader) * EPOCHS)
model.train()
for ep in range(EPOCHS):
    total = 0.0
    for batch in loader:
        with torch.autocast(dev, dtype=torch.float16, enabled=dev == "cuda"):
            loss = model(**{k: v.to(dev) for k, v in batch.items()}).loss
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step(), sched.step(), opt.zero_grad()
        total += loss.item()
    print(f"epoch {ep + 1}: loss {total / len(loader):.4f}")
model.eval()


def predict_spans(text):
    enc = tok(text, truncation=True, max_length=MAX, return_offsets_mapping=True, return_tensors="pt")
    offs = enc.pop("offset_mapping")[0].tolist()
    with torch.no_grad():
        ids = model(**{k: v.to(dev) for k, v in enc.items()}).logits[0].argmax(-1).tolist()
    found, cur = [], None
    for (s, e), i in zip(offs, ids):
        if s == e:
            continue
        tag = labels[i]
        if tag == "O":
            cur = None
            continue
        kind = tag[2:]
        if tag.startswith("I-") and cur and cur["label"] == kind:
            cur["end"] = e
        else:
            cur = {"label": kind, "start": s, "end": e}
            found.append(cur)
    for f in found:  # offsets from a prefix-space tokenizer can start on the space
        while f["start"] < f["end"] and text[f["start"]] == " ":
            f["start"] += 1
    return found


if spans:
    tp, fp, fn = Counter(), Counter(), Counter()
    for r in test_rows:
        gold = {(e["label"], e["start"], e["end"]) for e in r["entities"]}
        pred = {(e["label"], e["start"], e["end"]) for e in predict_spans(r["text"])}
        for g in gold:
            (tp if g in pred else fn)[g[0]] += 1
        for p in pred - gold:
            fp[p[0]] += 1
    f1 = lambda t, p, n: (2 * t / (2 * t + p + n)) if t else 0.0
    per = {t: {"precision": round(tp[t] / (tp[t] + fp[t]), 3) if tp[t] + fp[t] else 0.0,
               "recall": round(tp[t] / (tp[t] + fn[t]), 3) if tp[t] + fn[t] else 0.0,
               "f1": round(f1(tp[t], fp[t], fn[t]), 3), "support": tp[t] + fn[t]} for t in types}
    T, P, N = sum(tp.values()), sum(fp.values()), sum(fn.values())
    scores = {"metric": "exact entity match", "micro_precision": round(T / (T + P), 3), "micro_recall": round(T / (T + N), 3),
              "micro_f1": round(f1(T, P, N), 3), "per_type": per}
else:
    pred, gold = [], []
    for r in test_rows:
        enc = tok(r["text"], truncation=True, max_length=MAX, return_tensors="pt")
        with torch.no_grad():
            pred.append(labels[model(**{k: v.to(dev) for k, v in enc.items()}).logits[0].argmax().item()])
        gold.append(r["label"])
    acc = sum(p == g for p, g in zip(pred, gold)) / len(gold)
    per = {}
    for l in labels:
        t = sum(p == g == l for p, g in zip(pred, gold))
        p_ = sum(p == l for p in pred)
        g_ = sum(g == l for g in gold)
        per[l] = {"f1": round(2 * t / (p_ + g_), 3) if p_ + g_ else 0.0, "support": g_}
    scores = {"metric": "accuracy and macro F1", "accuracy": round(acc, 3),
              "macro_f1": round(sum(v["f1"] for v in per.values() if v["support"]) / sum(1 for v in per.values() if v["support"]), 3),
              "per_label": per}

scores.update({"base_model": BASE, "train_rows": len(train_rows), "test_rows": len(test_rows), "epochs": EPOCHS, "device": dev})
(folder / "scores.json").write_text(json.dumps(scores, indent=2) + "\n", encoding="utf-8")
print(json.dumps({k: v for k, v in scores.items() if k not in ("per_type", "per_label")}))
model.save_pretrained(folder / "out", safe_serialization=True)
tok.save_pretrained(folder / "out")