"""mDeBERTa-v3 cross-encoder: fine-tuned on the ce split (disjoint from stack/val, so its scores are
leak-free features downstream). Resumable training; sharded multi-GPU inference."""
import math
import os
import subprocess
import sys
import time

import numpy as np
import pandas as pd
import torch

from common import Progress, get_logger, gpu_patient, load_json, save_json, save_npy, save_parquet, shard_path
from config import CFG, PATHS, WORK
from prep import load_prep, load_splits
from retrieve_dense import n_shards

log = get_logger("ce")
CE = WORK / "ce"
os.environ.setdefault("HF_HOME", str(PATHS["hf_cache"]))


def _texts(split, side):
    d = load_prep(split, side, ["raw_name", "raw_addr"])
    return (d.raw_name + " | " + d.raw_addr).values


def _amp_dtype():
    """bf16 on Ampere+ (A100/L4/H100, same as the server run); fp16 + loss scaling on T4/V100."""
    return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16


def _tok_model():
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    return AutoTokenizer, AutoModelForSequenceClassification


def build_pairs():
    path = CE / "pairs.parquet"
    if path.exists():
        return pd.read_parquet(path)
    sp = load_splits()
    K = load_json(WORK / "reports" / "k_choice.json")["K"]
    parts = []
    for i in range(n_shards("train")):
        d = pd.read_parquet(shard_path("block", "train", i), columns=["q", "c", "y", "rrf_rank"])
        parts.append(d[np.isin(d.q.values, sp["g_ce"]) & (d.rrf_rank < K)])
    d = pd.concat(parts, ignore_index=True)
    rng = np.random.default_rng(CFG["seed"])
    pos = d[d.y == 1]
    neg = d[d.y == 0].sort_values(["q", "rrf_rank"])
    hard = neg.groupby("q").head(CFG["ce_neg_hard"])
    rest = neg.drop(hard.index)
    rnd = rest.assign(r=rng.random(len(rest))).sort_values(["q", "r"]).groupby("q").head(CFG["ce_neg_rand"])
    pairs = pd.concat([pos, hard, rnd[pos.columns]], ignore_index=True)[["q", "c", "y"]]
    pairs = pairs.sample(frac=1.0, random_state=CFG["seed"]).reset_index(drop=True)
    if len(pairs) > CFG["ce_max_pairs"]:          # shared GPU budget: a uniform sample keeps the pos/neg mix
        pairs = pairs.head(CFG["ce_max_pairs"])
    pairs["dev"] = ((pairs.q.values.astype(np.int64) * 2654435761 % 2 ** 32) % 30 == 0)
    save_parquet(pairs, path)
    log.info(f"ce pairs: {len(pairs)} (pos {int(pairs.y.sum())}), dev {int(pairs.dev.sum())}")
    return pairs


def _lr_lambda(total, warm):
    return lambda s: min(1.0, (s + 1) / max(1, warm)) * max(0.0, (total - s) / max(1, total - warm))


def run_ce_train():
    final = CE / "final"
    if (final / "config.json").exists():
        return
    AutoTokenizer, AutoModel = _tok_model()
    pairs = build_pairs()
    tr, dv = pairs[~pairs.dev].reset_index(drop=True), pairs[pairs.dev].reset_index(drop=True)
    ta, tb = _texts("train", "s1"), _texts("train", "pool")
    tok = AutoTokenizer.from_pretrained(CFG["ce_model"])
    # transformers>=5 loads the stored dtype (fp16 for mDeBERTa) by default; fp16 weights + AdamW -> NaN after 1 step
    model = AutoModel.from_pretrained(CFG["ce_model"], num_labels=1, dtype=torch.float32)
    model = gpu_patient(lambda: model.cuda(), log, "load cross-encoder")
    B = CFG["ce_batch"]
    spe = math.ceil(len(tr) / B)
    total = spe * CFG["ce_epochs"]
    if CFG["ce_max_steps"]:
        total = min(total, CFG["ce_max_steps"])
    opt = torch.optim.AdamW(model.parameters(), lr=CFG["ce_lr"], weight_decay=0.01)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, _lr_lambda(total, int(0.05 * total)))
    ck = CE / "ckpt.pt"
    amp = _amp_dtype()
    scaler = torch.amp.GradScaler("cuda", enabled=amp == torch.float16)
    log.info(f"cross-encoder mixed precision: {amp}")
    step = 0
    if ck.exists():
        st = torch.load(ck, map_location="cuda", weights_only=False)
        model.load_state_dict(st["model"]); opt.load_state_dict(st["opt"]); sched.load_state_dict(st["sched"])
        if st.get("scaler") and scaler.is_enabled():
            scaler.load_state_dict(st["scaler"])
        step = st["step"]
        log.info(f"resumed cross-encoder from step {step}/{total}")
    lossf = torch.nn.BCEWithLogitsLoss()
    model.train()
    t0, run = time.time(), []
    perms = {}
    while step < total:
        e, b = divmod(step, spe)
        if e not in perms:
            perms = {e: np.random.default_rng(CFG["seed"] + e).permutation(len(tr))}
        idx = perms[e][b * B:(b + 1) * B]
        q, c, y = tr.q.values[idx], tr.c.values[idx], tr.y.values[idx]
        swap = ((q.astype(np.int64) + c + e) % 2) == 1
        a, bb = ta[q], tb[c]
        x1, x2 = np.where(swap, bb, a).tolist(), np.where(swap, a, bb).tolist()
        try:
            enc = tok(x1, x2, truncation=True, max_length=CFG["ce_max_len"], padding=True, return_tensors="pt").to("cuda")
            with torch.autocast("cuda", dtype=amp):
                logits = model(**enc).logits.squeeze(-1)
            loss = lossf(logits.float(), torch.tensor(y, dtype=torch.float32, device="cuda"))
            scaler.scale(loss).backward()
        except torch.cuda.OutOfMemoryError:     # shared GPU filled up: drop this attempt, wait, redo the step
            opt.zero_grad(set_to_none=True)
            enc = logits = loss = None
            torch.cuda.empty_cache()
            log.warning(f"OOM at step {step}; waiting 60s for the shared GPU")
            time.sleep(60)
            continue
        scaler.unscale_(opt)
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if not torch.isfinite(gn) and not scaler.is_enabled():
            raise RuntimeError(f"non-finite gradient at step {step}; refusing to save a corrupted checkpoint")
        scaler.step(opt)                # with fp16 an overflowing step is skipped and the scale lowered
        scaler.update(); sched.step(); opt.zero_grad(set_to_none=True)
        if not all(torch.isfinite(p_).all() for p_ in (model.classifier.weight,)):
            raise RuntimeError(f"weights became non-finite at step {step}")
        step += 1
        run.append(loss.item())
        if step % 100 == 0:
            log.info(f"ce step {step}/{total} loss {np.mean(run[-100:]):.4f} {(time.time() - t0) / 60:.1f}m")
        if step % CFG["ce_ckpt_every"] == 0 or step == total:
            tmp = ck.with_suffix(".tmp")
            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "sched": sched.state_dict(),
                        "scaler": scaler.state_dict(), "step": step}, tmp)
            os.replace(tmp, ck)
    model.eval()
    p = _score(model, tok, ta[dv.q.values], tb[dv.c.values])
    from sklearn.metrics import roc_auc_score
    auc = roc_auc_score(dv.y.values, p) if dv.y.nunique() > 1 else float("nan")
    tmpd = CE / "final.tmp"
    model.save_pretrained(tmpd); tok.save_pretrained(tmpd)
    os.replace(tmpd, final)
    save_json({"dev_auc": auc, "steps": total, "n_train": len(tr)}, WORK / "reports" / "ce_train.json")
    log.info(f"cross-encoder dev AUC {auc:.5f}")


