"""Reproduces our final Codabench submission (Gemma-4-E4B-it + one LoRA adapter, bf16, deterministic).

Everything is downloaded from Hugging Face inside this script (base model, adapter, optional test data).
No local paths, no retrieval, no internet other than Hugging Face, no closed API.

Usage (fresh Colab / Kaggle, 1 GPU with >= 24 GB):
    pip install -r requirements.txt
    python inference.py --input <folder or .jsonl with the questions> --output predictions.jsonl
    # or, if the organisers host the questions on Hugging Face: set DATA_REPO below (or pass --data_repo)

perm4 = deterministic option rotation (same model + same adapter, 4 scoring passes, scores summed; no extra parameters,
no sampling, fixed seed). Use --perm 1 for the single-pass variant.
"""
import argparse
import subprocess
import sys
import tempfile
from pathlib import Path

from huggingface_hub import snapshot_download

# ---- team settings -------------------------------------------------------------------------------------------
BASE_MODEL = "google/gemma-4-e4b-it"                                    # downloaded by transformers from the Hub
MODEL_ID = "SinhalaMMLU-SharedTask/team-ape-llama-gemma4-e4b-lora"      # our adapter repo (weights + adapter_config.json)
DATA_REPO = None                                                         # organisers: HF dataset repo with the questions
# --------------------------------------------------------------------------------------------------------------

HERE = Path(__file__).resolve().parent


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", help="folder (easy/medium/hard/*.json) or .jsonl with the questions")
    ap.add_argument("--data_repo", default=DATA_REPO, help="HF dataset repo with the questions (used when --input is not given)")
    ap.add_argument("--output", default="predictions.jsonl")
    ap.add_argument("--perm", type=int, default=4, choices=[1, 4], help="4 = option rotation (submitted); 1 = single pass")
    args = ap.parse_args()

    if args.input:
        questions = args.input
    elif args.data_repo:
        questions = snapshot_download(args.data_repo, repo_type="dataset")
    else:
        sys.exit("give --input or set DATA_REPO")

    adapter_dir = snapshot_download(MODEL_ID)  # adapter_config.json + adapter_model.safetensors
    cfg = (HERE / "configs" / "final_bf16_perm4.yaml").read_text(encoding="utf-8")
    cfg = cfg.replace("__ADAPTER_DIR__", adapter_dir).replace("permutations: 4", f"permutations: {args.perm}")
    with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False, encoding="utf-8") as f:
        f.write(cfg)
        cfg_path = f.name

    cmd = [sys.executable, str(HERE / "predict.py"), "--config", cfg_path, "--input", questions, "--output", args.output]
    print("running:", " ".join(cmd))
    return subprocess.call(cmd, cwd=str(HERE))


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