"""Blend encoder + LLM (+ALLaM) scores once, then sweep none_bias to see what
label distribution comes out.
"""
import argparse
import collections
import os

import numpy as np
import torch

from src.data import ID2LABEL, LABEL2ID, load_split
from src.predict import model_probs
from src.scorer import validate_submission

LAB = ["Against", "Favor", "None"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", required=True)
    ap.add_argument("--models", nargs="*", default=[])
    ap.add_argument("--enc_weight", type=float, default=0.0)
    ap.add_argument("--llm_probs", nargs="+", required=True,
                    help="one or more npy files (LLM, ALLaM, ...)")
    ap.add_argument("--llm_weights", nargs="+", type=float, default=None,
                    help="weights for each --llm_probs (default: equal)")
    ap.add_argument("--sweep", nargs="+", type=float,
                    default=[0.0, -0.1, -0.2, -0.3, -0.5, -1.0])
    ap.add_argument("--none_bias", type=float, default=None,
                    help="if set, write a submission at this bias")
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    n = len(load_split(args.csv, "preserve", has_labels=False))

    # LLM side: weighted average of the provided npy prob files
    llms = [np.load(f) for f in args.llm_probs]
    for f, a in zip(args.llm_probs, llms):
        if len(a) != n:
            raise SystemExit(f"{f}: {len(a)} rows != csv {n}")
    w = args.llm_weights or [1.0] * len(llms)
    w = np.array(w) / sum(w)
    llm = sum(wi * a for wi, a in zip(w, llms))

    if args.models:
        defaults = {"prep_mode": "preserve", "use_description": False,
                    "max_len": 128}
        enc = np.mean([model_probs(m, args.csv, device, defaults)
                       for m in args.models], axis=0)
        base = args.enc_weight * enc + (1 - args.enc_weight) * llm
    else:
        base = llm

    print(f"rows={n}  llm_files={args.llm_probs} weights={list(w.round(3))} "
          f"enc_weight={args.enc_weight if args.models else 0}")
    for nb in args.sweep:
        p = base.copy()
        p[:, LABEL2ID["None"]] += nb
        preds = [ID2LABEL[i] for i in p.argmax(1)]
        c = collections.Counter(preds)
        none_pct = 100 * c.get("None", 0) / n
        print(f"  none_bias {nb:+.2f} -> Against={c.get('Against',0):3d} "
              f"Favor={c.get('Favor',0):3d} None={c.get('None',0):3d} "
              f"({none_pct:.0f}%)")

    if args.none_bias is not None and args.out:
        p = base.copy()
        p[:, LABEL2ID["None"]] += args.none_bias
        preds = [ID2LABEL[i] for i in p.argmax(1)]
        ok, msg = validate_submission(preds, n)
        print(f"[validate] {msg}")
        if not ok:
            raise SystemExit(1)
        os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
        with open(args.out, "w", encoding="utf-8") as f:
            f.write("\n".join(preds) + "\n")
        print(f"[write] none_bias {args.none_bias:+.2f} -> {args.out}")


if __name__ == "__main__":
    main()