"""IOL-AI 2026 submission — linguistic-reasoning SFT on M1 (+/think).

Decode: sample at temperature 0.6 (not greedy), force-close <|END_THINKING|>,
then answer continuation. Instructions live in the **user** prompt (empty
system) — better than system-prompt instructions for this model.

Token budget: think 2000 + answer 512 — Space hard-kills at 1800s (T4).
4096 timed out; 2000 is ~OK for ~14 IOL rows. CoT fallback 1024, same temp.
Writes submission.csv after every row so a late timeout still leaves a file.
"""

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 = "."

# M1 multilingual / linguistic-reasoning SFT: append /think on the user turn
USER_THINK_TOKEN = "/think"

import json
import re

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

END_THINKING = "<|END_THINKING|>"
START_THINKING = "<|START_THINKING|>"

# Space limit = 1800s. 4096 timed out; 2000 estimated OK for ~14 IOL rows on T4.
THINKING_BUDGET = 2000
ANSWER_CONTINUATION_TOKENS = 512
COT_MAX_NEW_TOKENS = 1024
TEMPERATURE = 0.6
TOP_P = 0.95

# Minimal system — task instructions go on the user turn
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
- 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, never option text.
- Never append glosses like "form – meaning" or "word - gloss"; bare answers only.
- 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."
)

# --- parser (inlined from parse_iol.py) ---
_MD_PREFIX = r"(?:[#*_=\-\s`>]*)"
_MARKER = re.compile(
    rf"(?im)^{_MD_PREFIX}final\s+answers?{_MD_PREFIX}:?{_MD_PREFIX}\s*(.*)$"
)
_NUMBERING = re.compile(r"^\s*(?:\d+[.)]|[-*•])\s*")
_TURN_NOISE = re.compile(
    r"<\|/?END_OF_TURN_TOKEN\|>|<\|/?START_OF_TURN_TOKEN\|>|"
    r"<\|CHATBOT_TOKEN\|>|<EOS_TOKEN>|<BOS_TOKEN>"
)
_RESPONSE_BLOCK = re.compile(
    r"<\|START_RESPONSE\|>(.*?)<\|END_RESPONSE\|>",
    flags=re.S,
)
_MD_WRAP = re.compile(r"^[*_`#\s]+|[*_`#\s]+$")
_TRAILING_LETTER = re.compile(
    r"(?:[–—\-]|→|->)\s*([A-Za-z])(?:\s*[.)]|)\s*$"
)
_LEADING_LETTER_OPT = re.compile(r"^([A-Za-z])\s*[.):\-–—]\s+\S")
_WORD_THEN_LETTER = re.compile(r"^.+\s([A-Za-z])\s*$")
_REFUSAL = re.compile(
    r"(?i)\b("
    r"i'?m sorry|i am sorry|i don'?t have|i cannot|i can'?t|"
    r"unable to|not able to|no reliable|cannot supply|can'?t supply|"
    r"as an ai|i apologize"
    r")\b"
)


def _clean_line(line: str) -> str:
    line = _NUMBERING.sub("", line).strip()
    line = _MD_WRAP.sub("", line).strip()
    line = line.replace("\u202f", " ").replace("\xa0", " ")
    return line.strip()


def after_thinking(text: str) -> str:
    """Prefer content after the last <|END_THINKING|>; else drop an unclosed think block."""
    if END_THINKING in text:
        text = text.rsplit(END_THINKING, 1)[-1]
    elif START_THINKING in text:
        text = ""
    return _TURN_NOISE.sub("", text)


def _as_option_letter(line: str) -> str | None:
    line = _clean_line(line)
    if not line:
        return None
    if len(line) == 1 and line.isalpha():
        return line.upper()
    m = _LEADING_LETTER_OPT.match(line)
    if m:
        return m.group(1).upper()
    m = _TRAILING_LETTER.search(line)
    if m:
        return m.group(1).upper()
    if len(line) <= 40:
        m = _WORD_THEN_LETTER.match(line)
        if m:
            return m.group(1).upper()
    return None


