"""Answer every question in a file and write predictions.jsonl.

Examples
--------
# baseline on the dev held-out split, with scores logged to results/results.csv
python predict.py --config configs/qwen25_7b.yaml --input data/dev_heldout.jsonl --evaluate --who Chamani

# test predictions + the zip Codabench expects
python predict.py --config configs/final.yaml --input data/test.jsonl --zip

# quick check that everything runs (first 5 questions only)
python predict.py --config configs/smoke_test.yaml --input data/sample/sample_dev.jsonl --limit 5
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

from sinmmlu.config import load_config
from sinmmlu.data import has_gold, load_questions, prediction_row, write_jsonl
from sinmmlu.metrics import accuracy_by_level, format_report
from sinmmlu.modeling import count_parameters, describe_budget, load_model, load_tokenizer, set_determinism
from sinmmlu.runinfo import append_result, write_run_info
from sinmmlu.scoring import Scorer


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--config", required=True)
    ap.add_argument("--input", required=True, help="jsonl of questions (dev, held-out or test)")
    ap.add_argument("--output", help="default: outputs/<config name>/<input name>.predictions.jsonl")
    ap.add_argument("--limit", type=int, help="only the first N questions (debugging)")
    ap.add_argument("--evaluate", action="store_true", help="input has gold answers: print + log scores")
    ap.add_argument("--zip", action="store_true", help="also write predictions.zip for Codabench")
    ap.add_argument("--who", default="", help="your name, for the results log")
    ap.add_argument("--note", default="", help="free text for the results log")
    args = ap.parse_args(argv)

    cfg = load_config(args.config)
    set_determinism(cfg["seed"])

    records = load_questions(args.input, cfg["data"]["answer_base"])
    if args.limit:
        records = records[: args.limit]
    print(f"[data] {len(records)} questions from {args.input}")

    pool = []
    if cfg["prompt"]["fewshot_k"] > 0:
        if not cfg["prompt"]["fewshot_file"]:
            sys.exit("prompt.fewshot_k > 0 needs prompt.fewshot_file (a dev TRAIN split)")
        if Path(cfg["prompt"]["fewshot_file"]).resolve() == Path(args.input).resolve():
            sys.exit("few-shot file must not be the file being predicted")
        pool = load_questions(cfg["prompt"]["fewshot_file"], cfg["data"]["answer_base"])
        if not has_gold(pool):
            sys.exit("few-shot file has no gold answers")

    t0 = time.time()
    tok = load_tokenizer(cfg["model"])
    model = load_model(cfg["model"])
    n_params = count_parameters(model)
    print(f"[model] {cfg['model']['name_or_path']} : {describe_budget(n_params)}")
    if cfg["model"]["load_in_4bit"]:
        print("[model] 4-bit: count above is understated; use tools/count_params.py for the manifest")

    scorer = Scorer(model, tok, cfg, pool)
    preds = scorer.predict(records)
    elapsed = time.time() - t0

    base = int(cfg["output"]["answer_base"])
    rows = [prediction_row(r, str(p + base)) for r, p in zip(records, preds)]
    out = Path(args.output) if args.output else Path("outputs") / cfg["name"] / f"{Path(args.input).stem}.predictions.jsonl"
    write_jsonl(rows, out)
    print(f"[out] {len(rows)} predictions -> {out}  ({elapsed:.0f}s)")

    extra = {"input": args.input, "n": len(records), "seconds": round(elapsed, 1), "loaded_params": n_params}
    if args.evaluate:
        if not has_gold(records):
            sys.exit("--evaluate: input file has no gold answers")
        result = accuracy_by_level([r["level"] for r in records], preds, [r["gold"] for r in records])
        print(format_report(result))
        if not args.limit:
            append_result(cfg, args.input, result, args.who, args.note)
            print("[log] appended to results/results.csv")
        extra["result"] = result
    write_run_info(out.with_suffix(".run.json"), cfg, extra)

    if args.zip:
        from tools.make_submission import make_zip

        make_zip(out, out.parent / "predictions.zip", expected=records)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())