"""
Sinhala MMLU Shared Task - Final inference script
Team: <TEAM NAME>

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

Rules
  * Everything (models, adapters, extra data) must be downloaded from Hugging Face inside
    this script. No local paths, no Google Drive / Kaggle paths.
  * No API keys or tokens in the code. If something is gated, the organizers log in themselves.
  * No closed or paid APIs. Only your open-weight system (8B total inference parameters or fewer).
  * Deterministic: no sampling (do_sample=False) and a fixed seed, so a re-run gives the same answers.
  * Do not rely on file names, question counts or answers in the dataset. Load every .json file
    in DATA_REPO, as section 1 does. The final dataset has the same structure as the test dataset.
"""

# ======================= 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-turing-labs-gemma-4-lora"  # your model/adapter repo
ADAPTER_SUBFOLDER = "lora_adapter_v1"
SEED = 42
BATCH = 16
MAX_NEW_TOKENS = 16

import json
import os
import random
import re

import numpy as np
import torch
from huggingface_hub import HfFileSystem, list_repo_files
from peft import PeftConfig, PeftModel
from transformers import AutoModelForMultimodalLM, AutoProcessor, BitsAndBytesConfig


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):
    dataset_dir = os.path.join("training", "dataset")
    files = []
    for root, _, names in os.walk(dataset_dir):
        for name in names:
            if name.endswith(".json"):
                files.append(os.path.join(root, name))
    files.sort()
    rows = []
    for f in files:
        with open(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"] = os.path.relpath(f, dataset_dir).replace(os.sep, "/")
            rows.append(r)
    return rows


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


# =============== 2. YOUR SYSTEM (adapted for A/B/C/D training setup) ===============
def load_system():
    
    peft_config = PeftConfig.from_pretrained(MODEL_ID, subfolder=ADAPTER_SUBFOLDER)
    proc = AutoProcessor.from_pretrained(peft_config.base_model_name_or_path)
    base_model = AutoModelForMultimodalLM.from_pretrained(
        peft_config.base_model_name_or_path,
        dtype=torch.bfloat16, 
        device_map="auto",
    )
    model = PeftModel.from_pretrained(
        base_model, MODEL_ID, subfolder=ADAPTER_SUBFOLDER
    )
    model.eval()
    return proc, model


def make_prompt(r):
    letters = ["A", "B", "C", "D"]
    choices_text = "\n".join([f"{letters[idx]}) {choice}" for idx, choice in enumerate(r["choices"])])
    
    # Matching the exact prompt format used during training
    return (
        "Answer the following multiple-choice question. Reply with the letter of the correct option.\n\n"
        f"Question: {r['question']}\n"
        f"{choices_text}\n\n"
        "Answer:"
    )


def extract_answer(text):
    # Extract the first letter A, B, C, or D generated by the model and map it to 1, 2, 3, or 4
    m = re.search(r"[A-Da-d]", text)
    if m:
        letter = m.group(0).upper()
        mapping = {"A": "1", "B": "2", "C": "3", "D": "4"}
        return mapping.get(letter, "1")
    return "1"


def predict_all(system, rows):
    proc, model = system
    tok = proc.tokenizer
    tok.padding_side = "left"
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    answers = [None] * len(rows)
    for b in range(0, len(rows), BATCH):
        batch = rows[b:b + BATCH]
        texts = [
            proc.apply_chat_template(
                [{"role": "user", "content": [{"type": "text", "text": make_prompt(r)}]}],
                tokenize=False, 
                add_generation_prompt=True
            ) for r in batch
        ]
        enc = tok(texts, return_tensors="pt", padding=True, add_special_tokens=False).to(model.device)
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
        decoded = tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        for j, d in enumerate(decoded):
            answers[b + j] = extract_answer(d)
        print(f"{min(b + BATCH, len(rows))}/{len(rows)}", end="\r")
    return answers


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