"""
Fine tune DistilBERT on the short-message dataset.

  source .venv/bin/activate
  python scripts/train_v2.py

Same architecture as v1 so the data is the only变 variable.
"""
import csv, os, random, time
import torch
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL = "distilbert-base-uncased"
OUTDIR = "models/distilbert_v2"
MAXLEN = 128
BATCH = 32
EPOCHS = 3
LR = 2e-5
SEED = 42
LABELS = {"legitimate": 0, "scam": 1}

random.seed(SEED); torch.manual_seed(SEED)
device = "mps" if torch.backends.mps.is_available() else "cpu"
print("device:", device)

def load(split):
    rows = list(csv.DictReader(open("data/processed/%s.csv" % split, newline="", encoding="utf-8")))
    return [r["text"] for r in rows], [LABELS[r["label"]] for r in rows]

tok = AutoTokenizer.from_pretrained(MODEL)

def encode(texts, labels):
    enc = tok(texts, truncation=True, padding="max_length", max_length=MAXLEN, return_tensors="pt")
    return TensorDataset(enc["input_ids"], enc["attention_mask"], torch.tensor(labels))

tr_x, tr_y = load("train")
va_x, va_y = load("val")
te_x, te_y = load("test")
print("train %d  val %d  test %d" % (len(tr_x), len(va_x), len(te_x)))

train_dl = DataLoader(encode(tr_x, tr_y), batch_size=BATCH, shuffle=True)
val_dl   = DataLoader(encode(va_x, va_y), batch_size=64)
test_dl  = DataLoader(encode(te_x, te_y), batch_size=64)

n_legit = tr_y.count(0); n_scam = tr_y.count(1); total = len(tr_y)
weights = torch.tensor([total/(2.0*n_legit), total/(2.0*n_scam)], dtype=torch.float).to(device)
print("class weights: legitimate %.3f scam %.3f" % (weights[0], weights[1]))

model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=LR)
lossf = torch.nn.CrossEntropyLoss(weight=weights)

def evaluate(dl):
    model.eval()
    tp = fp = fn = tn = 0
    with torch.no_grad():
        for ids, mask, y in dl:
            ids, mask, y = ids.to(device), mask.to(device), y.to(device)
            pred = model(input_ids=ids, attention_mask=mask).logits.argmax(1)
            tp += int(((pred == 1) & (y == 1)).sum())
            fp += int(((pred == 1) & (y == 0)).sum())
            fn += int(((pred == 0) & (y == 1)).sum())
            tn += int(((pred == 0) & (y == 0)).sum())
    prec = tp / max(1, tp + fp)
    rec = tp / max(1, tp + fn)
    f1 = 2 * prec * rec / max(1e-9, prec + rec)
    acc = (tp + tn) / max(1, tp + tn + fp + fn)
    return acc, prec, rec, f1, (tp, fp, fn, tn)

best = -1
for ep in range(1, EPOCHS + 1):
    model.train()
    t0, run = time.time(), 0.0
    for i, (ids, mask, y) in enumerate(train_dl, 1):
        ids, mask, y = ids.to(device), mask.to(device), y.to(device)
        opt.zero_grad()
        loss = lossf(model(input_ids=ids, attention_mask=mask).logits, y)
        loss.backward(); opt.step()
        run += loss.item()
        if i % 50 == 0:
            print("  epoch %d  batch %d/%d  loss %.4f" % (ep, i, len(train_dl), run / i))
    acc, prec, rec, f1, _ = evaluate(val_dl)
    print("epoch %d done in %.0fs  val acc %.4f  prec %.4f  rec %.4f  F1 %.4f" % (
        ep, time.time() - t0, acc, prec, rec, f1))
    if f1 > best:
        best = f1
        os.makedirs(OUTDIR, exist_ok=True)
        model.save_pretrained(OUTDIR); tok.save_pretrained(OUTDIR)
        print("  saved (best so far)")

print("")
model = AutoModelForSequenceClassification.from_pretrained(OUTDIR).to(device)
acc, prec, rec, f1, (tp, fp, fn, tn) = evaluate(test_dl)
print("=" * 60)
print("HELD OUT TEST")
print("=" * 60)
print("accuracy %.4f  precision %.4f  recall %.4f  F1 %.4f" % (acc, prec, rec, f1))
print("true scam %d   false alarm %d   missed scam %d   true legit %d" % (tp, fp, fn, tn))
print("")
print("This number is NOT the Nigerian number. Run the Nigerian evaluation")
print("and the length experiment before believing anything.")