"""Fine-tuned cross-encoder producing P(match) per pair.

Default model is BAAI/bge-reranker-v2-m3 (Apache-2.0, 568M, XLM-R large backbone).
It is ALREADY a trained reranker with a 1-logit head, so from_pretrained(num_labels=1)
warm-starts that head rather than randomising it -- verified: no "newly initialized"
warning. Zero-shot it already scores a Malayalam/English true pair at +10.75 and a
same-name/different-city false pair at -3.75.

CHECKPOINTING (added 2026-09-26)
  Training and inference are both resumable, because these stages run for hours on a
  shared box and a dropped connection or an OOM from another user must not cost the run.

  * training  -> <out_dir>/ckpt/  every ce_ckpt_steps optimiser steps: model weights,
                 optimiser, scheduler, epoch and step-within-epoch. On restart the
                 batch order is reproduced (per-epoch seeded generator) and already
                 consumed batches are skipped without a forward pass.
  * inference -> <cache>.part.npy + <cache>.done.json every ce_pred_ckpt_rows rows.
                 A restart reloads the finished prefix and scores only the remainder.

  Both are keyed on a fingerprint of the inputs, so a stale checkpoint from a
  different run can never be silently resumed into.
"""
import hashlib
import json
import math
import os
import random
import time

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup


def pair_text(name, addr, country):
    return f"name: {name} | addr: {addr} | country: {country}"


# --------------------------------------------------------------------------- utils
def _fp(*parts):
    h = hashlib.sha256()
    for p in parts:
        h.update(repr(p).encode("utf-8", "ignore"))
        h.update(b"\x00")
    return h.hexdigest()[:16]


