"""
IOL-AI Challenge 2026 — submission script (OFFLINE / Mode B).

Runtime facts (Space Submission tab):
  * T4 medium, 16 GB VRAM, Python 3.10, 30-min wall clock.
  * NO internet: cannot pip install or download anything. Model weights must be
    committed into THIS repo (the working dir) and loaded from ".". Only the
    pre-installed libraries/versions are available (torch 2.4.0, transformers
    4.44.1, accelerate 0.34.2, bitsandbytes 0.43.3, autoawq 0.2.7, pandas 2.2.2,
    numpy 2.1.3, ...). Do NOT pin different majors of torch/transformers/numpy.
  * Read hidden test set from /tmp/data/test.csv; write submission.csv here.
  * pred = JSON list, one entry per numbered item, in query order.

Ship the model in the repo with build_repo.py. This script loads it from "." with
bitsandbytes 4-bit by default so a ~7B fits 16 GB. T4 has no bf16 -> use float16.

Local dev: set IOL_TEST_CSV to a mock file. Quantization auto-disables if there's
no CUDA so the plumbing can be exercised on CPU with a tiny model.
"""

import os
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import re
import csv
import json

MODEL_DIR = os.environ.get("IOL_MODEL_DIR", ".")        # weights live in the repo
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
MAX_NEW_TOKENS = int(os.environ.get("IOL_MAX_NEW_TOKENS", "1024"))
# "4bit" (bitsandbytes), "awq" (weights already AWQ-quantized), or "fp16".
QUANT = os.environ.get("IOL_QUANT", "4bit")

SYSTEM_PROMPT = (
    "You are an expert competitor at the International Linguistics Olympiad. "
    "Each problem gives data from a language you have never seen; deduce its "
    "grammar and vocabulary using ONLY the data and hints in the problem. "
    "Think step by step, then give your final answers.\n\n"
    "OUTPUT FORMAT (strict): after any reasoning, output a line containing only "
    "the token <ANSWERS>, then one answer per numbered item, in order, each on "
    "its own line, with NO item numbers and NO extra commentary. Answer each item "
    "in the language the query asks for (matching items: the option letter; number "
    "items: digits or the written-out number as asked). Give your single best "
    "answer for every item — never leave one blank."
)


def count_items(query):
    """Number of numbered items in a query, e.g. '17. .. 18. ..' -> 2."""
    nums = re.findall(r"(?m)^\s*(\d+)[\.\)]", query)
    return len(nums) if nums else 1


def parse_answers(text, n_items):
    """Pull the final answer block and normalise to exactly n_items lines."""
    if "<ANSWERS>" in text:
        text = text.rsplit("<ANSWERS>", 1)[1]
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    cleaned = [re.sub(r"^\s*(\d+[\.\)]|[-*])\s*", "", ln).strip() for ln in lines]
    cleaned = [c for c in cleaned if c]
    if len(cleaned) < n_items:
        cleaned += [cleaned[-1] if cleaned else ""] * (n_items - len(cleaned))
    return cleaned[:n_items]


def _already_quantized(model_dir):
    """True if the shipped weights are pre-quantized (e.g. AWQ) — then transformers
    auto-detects the config and we must NOT stack bitsandbytes on top."""
    cfg = os.path.join(model_dir, "config.json")
    try:
        with open(cfg, encoding="utf-8") as f:
            return "quantization_config" in json.load(f)
    except Exception:
        return False


def load_model():
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM

    tok = AutoTokenizer.from_pretrained(MODEL_DIR)
    if not torch.cuda.is_available():
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_DIR, torch_dtype=torch.float32).eval()   # CPU dev fallback
        return tok, model

    kwargs = dict(torch_dtype=torch.float16, device_map="auto")  # T4 has no bf16
    if _already_quantized(MODEL_DIR):
        pass  # AWQ/pre-quant: transformers reads quantization_config from config.json
    elif QUANT == "4bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )
    model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, **kwargs).eval()
    return tok, model


def main():
    import torch
    tok, model = load_model()

    with open(TEST_CSV, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))

    dev = model.device if hasattr(model, "device") else "cpu"
    out = []
    for i, r in enumerate(rows):
        context = (r.get("context") or "").strip()
        query = (r.get("query") or "").strip()
        n_items = count_items(query)
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": context + "\n\n" + query},
        ]
        ids = tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt"
        ).to(dev)
        with torch.no_grad():
            gen = model.generate(
                ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=False,
                pad_token_id=tok.eos_token_id,
            )
        text = tok.decode(gen[0][ids.shape[-1]:], skip_special_tokens=True).strip()
        answers = parse_answers(text, n_items)
        out.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
        print("%d/%d done" % (i + 1, len(rows)), flush=True)

    with open(OUT_CSV, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["id", "pred"])
        w.writeheader()
        w.writerows(out)
    print("wrote %s (%d rows)" % (OUT_CSV, len(out)), flush=True)


if __name__ == "__main__":
    main()