"""IOL Space script — math-ling-hyp ckpt-5327 (M1 base) + soft end-think.

Decode: user-prompt instructions, /think, soft P(END_THINKING)>=0.55,
think T=0.4 budget 2048, answer T=0.6, CoT fallback.
Quality: strip gloss keeping form; reject alphabet dumps / essays and resample.
"""


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

USER_THINK_TOKEN = "/think"

import json
import re
from collections import Counter

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

try:
    from transformers.generation.logits_process import LogitsProcessor
except ImportError:  # pragma: no cover
    from transformers import LogitsProcessor

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

THINKING_BUDGET = 2048
ANSWER_CONTINUATION_TOKENS = 512
COT_MAX_NEW_TOKENS = 1024
THINK_TEMPERATURE = 0.4
THINK_TOP_P = 0.95
ANSWER_TEMPERATURE = 0.0
ANSWER_TOP_P = 0.95
# Soft-exit when P(END_THINKING) is already high
SOFT_END_PROB = 0.55
SOFT_END_MIN_TOKENS = 96
# Detect repetition loops → soft exit
LOOP_WINDOW = 24
LOOP_REPEAT = 3
# 1 = single sample; 3 = majority vote (use only if time allows)
MAJORITY_K = 1
# Reject bad format / essay / alphabet and resample full think+answer
QUALITY_RESAMPLES = 4

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.
- Never append junk tokens or codes (e.g. _GCY) to answers."""

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 _strip_gloss_keep_form(line: str) -> str:
    """Keep linguistic form; drop English gloss / _GCY_ tails (fix 1)."""
    s = line.strip()
    # [ipa] _GCY_ | meaning  OR  form _NS_ | gloss
    s = re.split(r"\s+_?(?:GCY|NS|N/A)_?\s*\|\s*", s, maxsplit=1, flags=re.I)[0]
    s = re.sub(r"\s+_?(?:GCY|NS|N/A)_?\b.*$", "", s, flags=re.I).strip()
    # form - to be / means ...
    m = re.match(
        r"^(.+?)\s+[-–—]\s+((?:to|the|a|an|in|of|for|being|means?)\b.*)$",
        s,
        flags=re.I,
    )
    if m:
        s = m.group(1).strip()
    # Dedup "[a] [a]"
    m_dup = re.fullmatch(r"(\[[^\]]+\])\s+\1", s)
    if m_dup:
        s = m_dup.group(1)
    return s.strip()


def _expand_line(line: str) -> list[str]:
    line = _clean_line(line)
    if not line:
        return []
    if re.match(r"(?i)^(verification|notes?|explanation|reasoning)\s*:", line):
        return []
    line = _strip_gloss_keep_form(line)
    if not line:
        return []
    # Space-separated letter blob: "A J L M O P Y"
    if re.fullmatch(r"(?:[A-Za-z]\s+)+[A-Za-z]", line.strip()):
        return [p.upper() for p in line.split()]
    # Comma-joined multi answers (Julia 204-style)
    if "," in line and not re.search(r"\d,\d", line):
        parts = [p.strip() for p in line.split(",") if p.strip()]
        if len(parts) >= 3:
            return [_strip_gloss_keep_form(p) for p in parts if _strip_gloss_keep_form(p)]
    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:
    blob = " ".join(answers)
    if any(len(a) > 120 for a in answers):
        return True
    if re.search(
        r"(?i)\b(colors are expressed|step by step|in this language|verification|"
        r"the following rules|as an ai)\b",
        blob,
    ):
        return True
    return False


def _looks_like_gloss_leak(answers: list[str]) -> bool:
    return any(re.search(r"(?i)_GCY_|_NS_\s*\|", a) for a in answers)


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 _looks_like_gloss_leak(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


class SoftEndThinkProcessor(LogitsProcessor):
    """When P(END_THINKING) is already high (or a loop is detected), force it."""

    def __init__(
        self,
        end_id: int,
        *,
        threshold: float,
        min_tokens: int,
        loop_window: int,
        loop_repeat: int,
    ):
        self.end_id = int(end_id)
        self.threshold = float(threshold)
        self.min_tokens = int(min_tokens)
        self.loop_window = int(loop_window)
        self.loop_repeat = int(loop_repeat)
        self.prompt_len = None
        self.forced = False

    def set_prompt_len(self, n: int) -> None:
        self.prompt_len = int(n)
        self.forced = False

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
        if self.forced:
            scores[:] = torch.finfo(scores.dtype).min
            scores[:, self.end_id] = 0.0
            return scores

        gen_len = 0
        if self.prompt_len is not None:
            gen_len = int(input_ids.shape[-1] - self.prompt_len)

        # Repetition loop: same window repeated
        if gen_len >= self.loop_window * self.loop_repeat:
            seq = input_ids[0, -self.loop_window * self.loop_repeat :].tolist()
            w = self.loop_window
            chunk = seq[-w:]
            if all(seq[i : i + w] == chunk for i in range(0, len(seq) - w, w)):
                self.forced = True
                scores[:] = torch.finfo(scores.dtype).min
                scores[:, self.end_id] = 0.0
                return scores

        if gen_len < self.min_tokens:
            return scores

        # Soft exit on high END_THINKING probability
        probs = torch.softmax(scores[0].float(), dim=-1)
        p_end = float(probs[self.end_id].item())
        if p_end >= self.threshold:
            self.forced = True
            scores[:] = torch.finfo(scores.dtype).min
            scores[:, self.end_id] = 0.0
        return scores


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 _clean_pred_line(s: str) -> str:
    s = (s or "").strip()
    s = _strip_gloss_keep_form(s)
    s = re.sub(r"\s+_+\w{2,6}$", "", s).strip()
    return s


def _post_answers(answers: list[str], n_expected: int | None) -> list[str]:
    answers = [_clean_pred_line(a) for a in answers if _clean_pred_line(a) or a == ""]
    answers = [a for a in answers if a != ""]
    if n_expected and n_expected > 0:
        if len(answers) > n_expected:
            answers = answers[:n_expected]
    return answers


def majority_vote(samples: list[list[str]], n_expected: int | None) -> list[str]:
    if not samples:
        return []
    n = n_expected or max((len(s) for s in samples), default=0)
    if n <= 0:
        return samples[0]
    padded = [(list(s) + [""] * n)[:n] for s in samples]
    tup_counts = Counter(tuple(s) for s in padded)
    best_tup, best_c = tup_counts.most_common(1)[0]
    if best_c >= 2:
        return list(best_tup)
    return [
        Counter(s[i] for s in padded).most_common(1)[0][0] for i in range(n)
    ]


@torch.inference_mode()
def generate_with_soft_end_think(model, tok, prompt_ids, end_id: int, soft: SoftEndThinkProcessor):
    device = next(model.parameters()).device
    prompt_ids = prompt_ids.to(device)
    prompt_len = prompt_ids.shape[-1]
    soft.set_prompt_len(prompt_len)

    think_out = model.generate(
        prompt_ids,
        max_new_tokens=THINKING_BUDGET,
        do_sample=True,
        temperature=THINK_TEMPERATURE,
        top_p=THINK_TOP_P,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
        logits_processor=[soft],
    )[0]
    gen_ids = think_out[prompt_len:].tolist()
    if end_id not in gen_ids:
        # budget cap fallback — still force-close so answer phase can start
        cont = torch.cat(
            [think_out, torch.tensor([end_id], device=device, dtype=think_out.dtype)]
        )
        soft_forced = False
    else:
        cont = think_out
        soft_forced = soft.forced or True

    # Answer continuation: sample at T=0.6 (user request) + quality resample upstream
    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(), soft.forced, len(gen_ids)


@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()


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()
soft_proc = SoftEndThinkProcessor(
    end_id,
    threshold=SOFT_END_PROB,
    min_tokens=SOFT_END_MIN_TOKENS,
    loop_window=LOOP_WINDOW,
    loop_repeat=LOOP_REPEAT,
)

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

    samples = []
    soft_flags = []
    think_lens = []
    for k in range(max(1, MAJORITY_K)):
        best_answers: list[str] = []
        soft_hit = False
        tlen = 0
        for attempt in range(QUALITY_RESAMPLES):
            torch.manual_seed(1000 + int(i) * 97 + k * 13 + attempt * 31)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(1000 + int(i) * 97 + k * 13 + attempt * 31)
            ids = _build_prompt_ids(tok, SYSTEM, user_think, thinking=True)
            text, soft_hit, tlen = generate_with_soft_end_think(
                model, tok, ids, end_id, soft_proc
            )
            answers = _post_answers(
                parse_answers(text, n_expected=n_exp, task_type=task), 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} "
                f"id={r['id']} n={len(answers)}",
                flush=True,
            )
        answers = best_answers
        if not has_usable_answer(answers, n_expected=n_exp, task_type=task):
            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 = _post_answers(
                parse_answers(cot_text, n_expected=n_exp, task_type=task), n_exp
            )
            if has_usable_answer(cot_answers, n_expected=n_exp, task_type=task) or (
                cot_answers and not answers
            ):
                answers = cot_answers
        samples.append(answers)
        soft_flags.append(soft_hit)
        think_lens.append(tlen)

    if MAJORITY_K > 1:
        answers = majority_vote(samples, n_exp)
    else:
        answers = samples[0]

    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)} soft={soft_flags} "
        f"tlen={think_lens} k={MAJORITY_K}",
        flush=True,
    )

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