"""IOL-AI 2026 v3: Qwen2.5-7B-Instruct in 4-bit, per-item prompting + reasoning.

Runs inside the competition sandbox (no internet). The submission repo is the
working directory; model weights are shipped in it and loaded from ".".
Reads /tmp/data/test.csv, writes submission.csv (columns: id, pred).

v3 over v2: Qwen2.5-7B-Instruct instead of 1.5B, quantized to 4-bit (NF4) at
load time with bitsandbytes. (AWQ was a dead end: autoawq pins an exact torch
version and the sandbox's package index serves no torch wheels at all, so the
preinstalled torch must be used as-is.)

v2 over the baseline:
- One generation per numbered item (not one per problem), so a misparsed line
  can't shift every later answer.
- The model reasons first, then emits "FINAL: <answer>"; only that is kept.
- Task-type-specific answer-format instructions (exact match is unforgiving).
- A one-shot worked example demonstrating the reasoning + FINAL format.
- A global time budget that shrinks reasoning space as the deadline nears.
"""

import os
import subprocess
import sys
import time

START_TIME = time.time()

# IOL_SKIP_PIP=1 skips dependency install for local testing (where torch etc.
# are already set up); the sandbox always installs.
if os.environ.get("IOL_SKIP_PIP") != "1":
    # Do NOT list torch here: the sandbox's package index has no torch wheels
    # (any resolver attempt to fetch one fails), so the preinstalled torch
    # must satisfy every requirement as-is.
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q",
         "transformers>=4.43", "accelerate>=0.30", "pandas"],
        check=True,
    )
    # bitsandbytes installed separately and non-fatally: if the sandbox index
    # lacks it, the script falls back to fp16 loading below.
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "bitsandbytes"],
        check=False,
    )

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

# IOL_TEST_CSV / IOL_MODEL_ID are only for local testing; in the sandbox the
# defaults are used.
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
MODEL_ID = os.environ.get("IOL_MODEL_ID", ".")
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
# Stop reasoning generously before the sandbox's 30-minute hard kill; the
# budget also has to cover pip install and model load (already elapsed).
TIME_BUDGET_S = int(os.environ.get("IOL_TIME_BUDGET_S", 25 * 60))

import json
import re

# Progress prints contain IPA characters; don't let a non-UTF-8 console
# (Windows local testing) crash the run.
sys.stdout.reconfigure(encoding="utf-8", errors="replace")

import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

SYSTEM_PROMPT = (
    "You solve International Linguistics Olympiad problems: given a few "
    "examples from an unfamiliar language, you deduce its rules and apply "
    "them. Think step by step, briefly. Then give your answer on a new line "
    "starting with 'FINAL:'. The FINAL line must contain only the answer "
    "itself: no explanation, no quotes, no numbering."
)

# One worked example (Hakhun, from the competition page) demonstrating the
# reasoning style and the FINAL format.
ONE_SHOT_USER = """DATA
Here are sentences in Hakhun with their English translations:
ŋa ka kɤ ne | Do I go?
nɤ ʒip tuʔ ne | Did you sleep?
ati ka la | He goes.
ŋa ʒip ku | I sleep.

TASK
Translate into English.
Answer format: only the translation, nothing else.

ITEM
nɤ ʒip ku ne"""

ONE_SHOT_ASSISTANT = """From the data: ŋa = I, nɤ = you, ati = he. ka = go, ʒip = sleep. \
Sentence-final ne marks a question. The verb suffix differs by subject and tense: \
kɤ/ku are present forms (kɤ with ŋa+ka, ku with ŋa+ʒip), tuʔ is past, la is 3rd person present. \
Here: nɤ (you) + ʒip (sleep) + ku (present) + ne (question) = present-tense question about you sleeping.
FINAL: Do you sleep?"""

# Answer-format hints per task_type; unknown types get the generic hint.
TASK_HINTS = {
    "translation": "only the translation, nothing else",
    "fill_blanks": "only the word or words that fill the blank",
    "match_letters": "only the letter or letters of the matching choice",
    "text_to_num": "only the number in digits",
    "num_to_text": "only the number written out in the problem language",
}
GENERIC_HINT = "only the answer itself, as short as possible"

