"""IOL-AI 2026 — M1 iterate on best private~0.083 (temp0.6 + user-instr).

Keep: user-prompt instructions, /think, sample think @ TEMPERATURE.
CSV fixes from that run: strip _GCY/gloss; reject essay+alphabet+resample;
stronger user prompt; think 2048; greedy answer continuation.
"""

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

# "" for A1 (reasoning_options only); "/think" for M1 multilingual
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|>"

THINKING_BUDGET = 2048
ANSWER_CONTINUATION_TOKENS = 512
COT_MAX_NEW_TOKENS = 1024
TEMPERATURE = 0.6
TOP_P = 0.95
QUALITY_RESAMPLES = 4
ANSWER_GREEDY = True  # format-sensitive answer phase

# Empty system — instructions go on the user turn (decode ablation vs v4)
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."
)

# --- 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 _strip_gloss_keep_form(line: str) -> str:
    s = (line or "").strip()
    s = re.sub(r"\*\*", "", s)
    s = re.split(r"\s+_?(?:GCY|NS|N/A)_?\b", s, maxsplit=1, flags=re.I)[0].strip()
    s = re.sub(r"\s+_?(?:GCY|NS|N/A)_?\s*$", "", s, flags=re.I).strip()
    m = re.match(
        r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being|means?|black|white|red|green|yellow)\b.*)$",
        s,
        flags=re.I,
    )
    if m:
        s = m.group(1).strip()
    return s.strip()


def _clean_line(line: str) -> str:
    line = _NUMBERING.sub("", line).strip()
    line = _MD_WRAP.sub("", line).strip()
    line = line.replace("\u202f", " ").replace("\xa0", " ")
    return _strip_gloss_keep_form(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) < 8:
        return False
    seq = 0
    for i, L in enumerate(letters):
        if ord(L) == ord("A") + i:
            seq += 1
        else:
            break
    return seq >= 8


def _looks_like_essay(answers: list[str]) -> bool:
    if any(len(str(a)) > 100 for a in answers):
        return True
    blob = " ".join(map(str, answers))
    return bool(
        re.search(
            r"(?i)\b(colors? are expressed|systematic set of lexical|"
            r"these stems are combined|step by step|as an ai|verification)\b",
            blob,
        )
    )


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 _looks_like_essay(answers):
        return False
    if any(re.search(r"(?i)_GCY\b|_NS\b", str(a)) for a in answers):
        return False
    if task_type != "match_letters" and len(answers) >= 4:
        if all(re.fullmatch(r"[ivxlcdm]+", str(a).strip(), flags=re.I) for a in answers):
            return False
    if n_expected is not None and n_expected > 0 and abs(len(answers) - n_expected) > max(
        2, n_expected // 2
    ):
        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

    ans_kw = dict(do_sample=False) if ANSWER_GREEDY else gen_kw
    full = model.generate(
        cont.unsqueeze(0),
        max_new_tokens=ANSWER_CONTINUATION_TOKENS,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
        **ans_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="",
    )

    best_answers: list[str] = []
    used_cot = False
    for attempt in range(QUALITY_RESAMPLES):
        torch.manual_seed(1000 + int(i) * 97 + attempt * 31)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(1000 + int(i) * 97 + attempt * 31)
        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)
        answers = [_strip_gloss_keep_form(a) for a in answers if _strip_gloss_keep_form(a)]
        if n_exp and n_exp > 0 and len(answers) > n_exp:
            answers = answers[:n_exp]
        if has_usable_answer(answers, n_expected=n_exp, task_type=task):
            best_answers = answers
            break
        if answers and not best_answers:
            best_answers = answers
        print(
            f"  quality resample {attempt + 1}/{QUALITY_RESAMPLES} id={r['id']} n={len(answers)}",
            flush=True,
        )
    answers = best_answers

    if not has_usable_answer(answers, n_expected=n_exp, task_type=task):
        used_cot = True
        # CoT: no /think, thinking channel off, capped budget
        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)
        cot_answers = [
            _strip_gloss_keep_form(a) for a in cot_answers if _strip_gloss_keep_form(a)
        ]
        if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task) or (
            cot_answers and not answers
        ):
            answers = cot_answers

    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    print(
        f"[{i + 1}/{len(df)}] {len(answers)} answers cot={used_cot}",
        flush=True,
    )

pd.DataFrame(rows).to_csv("submission.csv", index=False)
print("wrote submission.csv", flush=True)