"""Fine-tune an encoder as a multi-label toxicity classifier with the 7 Detoxify labels (sigmoid, BCE on the fractional
annotator labels, as Detoxify does).   python tox/train_tox.py --base jhu-clsp/mmBERT-small --data $DATA --out $OUTPUT_DIR/model
Data: tox/build_tox.py. Checkpoint selection on val (held-out source comments and their translations): mean ROC AUC over
the labels at >= 0.5 (labels with >= 20 val positives).
"""
import argparse, json, os, random, time
import numpy as np, pandas as pd, torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_cosine_schedule_with_warmup

ap = argparse.ArgumentParser()
ap.add_argument("--base", default="jhu-clsp/mmBERT-small")
ap.add_argument("--data", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--max_len", type=int, default=256)
ap.add_argument("--pos_weight", type=float, default=1.0, help="BCE weight of positive labels")
ap.add_argument("--tokens_per_batch", type=int, default=32768)
ap.add_argument("--lr", type=float, default=5e-5)
ap.add_argument("--epochs", type=float, default=3)
ap.add_argument("--warmup", type=float, default=0.06)
ap.add_argument("--wd", type=float, default=0.01)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--eval_every", type=int, default=1000)
ap.add_argument("--labels", default="toxicity,severe_toxicity,obscene,threat,insult,identity_attack,sexual_explicit",
                help="label columns (hate/: hate,target_race,...)")
ap.add_argument("--regression", action="store_true", help="one raw regression output trained with MSE (edu/: FineWeb-Edu style 0-5 score); "
                "val metric = Spearman correlation")
ap.add_argument("--context", action="store_true", help="encode (context, text) pairs when the data has a non-empty context column (irony/)")
args = ap.parse_args()
random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed)
os.makedirs(args.out, exist_ok=True)

import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
LABELS = args.labels.split(",")
tok = AutoTokenizer.from_pretrained(args.base)
model = AutoModelForSequenceClassification.from_pretrained(args.base, num_labels=len(LABELS), id2label=dict(enumerate(LABELS)),
                                                           label2id={l: i for i, l in enumerate(LABELS)},
                                                           problem_type="regression" if args.regression else "multi_label_classification").cuda()
model.config.uses_context = bool(args.context)
tr = pd.read_parquet(f"{args.data}/train.parquet"); va = pd.read_parquet(f"{args.data}/val.parquet")
print("train", len(tr), "val", len(va), flush=True)


def encode(df):
    if args.context and "context" in df:  # (context, text) pairs where a context is given, else the text alone
        ctx = df.context.fillna("").tolist(); txt = df.text.tolist(); enc = [None] * len(df)
        single = [i for i, c in enumerate(ctx) if not c]; pair = [i for i, c in enumerate(ctx) if c]
        for i, e in zip(single, tok([txt[i] for i in single], truncation=True, max_length=args.max_len)["input_ids"] if single else []):
            enc[i] = e
        for i, e in zip(pair, tok([ctx[i] for i in pair], [txt[i] for i in pair], truncation="longest_first", max_length=args.max_len)["input_ids"] if pair else []):
            enc[i] = e
    else:
        enc = tok(df.text.tolist(), truncation=True, max_length=args.max_len)["input_ids"]
    return list(zip(enc, df[LABELS].values.astype(np.float32)))


tr_enc, va_enc = encode(tr), encode(va)


def batches(data, shuffle):
    idx = list(range(len(data)))
    if shuffle:
        random.shuffle(idx)
    out = []
    for s in range(0, len(idx), 4096):
        part = sorted(idx[s: s + 4096], key=lambda i: len(data[i][0]))
        cur, mx = [], 0
        for i in part:
            L = len(data[i][0])
            if cur and max(mx, L) * (len(cur) + 1) > args.tokens_per_batch:
                out.append(cur); cur, mx = [], 0
            cur.append(i); mx = max(mx, L)
        if cur:
            out.append(cur)
    if shuffle:
        random.shuffle(out)
    return out


