import os
# The evaluation sandbox has NO internet. The model's weights are shipped inside
# this repo and loaded from the local folder ".", with offline mode forced. We do
# NOT pip install anything: transformers, torch and pandas are already in the
# sandbox, and autoawq (needed for AWQ models) is preinstalled too.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import json, re, shutil, tempfile
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "."          # the model's weights ship inside this repo
MAX_NEW_TOKENS = 2048    # this is a 1.7B model, cheap to run -- room for a longer answer block


def load_tokenizer(model_id: str = "."):
    """Load the tokenizer, patching tokenizer.json if it uses a merges format
    the sandbox's (older) tokenizers build can't parse. Newer exports sometimes
    store BPE merges as [["a","b"], ...] (list-of-lists) instead of the older
    ["a b", ...] (space-joined strings), which raises:
    'data did not match any variant of untagged enum ModelWrapper'."""
    tokenizer_path = os.path.join(model_id, "tokenizer.json")
    with open(tokenizer_path, encoding="utf-8") as handle:
        data = json.load(handle)

    merges = data.get("model", {}).get("merges", [])
    if not merges or not isinstance(merges[0], list):
        return AutoTokenizer.from_pretrained(model_id)

    data["model"]["merges"] = [" ".join(piece) for piece in merges]
    tmpdir = tempfile.mkdtemp()
    for name in ("tokenizer_config.json", "special_tokens_map.json"):
        src = os.path.join(model_id, name)
        if os.path.isfile(src):
            shutil.copy(src, tmpdir)
    with open(os.path.join(tmpdir, "tokenizer.json"), "w", encoding="utf-8") as handle:
        json.dump(data, handle)
    return AutoTokenizer.from_pretrained(tmpdir)


# 1) Load the model shipped in this repo (float16 = the T4's native precision).
tok = load_tokenizer(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
).eval()

# 2) Read the hidden test set the platform mounts for us (one row per problem).
df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")

# 3) How we ask: let the model reason, then write its answers after a marker.
SYSTEM = (
    "You solve International Linguistics Olympiad problems by reasoning from the "
    "data in CONTEXT you are given to solve the problems in QUERY.\n"
    "There are common TASK TYPES that we specify below, but you may meet a TASK "
    "TYPE you have never seen: read the instruction and the examples, and answer "
    "the QUERY in the same form they use.\n\n"
    "Common TASK TYPES and what to return:\n"
    "`translation`: return the translated form only, in the language the task asks for;\n"
    "`fill_blanks`: return only the missing form for each indicated blank "
    "(this could be a word, part of a word, or a phonetic transcription -- pay close "
    "attention to what part of the CONTEXT is missing in QUERY);\n"
    "`match_letters`: return only the option letter (for example A, B, C);\n"
    "`text_to_num`: return the number in digits;\n"
    "`num_to_text`: return the number written out in words, in the language asked;\n"
    "any other type: return exactly what the instruction asks for, nothing else.\n\n"
    "First, reason step by step about (1) the linguistic rules that can be deduced "
    "from the examples in CONTEXT, and (2) how to apply them to the items in QUERY. "
    "Then write a draft answer, check it against the format requirements and the "
    "deduced rules, and make sure it has one answer for every item in QUERY. Correct "
    "it if needed.\n"
    "Finally, write a line that says exactly FINAL ANSWERS: and, below it, the "
    "answers to the items in QUERY (not those already given in CONTEXT), one "
    "answer per line, in the order the items are asked for -- the bare answer "
    "only, no numbering, no quotes, no extra text."
)

def expected_answer_count(query: str, task_type: str) -> int:
    if task_type == "match_letters":
        numbered = re.findall(r"^\s*\d+\.", query, re.MULTILINE)
        return len(numbered) or 1
    if "blanks" in query.lower():
        range_match = re.search(r"\((\d+)-(\d+)\)", query)
        if range_match:
            return int(range_match.group(2)) - int(range_match.group(1)) + 1
        return len(re.findall(r"\(\d+\)", query)) or 1
    numbered = re.findall(r"^\s*\d+[.)]", query, re.MULTILINE)
    return len(numbered) or 1

def split_single_line_answer(text, expected, task_type):
    text = text.strip()
    if expected <= 1:
        return [text]

    def try_split(pattern):
        parts = [p.strip() for p in re.split(pattern, text) if p.strip()]
        return parts if len(parts) == expected else None

    if task_type == "match_letters":
        for pattern in (r"\s+", r",\s*", r";\s*"):
            if result := try_split(pattern):
                return result
        letters = re.findall(r"[A-Za-z]", text)
        if len(letters) == expected:
            return [letter.upper() for letter in letters]
        return [text]

    for pattern in (r";\s*", r",\s*", r"\s+"):
        if result := try_split(pattern):
            return result
    return [text]

def parse_answers(text, query, task_type):
    """Keep only the lines after the last 'FINAL ANSWERS:' marker, one per line.
    We drop the reasoning above it and return the answers in order; the scorer
    lines our list up against the reference by position."""
    marker = list(re.finditer(r"(?im)^[^\w\n]*final answers?[^\w\n]*:?\s*$", text))
    if not marker:
        return []
    text = text[marker[-1].end():]

    answers = []
    for line in text.splitlines():
        line = line.strip("`").strip()
        if not line:
            continue
        numbered = re.match(r"^\s*\d+[.)]\s+(.*)", line)
        line = numbered.group(1).strip() if numbered else line
        line = re.sub(r"\*\*", "", line).strip()

        if task_type == "match_letters":
            parts = [p.strip("().[]") for p in re.split(r"[\s,;]+", line) if p.strip()]
            if not (len(parts) > 1 and all(re.fullmatch(r"[A-Za-z]", p) for p in parts)):
                m = re.match(r"^\s*(?:\(([A-Za-z])\)|\[([A-Za-z])\]|([A-Za-z]))\.?:?\s*(.*)$", line)
                if m:
                    line = (m.group(1) or m.group(2) or m.group(3)).upper()

        if line:
            answers.append(line)

    expected = expected_answer_count(query, task_type)
    if len(answers) == 1 and expected > 1:
        answers = split_single_line_answer(answers[0], expected, task_type)
    return answers

# 4) Answer every problem, in order, and write the submission file.
rows = []
for i, r in df.iterrows():
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": (
            f"CONTEXT:\n{r['context'].strip()}\n\n"
            f"TASK TYPE: `{r['task_type']}`\n\n"
            f"QUERY:\n{r['query'].strip()}"
        )},
    ]
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt",
        enable_thinking=False,  # Qwen3 supports a <think> mode; off keeps output short and predictable
    ).to(model.device)
    with torch.no_grad():
        out = model.generate(ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
    text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
    answers = parse_answers(text, r["query"], r["task_type"])
    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
    print(f"[{i + 1}/{len(df)}] {len(answers)} answers", flush=True)

pd.DataFrame(rows).to_csv("submission.csv", index=False)