"""
Sinhala MMLU Shared Task - Final inference script
Team: SinhalaNuwana

    pip install -r requirements.txt
    python inference.py

Downloads the questions from DATA_REPO, the base model and our LoRA adapter from Hugging Face,
runs predict.py six times and merge_votes.py once (both next to this file), and writes
predictions.jsonl and predictions.zip in the current folder. Greedy decoding, no sampling.
About 16 hours for the 2,238 test questions on 2 x T4.
"""

# ======================= ORGANIZERS CHANGE ONLY THIS LINE =======================
DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
# ================================================================================

import argparse
import hashlib
import json
import os
import subprocess
import sys
from pathlib import Path

from huggingface_hub import snapshot_download

MODEL = "google/gemma-4-E4B-it"
REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
ADAPTER = "SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora"
ADAPTER_SHA256 = "b82696e2541e0133d24bf14721caa20485dd70ab8ef78a7e40c1d2a77db1ba71"
ROOT = Path(__file__).resolve().parent

PASSES = [
    ("orders-02", ["--think", "512", "--rotations", "0", "2"]),
    ("orders-13", ["--think", "512", "--rotations", "1", "3"]),
    ("hard-2048", ["--level", "hard", "--think", "2048", "--score_batch_size", "8"]),
    ("easy-2048", ["--level", "easy", "--think", "2048", "--score_batch_size", "8"]),
    ("medium-rev02", ["--level", "medium", "--reverse", "--think", "512", "--rotations", "0", "2"]),
    ("medium-rev13", ["--level", "medium", "--reverse", "--think", "512", "--rotations", "1", "3"]),
]


def slice_questions(data_dir, out_dir, per_file):
    for path in sorted(data_dir.rglob("*.json")):
        if any(part.startswith(".") for part in path.relative_to(data_dir).parts):
            continue
        target = out_dir / path.relative_to(data_dir)
        target.parent.mkdir(parents=True, exist_ok=True)
        questions = json.loads(path.read_text(encoding="utf-8"))[:per_file]
        target.write_text(json.dumps(questions, ensure_ascii=False), encoding="utf-8")
    return out_dir


def main():
    parser = argparse.ArgumentParser(
        description="Answer the SinhalaMMLU questions with Gemma 4 E4B-it and the round 3 STaR LoRA adapter, "
        "and write OUT_DIR/predictions.jsonl and OUT_DIR/predictions.zip"
    )
    parser.add_argument("--data_dir", type=Path, help=f"local folder of the per-subject .json files; {DATA_REPO} when not given")
    parser.add_argument("--out_dir", type=Path, default=Path("."), help="where predictions.jsonl goes; each pass writes to OUT_DIR/runs/")
    parser.add_argument("--adapter", help=f"local adapter folder; {ADAPTER} from the Hub when not given")
    parser.add_argument("--per_file", type=int, help="answer only the first N questions of each file, for a quick check")
    args = parser.parse_args()

    snapshot_download(MODEL, revision=REVISION)
    adapter = Path(args.adapter or snapshot_download(ADAPTER))
    digest = hashlib.sha256((adapter / "adapter_model.safetensors").read_bytes()).hexdigest()
    if digest != ADAPTER_SHA256:
        raise SystemExit(f"error: {adapter}/adapter_model.safetensors is {digest}, expected {ADAPTER_SHA256}")

    data_dir = args.data_dir or Path(snapshot_download(DATA_REPO, repo_type="dataset"))
    runs = args.out_dir / "runs"
    if args.per_file:
        data_dir = slice_questions(data_dir, runs / "questions", args.per_file)

    env = {**os.environ, "HF_HUB_OFFLINE": "1"}
    long_env = {**env, "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
    common = [
        "--model", MODEL, "--revision", REVISION, "--dtype", "float16", "--batch_size", "32",
        "--adapter", str(adapter), "--data_dir", str(data_dir),
    ]
    for name, options in PASSES:
        print(f"== pass {name}", flush=True)
        subprocess.run(
            [sys.executable, str(ROOT / "predict.py"), *common, "--out_dir", str(runs / name), *options],
            check=True,
            env=long_env if "2048" in options else env,
        )

    run = {name: str(runs / name) for name, _ in PASSES}
    subprocess.run(
        [
            sys.executable, str(ROOT / "merge_votes.py"), run["orders-02"], run["orders-13"],
            "--override", run["hard-2048"], "--add", run["easy-2048"],
            "--add", run["medium-rev02"], "--add", run["medium-rev13"], "--out_dir", str(args.out_dir),
        ],
        check=True,
        env=env,
    )


if __name__ == "__main__":
    main()