def _texts_fp(a, b):
    """Cheap but collision-safe fingerprint of a pair list: length + sampled rows."""
    n = len(a)
    idx = list(range(0, n, max(1, n // 64)))[:64]
    return _fp(n, [a[i] for i in idx], [b[i] for i in idx])


def _atomic_json(path, obj):
    tmp = path + ".tmp"
    with open(tmp, "w") as f:
        json.dump(obj, f)
    os.replace(tmp, path)


class PairDS(Dataset):
    def __init__(self, a, b, y=None, swap_prob=0.0):
        self.a, self.b, self.y, self.swap = a, b, y, swap_prob

    def __len__(self):
        return len(self.a)

    def __getitem__(self, i):
        a, b = self.a[i], self.b[i]
        if self.swap and random.random() < self.swap:   # both orders A-B / B-A
            a, b = b, a
        return a, b, (self.y[i] if self.y is not None else 0.0)


def _collate(tok, max_len):
    def fn(batch):
        a, b, y = zip(*batch)
        enc = tok(list(a), list(b), truncation="longest_first", max_length=max_len,
                  padding=True, return_tensors="pt")
        enc["labels"] = torch.tensor(y, dtype=torch.float32)
        return enc
    return fn


def _amp_dtype():
    """bf16 on any card that supports it. mDeBERTa needed this because fp16 gave NaN
    losses; bge-reranker-v2-m3 is XLM-R and fp16-safe, but bf16 is free on Blackwell."""
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
        return torch.bfloat16
    return None


# ------------------------------------------------------------------------ training
def train_cross_encoder(a, b, y, cfg, out_dir, seed=42):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    ck_every = int(cfg.get("ce_ckpt_steps", 2000))
    ck_dir = os.path.join(out_dir, "ckpt")
    state_p = os.path.join(ck_dir, "state.json")
    fp = _fp(cfg["ce_model"], cfg["ce_lr"], cfg["ce_epochs"], cfg["ce_batch"],
             cfg["ce_max_len"], seed, _texts_fp(a, b), len(y))

    # already finished?
    if os.path.exists(os.path.join(out_dir, "config.json")) and not os.path.exists(state_p):
        print(f"[ce] {out_dir} already trained - skipping")
        return out_dir

    tok = AutoTokenizer.from_pretrained(cfg["ce_model"])
    resume = None
    if os.path.exists(state_p):
        try:
            st = json.load(open(state_p))
            if st.get("fp") == fp:
                resume = st
                print(f"[ce] resuming from checkpoint: epoch {st['epoch']} "
                      f"step {st['step']}/{st['total_steps']}")
            else:
                print("[ce] checkpoint fingerprint mismatch - starting fresh")
        except Exception as e:                                   # noqa: BLE001
            print(f"[ce] unreadable checkpoint ({e}) - starting fresh")

    src = ck_dir if resume else cfg["ce_model"]
    model = AutoModelForSequenceClassification.from_pretrained(src, num_labels=1).to(dev)

    ds = PairDS(a, b, y, swap_prob=0.5)
    no_decay = ["bias", "LayerNorm.weight", "layernorm", "norm"]
    params = [
        {"params": [p for n, p in model.named_parameters() if not any(k in n for k in no_decay)],
         "weight_decay": cfg["ce_weight_decay"]},
        {"params": [p for n, p in model.named_parameters() if any(k in n for k in no_decay)],
         "weight_decay": 0.0},
    ]
    opt = torch.optim.AdamW(params, lr=cfg["ce_lr"])

    def loader(epoch):
        g = torch.Generator()
        g.manual_seed(seed * 1000 + epoch)          # reproducible order for resume
        return DataLoader(ds, batch_size=cfg["ce_batch"], shuffle=True, generator=g,
                          num_workers=cfg["ce_num_workers"],
                          collate_fn=_collate(tok, cfg["ce_max_len"]), drop_last=True)

    steps_per_epoch = len(loader(0))
    total = steps_per_epoch * cfg["ce_epochs"]
    sch = get_linear_schedule_with_warmup(opt, int(cfg["ce_warmup"] * total), total)

    start_epoch, start_step, gstep = 0, 0, 0
    if resume:
        opt.load_state_dict(torch.load(os.path.join(ck_dir, "opt.pt"), map_location=dev))
        sch.load_state_dict(torch.load(os.path.join(ck_dir, "sch.pt"), map_location=dev))
        start_epoch, start_step, gstep = resume["epoch"], resume["step_in_epoch"], resume["step"]

    lossf = torch.nn.BCEWithLogitsLoss()
    amp = _amp_dtype()
    os.makedirs(ck_dir, exist_ok=True)

    # Both cards on this box are RTX PRO 6000 (97 GB). `model` stays the raw module
    # so save_pretrained/checkpointing are untouched; only the forward pass is
    # parallelised. DataParallel splits the batch, so ce_batch is the GLOBAL batch.
    ngpu = torch.cuda.device_count()
    use_dp = ngpu > 1 and cfg.get("ce_data_parallel", False)
    net = torch.nn.DataParallel(model) if use_dp else model
    if use_dp:
        print(f"[ce] training on {ngpu} GPUs (DataParallel), global batch {cfg['ce_batch']}")

    def save_ckpt(epoch, step_in_epoch, gstep):
        model.save_pretrained(ck_dir)
        tok.save_pretrained(ck_dir)
        torch.save(opt.state_dict(), os.path.join(ck_dir, "opt.pt"))
        torch.save(sch.state_dict(), os.path.join(ck_dir, "sch.pt"))
        _atomic_json(state_p, {"fp": fp, "epoch": epoch, "step_in_epoch": step_in_epoch,
                               "step": gstep, "total_steps": total,
                               "saved_at": time.strftime("%Y-%m-%d %H:%M:%S")})
        print(f"[ce] checkpoint @ epoch {epoch} step {gstep}/{total}", flush=True)

    model.train()
    t0 = time.time()
    for ep in range(start_epoch, cfg["ce_epochs"]):
        dl = loader(ep)
        skip = start_step if ep == start_epoch else 0
        if skip:
            print(f"[ce] skipping {skip:,} already-trained batches of epoch {ep}")
        run, seen = 0.0, 0
        for i, batch in enumerate(dl):
            if i < skip:                                  # replay order, no compute
                continue
            yb = batch.pop("labels").to(dev)
            batch = {k: v.to(dev) for k, v in batch.items()}
            with torch.autocast("cuda", dtype=amp, enabled=amp is not None):
                logits = net(**batch).logits.squeeze(-1)
            loss = lossf(logits.float(), yb)
            if not torch.isfinite(loss):
                raise RuntimeError("non-finite loss - check precision settings")
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); sch.step(); opt.zero_grad(set_to_none=True)
            run += loss.item(); gstep += 1; seen += 1
            if seen % 500 == 0:
                print(f"[ce] ep {ep} step {gstep}/{total} loss {run / 500:.4f} "
                      f"({time.time() - t0:.0f}s)", flush=True)
                run = 0.0
            if ck_every and gstep % ck_every == 0:
                save_ckpt(ep, i + 1, gstep)
        start_step = 0

    os.makedirs(out_dir, exist_ok=True)
    model.save_pretrained(out_dir)
    tok.save_pretrained(out_dir)
    if os.path.exists(state_p):
        os.remove(state_p)                                # mark training complete
    print(f"[ce] training done in {time.time() - t0:.0f}s -> {out_dir}")
    return out_dir


# ----------------------------------------------------------------------- inference
@torch.no_grad()
def predict_cross_encoder(a, b, model_dir, cfg, cache=None):
    """P(match) per pair. If `cache` is given, scoring is resumable: the finished
    prefix is reloaded and only the remainder is computed."""
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    n = len(a)
    fp = _texts_fp(a, b)
    done = 0
    out = np.empty(n, np.float32)

    part_p = done_p = None
    if cache:
        os.makedirs(os.path.dirname(cache) or ".", exist_ok=True)
        part_p, done_p = cache + ".part.npy", cache + ".done.json"
        if os.path.exists(done_p) and os.path.exists(part_p):
            try:
                st = json.load(open(done_p))
                if st.get("fp") == fp and st.get("n") == n:
                    prev = np.load(part_p)
                    done = int(st.get("done_sorted", st.get("done", 0)))
                    out[:len(prev)] = prev[:len(prev)]
                    print(f"[ce] resuming scoring: {done:,}/{n:,} already done")
                else:
                    print("[ce] prediction cache mismatch - rescoring")
            except Exception as e:                            # noqa: BLE001
                print(f"[ce] unreadable prediction cache ({e}) - rescoring")
    if done >= n:
        return out

    tok = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(dev).eval()
    ngpu = torch.cuda.device_count()
    use_dp = ngpu > 1 and cfg.get("ce_data_parallel", False)
    net = torch.nn.DataParallel(model) if use_dp else model
    if use_dp:
        print(f"[ce] scoring on {ngpu} GPUs (DataParallel)")
    amp = _amp_dtype()
    bs = cfg["ce_eval_batch"]
    ck_rows = int(cfg.get("ce_pred_ckpt_rows", 2_000_000))
    col = _collate(tok, cfg["ce_max_len"])

    # length-sorted within the remaining slice: big speedup, and keeps the resume
    # boundary a simple prefix of the ORIGINAL order.
    order = np.argsort([len(a[i]) + len(b[i]) for i in range(n)], kind="stable")
    buf = np.empty(n, np.float32)
    if done:
        buf[:] = out[:]                       # scored values for order[:done]
    start = done                              # resume at this position in `order`

    t0, scored, since = time.time(), 0, 0
    for s in range(start, len(order), bs):
        idx = order[s:s + bs]
        enc = col([(a[i], b[i], 0.0) for i in idx])
        enc.pop("labels")
        enc = {k: v.to(dev) for k, v in enc.items()}
        with torch.autocast("cuda", dtype=amp, enabled=amp is not None):
            lg = net(**enc).logits.squeeze(-1).float()
        buf[idx] = torch.sigmoid(lg).cpu().numpy()
        scored += len(idx); since += len(idx)
        if scored % (bs * 2000) == 0:
            rate = scored / max(time.time() - t0, 1e-9)
            print(f"[ce] scored {scored:,}/{len(order):,} ({rate:,.0f}/s, "
                  f"{time.time()-t0:.0f}s)", flush=True)
        # Progress is a prefix of the LENGTH-SORTED order, not the original one.
        # `order` is a deterministic stable argsort of the same inputs, so it is
        # recomputed identically on resume. (The previous version demanded a
        # contiguous prefix of the original order, which after 2M scored rows still
        # resolved to row ~35 - correct but useless for resuming.)
        if cache and since >= ck_rows:
            k = s + len(idx)
            np.save(part_p + ".tmp.npy", buf)
            os.replace(part_p + ".tmp.npy", part_p)
            _atomic_json(done_p, {"fp": fp, "n": n, "done_sorted": int(k)})
            print(f"[ce] scoring checkpoint: {k:,}/{n:,} rows persisted "
                  f"({100*k/max(n,1):.1f}%)", flush=True)
            since = 0

    if cache:
        np.save(part_p + ".tmp.npy", buf)
        os.replace(part_p + ".tmp.npy", part_p)
        _atomic_json(done_p, {"fp": fp, "n": n, "done": n})
    return buf