# A numbered item looks like "17. ..." or "17) ..." at the start of a line.
ITEM_RE = re.compile(r"^\s*\(?(\d+)\s*[.)]\s*", re.MULTILINE)
# Models drift to variants like "Final Answer:" or "Answer:"; accept them all.
FINAL_RE = re.compile(r"(?:FINAL(?:\s+ANSWER)?|ANSWER)\s*:\s*(.*)", re.IGNORECASE)


def split_query(query: str):
    """Split the query into (instruction, [(number, item_text), ...])."""
    matches = list(ITEM_RE.finditer(query))
    if not matches:
        return query.strip(), [("1", query.strip())]
    instruction = query[: matches[0].start()].strip()
    items = []
    for i, m in enumerate(matches):
        end = matches[i + 1].start() if i + 1 < len(matches) else len(query)
        items.append((m.group(1), query[m.end(): end].strip()))
    return instruction, items


QUOTED_RE = re.compile(r"[\"“”'']([^\"“”'']+)[\"“”'']")


def extract_final(text: str) -> str:
    """Take the last FINAL: line; fall back to the last non-empty line."""
    finals = FINAL_RE.findall(text)
    if finals:
        return finals[-1].strip().strip('"').strip()
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    if not lines:
        return ""
    last = lines[-1]
    # Without a FINAL marker the last line is often prose like
    # 'The sentence "X" translates to "Y".' — take the last quoted span.
    quoted = QUOTED_RE.findall(last)
    if quoted and len(last) > 2 * len(quoted[-1]):
        return quoted[-1].strip()
    return last


def main():
    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    # NF4 4-bit quantization: ~5 GB on the T4, fp16 compute. Well-supported
    # on T4 (no bf16 needed) and avoids the fp16-overflow garbage seen when
    # running full-precision Qwen2.5 in fp16.
    try:
        quant = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, quantization_config=quant, device_map="auto"
        ).eval()
        print("Loaded 4-bit (bitsandbytes NF4).", flush=True)
    except Exception as e:
        # Last resort if bitsandbytes is unavailable: fp16 barely fits the
        # 16 GB T4 for 7B and risks Qwen2.5 fp16 overflow, but a degraded
        # run beats no submission.csv at all.
        print(f"4-bit load failed ({e!r}); falling back to fp16.", flush=True)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, torch_dtype=torch.float16, device_map="auto"
        ).eval()

    df = pd.read_csv(TEST_CSV, dtype=str).fillna("")

    # Count items up front so the time budget can be split evenly.
    parsed = []
    total_items = 0
    for _, r in df.iterrows():
        instruction, items = split_query(r["query"])
        parsed.append((r, instruction, items))
        total_items += len(items)

    def generate(messages, max_new_tokens):
        prompt = tok.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )
        enc = tok(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            out = model.generate(
                **enc,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                temperature=None,
                top_p=None,
                top_k=None,
                pad_token_id=tok.eos_token_id,
            )
        return tok.decode(
            out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True
        ).strip()

    rows = []
    done_items = 0
    for r, instruction, items in parsed:
        hint = TASK_HINTS.get(r["task_type"], GENERIC_HINT)
        answers = []
        for num, item_text in items:
            remaining_s = TIME_BUDGET_S - (time.time() - START_TIME)
            remaining_items = total_items - done_items
            # Even share of the remaining budget, assuming ~10 tokens/s for
            # NF4 7B on a T4; clamp to a sane range. If time has essentially
            # run out, emit a 32-token answer-only guess.
            share_s = max(remaining_s, 1) / max(remaining_items, 1)
            max_new = int(min(384, max(32, share_s * 10)))

            user_msg = (
                f"DATA\n{r['context'].strip()}\n\n"
                f"TASK\n{instruction or r['query'].strip()}\n"
                f"Answer format: {hint}.\n\n"
                f"ITEM\n{item_text}"
            )
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": ONE_SHOT_USER},
                {"role": "assistant", "content": ONE_SHOT_ASSISTANT},
                {"role": "user", "content": user_msg},
            ]
            text = generate(messages, max_new)
            answer = extract_final(text)
            answers.append(answer)
            done_items += 1
            elapsed = time.time() - START_TIME
            print(f"[{r['id']}] item {num}: {answer!r} "
                  f"(max_new={max_new}, t={elapsed:.0f}s)", flush=True)

        rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})

    pd.DataFrame(rows).to_csv(OUT_CSV, index=False)
    print(f"Wrote {len(rows)} rows to {OUT_CSV} "
          f"in {time.time() - START_TIME:.0f}s")


if __name__ == "__main__":
    main()