def _expand_line(line: str) -> list[str]:
    line = _clean_line(line)
    if not line:
        return []
    if len(line) == 1 and line.isalpha():
        return [line]
    if _LEADING_LETTER_OPT.match(line) or _TRAILING_LETTER.search(line):
        letter = _as_option_letter(line)
        if letter:
            return [letter]
    if len(line) <= 40 and _WORD_THEN_LETTER.match(line):
        letter = _as_option_letter(line)
        if letter:
            return [letter]
    if "|" in line:
        parts = [p.strip() for p in line.split("|") if p.strip()]
        if len(parts) >= 2:
            if len(parts) >= 4 and len(parts) % 2 == 0:
                left, right = parts[0::2], parts[1::2]
                if sum(" " in r for r in right) >= max(1, len(right) // 2):
                    return [_clean_line(x) for x in left if _clean_line(x)]
            if len(parts) == 2:
                a, b = parts
                if (" " in b and " " not in a) or (
                    len(b) > 2 * max(len(a), 1) and " " in b
                ):
                    return [_clean_line(a)] if _clean_line(a) else []
            return [_clean_line(p) for p in parts if _clean_line(p)]
    return [line]


def _dedupe_runaway(parts: list[str]) -> list[str]:
    if len(parts) < 6:
        return parts
    out: list[str] = []
    run = 0
    prev = None
    for p in parts:
        if p == prev:
            run += 1
            if run >= 4:
                break
        else:
            run = 1
            prev = p
        out.append(p)
    return out


def _lines_from_region(region: str, *, allow_all_lines: bool) -> list[str]:
    markers = list(_MARKER.finditer(region))
    if markers:
        last = markers[-1]
        after_parts: list[str] = []
        same = _clean_line(last.group(1) or "")
        if same:
            after_parts.extend(_expand_line(same))
        for line in region[last.end() :].splitlines():
            after_parts.extend(_expand_line(line))
        if after_parts:
            return _dedupe_runaway(after_parts)
        before_parts: list[str] = []
        for line in region[: last.start()].splitlines():
            before_parts.extend(_expand_line(line))
        if before_parts:
            return _dedupe_runaway(before_parts)

    parts: list[str] = []
    for line in region.splitlines():
        parts.extend(_expand_line(line))
    if not parts:
        return []
    if allow_all_lines:
        return _dedupe_runaway(parts)
    return [parts[-1]]


def parse_answers(
    raw: str,
    *,
    n_expected: int | None = None,
    task_type: str = "",
) -> list[str]:
    text = after_thinking(raw)
    closed_blocks = _RESPONSE_BLOCK.findall(text)
    answers: list[str] = []
    if closed_blocks:
        for region in reversed(closed_blocks):
            answers = _lines_from_region(region.strip(), allow_all_lines=True)
            if answers:
                break
    if not answers:
        answers = _lines_from_region(text, allow_all_lines=False)

    if task_type == "match_letters":
        coerced: list[str] = []
        for a in answers:
            letter = _as_option_letter(a)
            coerced.append(letter if letter else a)
        answers = coerced

    if n_expected is not None and n_expected > 0 and len(answers) > n_expected:
        answers = answers[:n_expected]
    return answers


def _looks_like_alphabet_dump(answers: list[str]) -> bool:
    letters = [a.strip().upper() for a in answers if len(a.strip()) == 1 and a.strip().isalpha()]
    if len(letters) < 15:
        return False
    seq = 0
    for i, L in enumerate(letters):
        if ord(L) == ord("A") + i:
            seq += 1
        else:
            break
    return seq >= 15


def _looks_like_refusal(answers: list[str]) -> bool:
    blob = " ".join(answers)
    return bool(_REFUSAL.search(blob)) or len(blob) > 400 and "dictionary" in blob.lower()


def has_usable_answer(
    answers: list[str],
    *,
    n_expected: int | None = None,
    task_type: str = "",
) -> bool:
    if not answers or not any(a.strip() for a in answers):
        return False
    if _looks_like_refusal(answers):
        return False
    if _looks_like_alphabet_dump(answers):
        return False
    if n_expected is not None and n_expected > 0 and len(answers) != n_expected:
        return False
    if task_type == "match_letters":
        letters = [a for a in answers if len(a) == 1 and a.isalpha()]
        if len(letters) < max(1, int(0.8 * len(answers))):
            return False
    return True


def _n_items_guess(query: str) -> int:
    nums = re.findall(r"(?m)^\s*(?:\(?\d+[.)]|\d+\))", query)
    return len(nums) if nums else 0


def _end_thinking_id(tok) -> int:
    end_id = tok.convert_tokens_to_ids(END_THINKING)
    if end_id is None or end_id == tok.unk_token_id:
        ids = tok.encode(END_THINKING, add_special_tokens=False)
        if len(ids) == 1:
            end_id = ids[0]
    if end_id is None or end_id == tok.unk_token_id:
        raise RuntimeError(f"Tokenizer missing end-think token {END_THINKING!r}")
    return int(end_id)


def _build_prompt_ids(tok, system: str, user: str, *, thinking: bool):
    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",
            reasoning_options={"enabled": thinking},
        )
    except TypeError:
        return tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt"
        )


