"""IOL-AI 2026 submission — Tiny Aya bakeoff v4.

Force-close <|END_THINKING|>, 1536+512, parser v3, prompt v3.
CoT fallback (1024) when format/parse fails.
M1: set USER_THINK_TOKEN="/think" + instruction-following addendum.
A1: leave USER_THINK_TOKEN="".
"""

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

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 = 1536
ANSWER_CONTINUATION_TOKENS = 512
COT_MAX_NEW_TOKENS = 1024  # T4 30min: capped CoT, not another 1536+512

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

_M1_ADDENDUM = """

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

SYSTEM = _SYSTEM_BASE + (_M1_ADDENDUM if USER_THINK_TOKEN == "/think" else "")
SYSTEM_COT = (
    SYSTEM
    + "\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
    # sequential A,B,C… for a long prefix
    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 = [
        {"role": "system", "content": system},
        {"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"
        )


@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]

    think_out = model.generate(
        prompt_ids,
        max_new_tokens=THINKING_BUDGET,
        do_sample=False,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )[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,
        do_sample=False,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )[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,
        do_sample=False,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
    )[0]
    text = tok.decode(out[prompt_len:], skip_special_tokens=False)
    return _TURN_NOISE.sub("", text).strip()


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 = f"{r['context'].strip()}\n\n{r['query'].strip()}"
    if n_guess:
        user += f"\n\n(Emit exactly {n_guess} answer line(s) after FINAL ANSWERS:.)"
    user_think = user
    if USER_THINK_TOKEN:
        user_think += f"\n{USER_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: no /think, thinking channel off, capped budget
        cot_ids = _build_prompt_ids(tok, SYSTEM_COT, user, 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)})
    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)