def collate(data, b):
    L = max(len(data[i][0]) for i in b)
    ids = torch.full((len(b), L), tok.pad_token_id, dtype=torch.long); att = torch.zeros((len(b), L), dtype=torch.long)
    for j, i in enumerate(b):
        x = data[i][0]; ids[j, :len(x)] = torch.tensor(x); att[j, :len(x)] = 1
    return ids.cuda(), att.cuda(), torch.tensor(np.stack([data[i][1] for i in b])).cuda()


@torch.no_grad()
def evaluate():
    model.eval(); P = np.zeros((len(va_enc), len(LABELS)))
    for b in batches(va_enc, False):
        ids, att, _ = collate(va_enc, b)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            lo = model(input_ids=ids, attention_mask=att).logits.float()
            P[b] = (lo if args.regression else torch.sigmoid(lo)).cpu().numpy()
    model.train()
    T = np.stack([x[1] for x in va_enc]); aucs = []
    if args.regression:  # Spearman (Pearson of ranks) per output; key kept as macro_f1 for the selection code below
        rk = lambda a: pd.Series(a).rank().values
        sp = [float(np.corrcoef(rk(P[:, c]), rk(T[:, c]))[0, 1]) for c in range(len(LABELS))]
        return dict(macro_f1=float(np.mean(sp)), mse=float(((P - T) ** 2).mean()))
    for c in range(len(LABELS)):
        y = T[:, c] >= 0.5
        if y.sum() >= 20:
            s = P[:, c]; r = np.empty(len(s)); r[s.argsort()] = np.arange(1, len(s) + 1)
            aucs.append((r[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (len(y) - y.sum())))
    return dict(macro_f1=float(np.mean(aucs)), bce=float(-(T * np.log(np.clip(P, 1e-9, 1)) + (1 - T) * np.log(np.clip(1 - P, 1e-9, 1))).mean()))


steps = len(batches(tr_enc, True)); total = int(steps * args.epochs)
nd = ["bias", "norm", "LayerNorm"]
opt = torch.optim.AdamW([{"params": [p for n, p in model.named_parameters() if not any(k in n for k in nd)], "weight_decay": args.wd},
                         {"params": [p for n, p in model.named_parameters() if any(k in n for k in nd)], "weight_decay": 0.0}],
                        lr=args.lr, betas=(0.9, 0.98), eps=1e-6, fused=True)
sch = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total)
print(f"steps/epoch {steps}, total {total}", flush=True)
step, best, log, t0 = 0, -1, [], time.time()
model.train()
while step < total:
    for b in batches(tr_enc, True):
        ids, att, t = collate(tr_enc, b)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits = model(input_ids=ids, attention_mask=att).logits.float()
        loss = (torch.nn.functional.mse_loss(logits, t) if args.regression else
                torch.nn.functional.binary_cross_entropy_with_logits(logits, t, pos_weight=torch.full((len(LABELS),), args.pos_weight, device=t.device)))
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step(); sch.step(); opt.zero_grad(set_to_none=True); step += 1
        if step % 100 == 0:
            print(f"step {step}/{total} loss {loss.item():.4f} {time.time()-t0:.0f}s", flush=True)
        if step % args.eval_every == 0 or step == total:
            res = evaluate(); log.append(dict(step=step, **res)); print("EVAL", step, json.dumps(res), flush=True)
            if res["macro_f1"] > best:
                best = res["macro_f1"]; model.save_pretrained(args.out); tok.save_pretrained(args.out)
                json.dump(dict(step=step, **res), open(f"{args.out}/val_metrics.json", "w"), indent=1)
        if step >= total:
            break
json.dump(dict(args=vars(args), log=log), open(f"{args.out}/train_log.json", "w"), indent=1)
print("done; best val metric (mean AUC, or Spearman with --regression)", best)