def _sample_kwargs():
    return dict(
        do_sample=True,
        temperature=TEMPERATURE,
        top_p=TOP_P,
    )


@torch.inference_mode()
def generate_with_think_budget(model, tok, prompt_ids, end_id: int):
    device = next(model.parameters()).device
    prompt_ids = prompt_ids.to(device)
    prompt_len = prompt_ids.shape[-1]
    gen_kw = _sample_kwargs()

    think_out = model.generate(
        prompt_ids,
        max_new_tokens=THINKING_BUDGET,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
        **gen_kw,
    )[0]
    gen_ids = think_out[prompt_len:].tolist()
    if end_id not in gen_ids:
        cont = torch.cat(
            [think_out, torch.tensor([end_id], device=device, dtype=think_out.dtype)]
        )
    else:
        cont = think_out

    full = model.generate(
        cont.unsqueeze(0),
        max_new_tokens=ANSWER_CONTINUATION_TOKENS,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
        **gen_kw,
    )[0]
    text = tok.decode(full[prompt_len:], skip_special_tokens=False)
    return _TURN_NOISE.sub("", text).strip()


@torch.inference_mode()
def generate_plain(model, tok, prompt_ids, max_new_tokens: int):
    device = next(model.parameters()).device
    prompt_ids = prompt_ids.to(device)
    prompt_len = prompt_ids.shape[-1]
    out = model.generate(
        prompt_ids,
        max_new_tokens=max_new_tokens,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
        **_sample_kwargs(),
    )[0]
    text = tok.decode(out[prompt_len:], skip_special_tokens=False)
    return _TURN_NOISE.sub("", text).strip()


def _build_user(
    instructions: str,
    context: str,
    query: str,
    *,
    n_guess: int,
    think_token: str = "",
) -> str:
    parts = [
        instructions.strip(),
        "",
        context.strip(),
        "",
        query.strip(),
    ]
    if n_guess:
        parts.append("")
        parts.append(f"(Emit exactly {n_guess} answer line(s) after FINAL ANSWERS:.)")
    if think_token:
        parts.append(think_token.strip())
    return "\n".join(parts)


tok = AutoTokenizer.from_pretrained(MODEL_ID)
end_id = _end_thinking_id(tok)
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():
    n_guess = _n_items_guess(r["query"])
    task = str(r.get("task_type", "") or "")
    n_exp = n_guess or None

    user_think = _build_user(
        USER_INSTRUCTIONS,
        r["context"],
        r["query"],
        n_guess=n_guess,
        think_token=USER_THINK_TOKEN,
    )
    user_plain = _build_user(
        USER_INSTRUCTIONS_COT,
        r["context"],
        r["query"],
        n_guess=n_guess,
        think_token="",
    )

    ids = _build_prompt_ids(tok, SYSTEM, user_think, thinking=True)
    text = generate_with_think_budget(model, tok, ids, end_id)
    answers = parse_answers(text, n_expected=n_exp, task_type=task)
    used_cot = False

    if not has_usable_answer(answers, n_expected=n_exp, task_type=task):
        used_cot = True
        cot_ids = _build_prompt_ids(tok, SYSTEM, user_plain, thinking=False)
        cot_text = generate_plain(model, tok, cot_ids, COT_MAX_NEW_TOKENS)
        cot_answers = parse_answers(cot_text, n_expected=n_exp, task_type=task)
        if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task):
            answers = cot_answers

    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    # Flush after each row — Space timeout deletes unfinished runs otherwise
    pd.DataFrame(rows).to_csv("submission.csv", index=False)
    print(
        f"[{i + 1}/{len(df)}] {len(answers)} answers cot={used_cot}",
        flush=True,
    )

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