"""Run one or more models over a CSV and dump a submission file. Multiple
model dirs just get their softmax probabilities averaged. Checks the
submission format before writing it out.

    python -m src.predict --models outputs/t1_arabert \\
        --csv data/track1/dev.csv --gold data/track1/dev.csv \\
        --out submissions/t1_dev.txt
"""
import argparse
import json
import os

import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from src.data import ID2LABEL, LABEL2ID, StanceDataset, load_split
from src.scorer import load_gold, score, validate_submission


def model_settings(model_dir, defaults):
    """Read the preprocessing a model was trained with, if recorded."""
    cfg_path = os.path.join(model_dir, "best.json")
    if os.path.isfile(cfg_path):
        with open(cfg_path, encoding="utf-8") as f:
            cfg = json.load(f).get("config", {})
        return {
            "prep_mode": cfg.get("prep_mode", defaults["prep_mode"]),
            "use_description": cfg.get(
                "use_description", defaults["use_description"]
            ),
            "max_len": cfg.get("max_len", defaults["max_len"]),
        }
    return dict(defaults)


@torch.no_grad()
def model_probs(model_dir, csv_path, device, defaults, batch_size=64):
    s = model_settings(model_dir, defaults)
    df = load_split(csv_path, s["prep_mode"], has_labels=False)
    tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_dir, trust_remote_code=True
    ).to(device).eval()
    ds = StanceDataset(
        df, tok, s["max_len"], s["use_description"], has_labels=False
    )
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
    probs = []
    for batch in loader:
        batch = {k: v.to(device) for k, v in batch.items()}
        logits = model(**batch).logits.float()
        probs.append(F.softmax(logits, dim=-1).cpu().numpy())
    return np.concatenate(probs, axis=0)


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", nargs="+", required=True)
    ap.add_argument("--csv", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--gold", default=None)
    ap.add_argument("--max_len", type=int, default=128)
    ap.add_argument("--use_description", action="store_true")
    ap.add_argument("--prep_mode", default="preserve")
    ap.add_argument("--none_bias", type=float, default=0.0)
    ap.add_argument("--llm_probs", default=None,
                    help="npy of LLM class probabilities aligned to --csv")
    ap.add_argument("--enc_weight", type=float, default=1.0,
                    help="weight on encoder probs; LLM gets 1 - enc_weight")
    return ap.parse_args()


def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    defaults = {
        "prep_mode": args.prep_mode,
        "use_description": args.use_description,
        "max_len": args.max_len,
    }
    n_rows = len(load_split(args.csv, "preserve", has_labels=False))

    probs = np.mean(
        [model_probs(m, args.csv, device, defaults) for m in args.models],
        axis=0,
    )
    if args.llm_probs:
        llm = np.load(args.llm_probs)
        if len(llm) != n_rows:
            raise SystemExit(
                f"llm_probs rows {len(llm)} != csv rows {n_rows}"
            )
        probs = args.enc_weight * probs + (1 - args.enc_weight) * llm
    probs[:, LABEL2ID["None"]] += args.none_bias
    preds = [ID2LABEL[i] for i in probs.argmax(axis=1)]

    ok, msg = validate_submission(preds, n_rows)
    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] {len(preds)} predictions -> {args.out}")

    if args.gold:
        print("[score]")
        score(load_gold(args.gold), preds)


if __name__ == "__main__":
    main()