"""Blend the TF-IDF baseline with a fine-tuned encoder and evaluate honestly.

The two strongest models on the leak-free ``document`` split — TF-IDF + one-vs-rest
logistic regression and ModernBERT — score within 0.01 macro-F1 of each other but
make different mistakes: one is a bag-of-ngrams linear model, the other a
contextual encoder. A convex blend of their per-class probabilities, with the
blend weight and decision thresholds tuned **on dev only**, tests whether that
disagreement carries signal.

    python scripts/05_ensemble.py --scheme document --model modernbert

Nothing here is tuned on test. The blend weight alpha is chosen on dev, the
thresholds are chosen on dev, and only then is the blend scored on test.
"""

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import numpy as np  # noqa: E402
from torch.utils.data import DataLoader  # noqa: E402

from cti_attack import baselines, config, data, evaluate, modeling  # noqa: E402


def bert_scores(model_key: str, scheme: str, dv, te, labels):
    """Dev + test sigmoid probabilities from the saved best checkpoint."""
    out_dir = config.MODELS_DIR / f"{model_key}__{scheme}"
    model, tok = modeling.load_for_inference(out_dir)
    dev_ = modeling.device()
    model.to(dev_)
    amp = model_key not in config.FP32_ONLY_MODELS
    ds_dv = modeling.SentenceDataset(dv, labels, tok, config.MAX_LENGTH)
    ds_te = modeling.SentenceDataset(te, labels, tok, config.MAX_LENGTH)
    s_dv = modeling.predict_scores(model, DataLoader(ds_dv, batch_size=32), dev_, amp=amp)
    s_te = modeling.predict_scores(model, DataLoader(ds_te, batch_size=32), dev_, amp=amp)
    return s_dv, s_te


def tfidf_scores(tr, dv, te, Ytr, labels):
    txt = lambda recs: [r["sentence"] for r in recs]
    s = baselines.tfidf_lr_scores(txt(tr), Ytr, {"dev": txt(dv), "test": txt(te)})
    return s["dev"], s["test"]


def dev_macro_at_best_threshold(Ydv, scores):
    """Best macro-F1 achievable on dev under a single global threshold."""
    _, f = evaluate.tune_global_threshold(Ydv, scores)
    return f


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--scheme", default="document", choices=["document", "random"])
    ap.add_argument("--model", default="modernbert", choices=list(config.BASE_MODELS))
    args = ap.parse_args()

    labels = data.load_labels()
    tr = data.load_split(args.scheme, "train")
    dv = data.load_split(args.scheme, "dev")
    te = data.load_split(args.scheme, "test")

    Ytr = evaluate.to_matrix(tr, labels)
    Ydv = evaluate.to_matrix(dv, labels).astype("int8")
    Yte = evaluate.to_matrix(te, labels).astype("int8")

    print(f"{'=' * 62}\n  ensemble  tfidf_lr + {args.model}  ({args.scheme})\n{'=' * 62}")
    print(f"  train={len(tr)}  dev={len(dv)}  test={len(te)}  labels={len(labels)}")

    print("  fitting tfidf_lr …")
    t_dv, t_te = tfidf_scores(tr, dv, te, Ytr, labels)
    print(f"  loading {args.model} checkpoint and scoring …")
    b_dv, b_te = bert_scores(args.model, args.scheme, dv, te, labels)

    # ---- choose the blend weight alpha on dev ---------------------------------
    # scores = alpha * tfidf + (1 - alpha) * bert
    alphas = [round(0.05 * i, 2) for i in range(21)]
    dev_curve = []
    for a in alphas:
        blend = a * t_dv + (1 - a) * b_dv
        dev_curve.append((a, dev_macro_at_best_threshold(Ydv, blend)))
    best_alpha, best_dev = max(dev_curve, key=lambda x: x[1])
    print("\n  alpha sweep (dev macro-F1, global threshold):")
    for a, f in dev_curve:
        star = "  <-- best" if a == best_alpha else ""
        print(f"    alpha={a:.2f}  dev={f:.4f}{star}")

    blend_dv = best_alpha * t_dv + (1 - best_alpha) * b_dv
    blend_te = best_alpha * t_te + (1 - best_alpha) * b_te

    # ---- tune thresholds on dev, apply to test --------------------------------
    gt, _ = evaluate.tune_global_threshold(Ydv, blend_dv)
    pct = evaluate.tune_per_class_thresholds(Ydv, blend_dv)
    rep_g = evaluate.evaluate(Yte, evaluate.apply_thresholds(blend_te, gt), labels)
    rep_p = evaluate.evaluate(Yte, evaluate.apply_thresholds(blend_te, pct), labels)

    # ---- reference: the two components alone, same protocol -------------------
    def solo(name, s_dv, s_te):
        g, _ = evaluate.tune_global_threshold(Ydv, s_dv)
        p = evaluate.tune_per_class_thresholds(Ydv, s_dv)
        rg = evaluate.evaluate(Yte, evaluate.apply_thresholds(s_te, g), labels)
        rp = evaluate.evaluate(Yte, evaluate.apply_thresholds(s_te, p), labels)
        print(f"  {name:22} global={rg.macro_f1:.4f}  per-class={rp.macro_f1:.4f}")
        return rg.macro_f1, rp.macro_f1

    print("\n  --- test macro-F1 (all tuned on dev) ---")
    tf_g, tf_p = solo("tfidf_lr", t_dv, t_te)
    bt_g, bt_p = solo(args.model, b_dv, b_te)
    print(f"  {'ensemble (a=%.2f)' % best_alpha:22} global={rep_g.macro_f1:.4f}  per-class={rep_p.macro_f1:.4f}")

    best_solo = max(tf_g, tf_p, bt_g, bt_p)
    best_ens = max(rep_g.macro_f1, rep_p.macro_f1)
    print(f"\n  best component={best_solo:.4f}   best ensemble={best_ens:.4f}   "
          f"delta={best_ens - best_solo:+.4f}")

    payload = {
        "run": f"ensemble_tfidf_{args.model}",
        "split_scheme": args.scheme,
        "alpha": best_alpha,
        "alpha_meaning": "scores = alpha*tfidf + (1-alpha)*bert",
        "dev_macro_f1_at_best_alpha": round(best_dev, 4),
        "global_threshold": {"threshold": gt, **rep_g.as_dict()},
        "per_class_threshold": {
            "thresholds": {l: float(t) for l, t in zip(labels, pct)},
            **rep_p.as_dict(),
        },
        "components_test_macro_f1": {
            "tfidf_lr": {"global": round(tf_g, 4), "per_class": round(tf_p, 4)},
            args.model: {"global": round(bt_g, 4), "per_class": round(bt_p, 4)},
        },
    }
    out = config.RESULTS_DIR / f"ensemble_tfidf_{args.model}__{args.scheme}.json"
    out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(f"  -> {out.relative_to(config.REPO_ROOT)}")


if __name__ == "__main__":
    main()