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

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

One model (google/gemma-4-E4B-it, 7,941,100,832 parameters loaded) is used in two modes:
  1. option scoring for every question (4 cyclic option orderings, no generation);
  2. thinking mode (reason, then "Answer: N") for every question, run in the groups below,
     with the option-scoring answer as the fallback when no answer can be read.
The model code is src/smmlu.py (downloaded from this repository when it is not next to this file).

Each step writes its own file under outputs/ and is skipped when that file already exists, so a run
that stops (e.g. a 12 h session limit; the full run takes ~13.5 h on 2x T4) continues from the next
step when `python inference.py` is started again in the same folder.
"""

# ======================= 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 = "google/gemma-4-E4B-it"
REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"  # pinned Hugging Face commit
CODE_REPO = "SinhalaMMLU-SharedTask/team-mission-x"
SEED = 42
BATCH = 16          # thinking mode; batch composition affects padding, so keep groups and batch sizes
SCORE_BATCH = 4     # option scoring
MAX_NEW_TOKENS = 2048
OUT_DIR = "outputs"

# Thinking-mode groups, in run order: tag -> routing flags (level, or "level:category").
GROUPS = {
    "hard": ["--gen-levels", "hard"],
    "med2": ["--gen-rules", "medium:social science", "medium:other"],
    "easy2": ["--gen-rules", "easy:stem", "easy:social science", "easy:other"],
    "easy3": ["--gen-rules", "easy:humanities"],
    "med3": ["--gen-rules", "medium:language"],
    "med4": ["--gen-rules", "medium:humanities"],
}

import json
import os
import random
import subprocess
import sys
from pathlib import Path

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. YOUR SYSTEM ===============
def load_system():
    """Returns the path of src/smmlu.py; the model is loaded inside each step (one process per step)."""
    local = Path(__file__).resolve().parent / "src" / "smmlu.py"
    if local.exists():
        return str(local)
    from huggingface_hub import hf_hub_download
    return hf_hub_download(CODE_REPO, "src/smmlu.py")


def to_record(r):
    """Same fields as the merged file used for the submission: level, subject and category come from
    the file path only when a record lacks them."""
    rec = {k: v for k, v in r.items() if k != "_file"}
    p = Path(r["_file"])
    stem_parts = p.stem.rsplit("_", 2)
    if not (rec.get("level") or rec.get("difficulty")):
        rec["level"] = p.parent.name.lower()
    if not rec.get("subject") and len(stem_parts) == 3:
        rec["subject"] = stem_parts[0]
    if not rec.get("category") and len(stem_parts) == 3:
        rec["category"] = stem_parts[1]
    rec["_source_file"] = p.name
    return rec


def step(cmd, output):
    if output.exists():
        print(f"skip (exists): {output}")
        return
    print("+", " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True, env={**os.environ, "SMMLU_REVISION": REVISION})


def predict_all(system, rows):
    smmlu, py = system, sys.executable
    out_dir = Path(OUT_DIR)
    out_dir.mkdir(parents=True, exist_ok=True)
    data = out_dir / "test.jsonl"
    with open(data, "w", encoding="utf-8", newline="\n") as f:
        for r in rows:
            f.write(json.dumps(to_record(r), ensure_ascii=False) + "\n")

    # 1. Option scoring for all questions (the fallback inside the thinking groups).
    scores = out_dir / "final_scoring_scores.jsonl"
    step([py, smmlu, "run", "--model", MODEL_ID, "--data", str(data), "--tag", "final_scoring",
          "--perms", "4", "--with-subject", "--batch-size", str(SCORE_BATCH), "--seed", str(SEED),
          "--out-dir", str(out_dir)], scores)

    # 2. Thinking-mode answers, one run per group (the model is spread over all visible GPUs).
    parts, levels, rules = [], [], []
    for tag, flags in GROUPS.items():
        part = out_dir / f"final_{tag}_s0_gen.jsonl"
        step([py, smmlu, "generate", "--thinking", "--with-subject", "--max-new-tokens", str(MAX_NEW_TOKENS),
              "--batch-size", str(BATCH), "--model", MODEL_ID, "--data", str(data), "--tag", f"final_{tag}",
              *flags, "--seed", str(SEED), "--out-dir", str(out_dir)], part)
        parts.append(str(part))
        (levels if flags[0] == "--gen-levels" else rules).extend(flags[1:])

    # 3. Combine: thinking answers, option scoring where none could be read.
    preds = out_dir / "final_predictions.jsonl"
    if preds.exists():
        preds.unlink()
    step([py, smmlu, "combine", "--data", str(data), "--tag", "final", "--parts", *parts,
          "--ll-scores", str(scores), "--gen-levels", *levels, "--gen-rules", *rules,
          "--out-dir", str(out_dir)], preds)

    out = [json.loads(line) for line in open(preds, encoding="utf-8")]
    assert len(out) == len(rows), "combine must return one answer per question"
    for r, p in zip(rows, out):
        assert (str(p["q_no"]), str(p["subject"])) == (str(r["q_no"]), str(to_record(r)["subject"])), \
            f"order mismatch at q_no={r.get('q_no')}"
    return [p["answer"] for p in out]


# =============== 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)