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

The organizers will run this file exactly as it is, changing ONLY the DATA_REPO line
below to point to the final Sinhala MMLU dataset:

    pip install -r requirements.txt
    python inference.py

System (see README.md): Gemma 4 E4B-it, text-only (7.463B parameters, vision/audio towers removed, weights unchanged),
no fine-tuning. Every question is answered with the model's thinking mode, greedy decoding only (no sampling).
A fast one-step answer (next-token scores over the labels 1-4, averaged over 4 cyclic option orders) is used only as a
fallback if a thinking answer cannot be read. Set SELF_CONSISTENCY = True for an earlier, sampling-based variant
(not the submitted system).

Section 2 runs the exact pipeline that produced our Codabench submission (inference/predict.py with
configs/submission_base_think_greedy.yaml), on the questions loaded by section 1.
"""

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

OUTPUT_FILE = "predictions.jsonl"

# ----------------------- TEAM SETTINGS (keep exactly as used for your final submission)
MODEL_ID = "SinhalaMMLU-SharedTask/team-nexus-gemma-4-e4b-it-text"  # text-only extraction of google/gemma-4-E4B-it
MODEL_REVISION = "8f686039aad1a865c9335c17874d9b870fdec766"  # exact commit of MODEL_ID used for our submission
SEED = 1234
SELF_CONSISTENCY = False  # False = our submitted system (greedy thinking only, no sampling); True = earlier sampling variant, NOT submitted

import json
import os
import random
import re

import numpy as np
import torch
from huggingface_hub import HfFileSystem, list_repo_files


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


# =============== 1. LOAD THE DATASET FROM HUGGING FACE (do not change) ===============
def load_questions(data_repo):
    fs = HfFileSystem()
    files = sorted(f for f in list_repo_files(data_repo, repo_type="dataset") if f.endswith(".json"))
    rows = []
    for f in files:
        with fs.open(f"datasets/{data_repo}/{f}", "r", encoding="utf-8") as fh:
            raw = json.load(fh)
        if isinstance(raw, dict):
            raw = next(v for v in raw.values() if isinstance(v, list))
        for r in raw:
            r["_file"] = f
            rows.append(r)
    return rows


def get_level(r):
    return (r.get("metadata") or {}).get("difficulty") or r["_file"].split("/")[0]


# =============== 2. OUR SYSTEM ===============
# predict_all() writes the loaded questions to a temporary folder (same file layout as the dataset repo), runs
# inference/predict.py on it offline, and maps its answers back to `rows` by (level, subject, q_no).

HERE = os.path.dirname(os.path.abspath(__file__))
CONFIG = os.path.join(HERE, "configs", "submission_base_think_gated_sc.yaml" if SELF_CONSISTENCY
                      else "submission_base_think_greedy.yaml")


def load_system():
    from huggingface_hub import snapshot_download

    return snapshot_download(MODEL_ID, revision=MODEL_REVISION)  # local folder with weights + tokenizer


def predict_all(system, rows):
    import subprocess
    import sys
    import tempfile
    import zipfile

    import yaml

    model_dir = system
    with tempfile.TemporaryDirectory() as tmp:
        by_file = {}
        for r in rows:
            by_file.setdefault(r["_file"], []).append({k: v for k, v in r.items() if k != "_file"})
        for f, recs in by_file.items():
            path = os.path.join(tmp, "data", f)
            os.makedirs(os.path.dirname(path), exist_ok=True)
            with open(path, "w", encoding="utf-8") as fh:
                json.dump(recs, fh, ensure_ascii=False)
        with open(CONFIG, encoding="utf-8") as fh:
            cfg = yaml.safe_load(fh)
        cfg["model"] = model_dir
        cfg["seed"] = SEED
        cfg_path = os.path.join(tmp, "config.yaml")
        with open(cfg_path, "w", encoding="utf-8") as fh:
            yaml.safe_dump(cfg, fh, allow_unicode=True)
        out_zip = os.path.join(tmp, "predictions.zip")
        subprocess.run([sys.executable, "-m", "inference.predict", "--config", cfg_path,
                        "--input", os.path.join(tmp, "data"), "--output", out_zip], cwd=HERE, check=True)
        with zipfile.ZipFile(out_zip) as z:
            preds = [json.loads(l) for l in z.read("predictions.jsonl").decode("utf-8").splitlines() if l.strip()]
    ans = {(p["level"], p["subject"], p["q_no"]): p["answer"] for p in preds}
    return [ans[(get_level(r), r["subject"], r["q_no"])] for r in rows]


# =============== 3. WRITE PREDICTIONS (do not change) ===============
if __name__ == "__main__":
    set_seed(SEED)
    rows = load_questions(DATA_REPO)
    print(len(rows), "questions loaded from", DATA_REPO)

    system = load_system()
    answers = predict_all(system, rows)

    assert len(answers) == len(rows), "predict_all() must return one answer per question"
    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:
        for r, a in zip(rows, answers):
            a = str(a).strip()
            assert a in {"1", "2", "3", "4"}, f"invalid answer {a!r} for q_no={r.get('q_no')}"
            f.write(json.dumps({
                "q_no": r["q_no"],
                "subject": r["subject"],
                "category": r["category"],
                "level": get_level(r),
                "answer": a,
            }, ensure_ascii=False, separators=(",", ":")) + "\n")
    print("\nSaved", OUTPUT_FILE)