@torch.no_grad()
def _score(model, tok, a, b, bs=None):
    bs = bs or CFG["ce_infer_batch"]
    order = np.argsort([len(x) + len(y) for x, y in zip(a, b)])
    out = np.empty(len(a), dtype=np.float32)
    i = 0
    while i < len(order):
        idx = order[i:i + bs]
        try:
            enc = tok([a[j] for j in idx], [b[j] for j in idx], truncation=True, max_length=CFG["ce_max_len"],
                      padding=True, return_tensors="pt").to("cuda")
            with torch.autocast("cuda", dtype=_amp_dtype()):
                out[idx] = model(**enc).logits.squeeze(-1).float().cpu().numpy()
            i += bs
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            bs = max(8, bs // 2)
            log.warning(f"OOM -> infer batch {bs}")
    return out


def _n_ce_shards(split):
    n = len(pd.read_parquet(WORK / "s2" / f"{split}.parquet", columns=["q"]))
    return n, max(1, math.ceil(n / CFG["ce_shard_pairs"]))


def run_ce_infer_worker(split, rank, world):
    AutoTokenizer, AutoModel = _tok_model()
    n, ns = _n_ce_shards(split)
    todo = [i for i in range(ns) if i % world == rank and not shard_path("ce_scores", split, i, "npy").exists()]
    if not todo:
        return
    s2 = pd.read_parquet(WORK / "s2" / f"{split}.parquet", columns=["q", "c"])
    ta, tb = _texts(split, "s1"), _texts(split, "pool")
    tok = AutoTokenizer.from_pretrained(CE / "final")
    model = AutoModel.from_pretrained(CE / "final", dtype=torch.float32)
    model = gpu_patient(lambda: model.cuda().eval(), log, "load cross-encoder")
    prog = Progress(log, len(todo), f"ce infer {split} rank {rank}")
    for i in todo:
        if shard_path("ce_scores", split, i, "npy").exists():   # finished by another worker/helper meanwhile
            prog.step()
            continue
        sl = slice(i * CFG["ce_shard_pairs"], (i + 1) * CFG["ce_shard_pairs"])
        q, c = s2.q.values[sl], s2.c.values[sl]
        save_npy(shard_path("ce_scores", split, i, "npy"), _score(model, tok, ta[q], tb[c]))
        prog.step()


def run_ce_infer(split):
    gpus = [g for g in os.environ.get("M4_ASSIGNED_GPUS", "0").split(",") if g != ""]
    if len(gpus) <= 1:
        run_ce_infer_worker(split, 0, 1)
    else:
        procs = []
        for r, g in enumerate(gpus):
            env = dict(os.environ, CUDA_VISIBLE_DEVICES=g)
            procs.append(subprocess.Popen([sys.executable, "-u", __file__.replace("cross_encoder.py", "stages.py"),
                                           "ce_infer_worker", split, str(r), str(len(gpus))], env=env))
        codes = [p.wait() for p in procs]
        if any(codes):
            raise RuntimeError(f"ce inference workers failed: {codes}")
    n, ns = _n_ce_shards(split)
    missing = [i for i in range(ns) if not shard_path("ce_scores", split, i, "npy").exists()]
    if missing:
        raise RuntimeError(f"missing ce shards {missing[:10]}")


def load_ce_scores(split):
    s2 = pd.read_parquet(WORK / "s2" / f"{split}.parquet", columns=["q", "c"])
    n, ns = _n_ce_shards(split)
    ce = np.concatenate([np.load(shard_path("ce_scores", split, i, "npy")) for i in range(ns)])
    assert len(ce) == len(s2)
    return s2.assign(ce=ce)