"""IOL-AI 2026 — notebook decode (prompt + force-close + T=0.6).

Matches linguini-test-iolai.ipynb cells 4–5:
- USER_INSTRUCTIONS (+ COT + /think)
- think 2048 @ T=0.6; force <|END_THINKING|><|START_RESPONSE|>
- answer 512 @ T=0.6
- marker-gated FINAL ANSWERS: parser
"""

from __future__ import annotations

import os
import subprocess
import sys


def _install_bundled_deps() -> None:
    wheels_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wheels")
    if not os.path.isdir(wheels_dir):
        return
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "-q",
            "--no-index",
            f"--find-links={wheels_dir}",
            "transformers==4.56.2",
        ],
        check=True,
    )


_install_bundled_deps()

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

import json
import re

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

SYSTEM = ""

USER_INSTRUCTIONS = """You solve International Linguistics Olympiad (IOL) problems from the data you are given.
You may see a task type you have never seen: follow the instruction and examples, and answer in the same form they use.
What to return by task type:
- translation: only the required form in the language the query asks for — do not add extra glosses or "form | meaning" unless asked
- fill_blanks: only the missing form for each blank — no extra glosses
- match_letters: ONLY the option letter (A, B, C, …), one letter per line — never copy option text, never arrows, never "A. word"
- text_to_num: the number in digits only
- num_to_text: the number written out in words, in the language asked
- kinship / sentence matching: the full required sentence or form — NOT roman numerals (i, ii, iii) and NOT an alphabet dump
- any other type: exactly what the instruction asks for, nothing else
Answer in the language and form the query asks for. Do not add glosses, translations, or explanations unless the instruction requires them.
Output rules:
- Put answers ONLY after a line that says exactly: FINAL ANSWERS:
- Never put answers before that marker.
- One answer per line; exactly as many lines as items asked in the query.
- Bare answers only: no numbering, no quotes, no commentary, no repeating the question.
Extra hard rules:
- Never refuse or apologize; always output FINAL ANSWERS: with your best guess.
- For match_letters: only bare letters (A, B, C, …) — never dump the alphabet (A B C D E F…), never option text.
- Never append tags or glosses: no "_GCY", "_NS", "form – meaning", "word - gloss", or markdown bold.
- Never write an essay or explanation of how the language works under FINAL ANSWERS: — only the answer strings.
- If the query asks for colour/color forms, output those forms only (one per line), not a linguistics write-up.
- Emit exactly as many answer lines as items asked — no more, no fewer."""

USER_INSTRUCTIONS_COT = (
    USER_INSTRUCTIONS
    + "\n\nThink step by step about the rules in the examples and how they apply to the query, "
    "then write FINAL ANSWERS: and the answer lines."
)

APPEND_THINK = True

THINK_BUDGET = 2048
ANSWER_BUDGET = 512
TEMPERATURE = 0.6


def build_user(context, query):
    instructions = USER_INSTRUCTIONS_COT if APPEND_THINK else USER_INSTRUCTIONS
    parts = [instructions.strip(), "", context.strip(), "", query.strip()]
    if APPEND_THINK:
        parts.append("/think")
    return "\n".join(parts)


_SPECIAL = re.compile(
    r"<\|START_RESPONSE\|>|<\|END_RESPONSE\|>|"
    r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|"
    r"<\|CHATBOT_TOKEN\|>|<\|/?START_THINKING\|>|<\|/?END_THINKING\|>|"
    r"<EOS_TOKEN>|<BOS_TOKEN>|<PAD>"
)
_MARKER = re.compile(r"(?im)\bfinal\s+answers?\b\s*:?\s*")
_TURN_NOISE = re.compile(
    r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|"
    r"<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>"
)


def _lines_after(region: str) -> list[str]:
    answers = []
    for line in region.splitlines():
        line = _SPECIAL.sub("", line)
        line = re.sub(r"^\s*(?:\(?\d+[.)]|[-*•])\s*", "", line).strip()
        if not line:
            if answers:
                break
            continue
        if re.fullmatch(r"(?i)final\s+answers?\s*:?", line):
            continue
        if answers and (
            len(line) > 120
            or line.lower().startswith(("however", "but the", "to solve", "### ", "**"))
        ):
            break
        answers.append(line)
    return answers


def parse_answers(text: str) -> list[str]:
    raw = text

    def _from_region(region: str) -> list[str] | None:
        markers = list(_MARKER.finditer(region))
        if not markers:
            return None
        for m in reversed(markers):
            answers = _lines_after(region[m.end() :])
            if answers:
                return answers
        return []

    region = raw
    if "<|END_THINKING|>" in region:
        region = region.rsplit("<|END_THINKING|>", 1)[-1]
    blocks = re.findall(
        r"<\|START_RESPONSE\|>(.*?)<\|END_RESPONSE\|>", region, flags=re.S
    )
    if blocks:
        region = blocks[-1]

    found = _from_region(region)
    if found is not None:
        return found

    found = _from_region(raw)
    if found is not None:
        return found

    return []


def _build_prompt(tok, system: str, user: str):
    messages = []
    if system.strip():
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": user})
    try:
        return tok.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
            reasoning_options={"enabled": True},
        )
    except TypeError:
        ids = tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt"
        )
        return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}


@torch.inference_mode()
def generate_with_budgets(model, tok, enc):
    device = next(model.parameters()).device
    enc = {k: v.to(device) if hasattr(v, "to") else v for k, v in enc.items()}
    prompt_len = enc["input_ids"].shape[-1]
    pad_id = tok.pad_token_id or tok.eos_token_id
    end_ids = tok.encode("<|END_THINKING|>", add_special_tokens=False)
    start_ids = tok.encode("<|START_RESPONSE|>", add_special_tokens=False)

    out = model.generate(
        **enc,
        max_new_tokens=THINK_BUDGET,
        do_sample=True,
        temperature=TEMPERATURE,
        eos_token_id=tok.eos_token_id,
        pad_token_id=pad_id,
    )
    gen = out[0]
    new_ids = gen[prompt_len:].tolist()
    forced = False
    finished = bool(new_ids) and new_ids[-1] == tok.eos_token_id

    if end_ids[0] not in new_ids:
        forced = True
        cont = torch.tensor(
            end_ids + start_ids, device=gen.device, dtype=gen.dtype
        )
        gen = torch.cat([gen, cont], dim=0)
        finished = False

    if not finished:
        ids = gen.unsqueeze(0)
        out2 = model.generate(
            input_ids=ids,
            attention_mask=torch.ones_like(ids),
            max_new_tokens=ANSWER_BUDGET,
            do_sample=True,
            temperature=TEMPERATURE,
            eos_token_id=tok.eos_token_id,
            pad_token_id=pad_id,
        )
        gen = out2[0]

    text = tok.decode(gen[prompt_len:], skip_special_tokens=False).strip()
    return _TURN_NOISE.sub("", text).strip(), forced


tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
).eval()

df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")

rows = []
for i, r in df.iterrows():
    user = build_user(r["context"], r["query"])
    enc = _build_prompt(tok, SYSTEM, user)
    text, forced = generate_with_budgets(model, tok, enc)
    answers = parse_answers(text)
    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    pd.DataFrame(rows).to_csv("submission.csv", index=False)
    print(f"[{i + 1}/{len(df)}] n={len(answers)} forced={forced}", flush=True)

print("wrote submission.csv", flush=True)