"""
IOL-AI 2026 submission script. Ships inside the model repo alongside the Qwen2.5-14B-Instruct-AWQ
weights. Reads /tmp/data/test.csv, writes submission.csv (id, pred, explanation) to the working
directory. No internet at runtime -- everything must load from local files (".").
"""
import os

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import json
import re
import time

START_TIME = time.time()
TIME_LIMIT_SECONDS = 30 * 60
SAFETY_MARGIN_SECONDS = 90              # reserve for CSV write + any per-row overrun
DEADLINE = START_TIME + TIME_LIMIT_SECONDS - SAFETY_MARGIN_SECONDS

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

SYSTEM = (
    "You solve International Linguistics Olympiad problems. You will be given data "
    "from a language you have never seen before, plus hints, and asked to answer "
    "numbered items about it. The ONLY source of truth is the data given to you in "
    "this problem -- do not rely on anything you think you know about the language "
    "if it conflicts with the examples given here.\n\n"
    "Work method:\n"
    "1. Go through every single example given, in order. For each word in the "
    "unfamiliar language, first split it into its likely component morphemes "
    "(stem plus any prefixes/suffixes) even if you are not fully sure of the "
    "boundaries -- treat it as a sequence of parts, not one opaque unit. Then, "
    "for each one, write down "
    "every distinct morpheme, word, particle, or structural pattern (word order, "
    "marking, alternation) it contains and what it appears to mean or mark. Do not "
    "skip any example, and do not stop early once you have a plausible-looking "
    "pattern -- an alternation that looks like it marks one thing (e.g. tense) may "
    "actually mark something else (e.g. person, number, or agreement with a "
    "different argument), and only the examples you skipped may reveal which.\n"
    "2. Write a section titled RULE TABLE: that lists, as a table or bullet list, "
    "every distinct piece you identified and its meaning/function -- this must "
    "cover every example from step 1, not just the ones similar to the query.\n"
    "3. Using ONLY entries from your RULE TABLE, work out the answer to each query "
    "item. If the table has no entry for something the query needs, say so and give "
    "your best-supported guess rather than leaving it blank.\n"
    "4. Before writing your final answers, re-read your RULE TABLE and re-derive "
    "each answer from it one more time, checking: did you actually apply every "
    "rule you stated (e.g. a plural marker, a tense marker) to every relevant "
    "answer, not just some of them? If two answers could plausibly be swapped "
    "(e.g. two options assigned to the wrong item), re-check the evidence that "
    "distinguishes them specifically.\n\n"
    "Output format by task type -- give exactly this, nothing more:\n"
    "- translation: the translated form only, in the language the task asks for.\n"
    "- fill_blanks: only the missing form for each blank.\n"
    "- match_letters: only the option letter (e.g. A, B, C).\n"
    "- text_to_num: the number in digits.\n"
    "- num_to_text: the number written out in words, in the language asked.\n"
    "- any other task type: give exactly what the instruction asks for, nothing else.\n\n"
    "Each final answer must be the bare form only -- no surrounding quotes, no "
    "trailing period or punctuation that isn't part of the answer itself, no "
    "parenthetical notes, no alternate options separated by '/' or 'or'. Pick "
    "one single best answer per item.\n\n"
    "Follow the work method above, showing your RULE TABLE. Then write a line that "
    "says exactly FINAL ANSWERS: and, below it, one answer per line in the order "
    "the items are asked -- the bare answer only, no numbering, no quotes, no "
    "extra commentary."
)


def count_expected_items(query):
    """Count numbered items (e.g. '17.', '18)') in the query -- the target answer count."""
    return len(re.findall(r"(?m)^\s*\d+[.)]", query))


MAX_EXPLANATION_CHARS = 2000


def parse_answers(text, expected_count=None):
    """Keep only the lines after the last 'FINAL ANSWERS:' marker, one answer per line.

    If expected_count is given, pad with "" or truncate so the row never silently
    drops points from a length mismatch against the scorer's positional alignment.
    """
    marker = list(re.finditer(r"(?im)^[#*\s]*final answers?[:#*\s]*", text))
    if marker:
        text = text[marker[-1].end():]
    answers = []
    for line in text.splitlines():
        line = re.sub(r"^\s*\d+[.)]\s*", "", line).strip()
        line = line.strip("\"'")
        if line:
            answers.append(line)

    if expected_count:
        if len(answers) < expected_count:
            answers = answers + [""] * (expected_count - len(answers))
        elif len(answers) > expected_count:
            answers = answers[:expected_count]
    return answers


def extract_explanation(text):
    """The reasoning/RULE TABLE portion before the final-answers marker, for the
    Human Evaluation track -- truncated so one long row can't bloat the CSV."""
    marker = re.search(r"(?im)^[#*\s]*final answers?[:#*\s]*$", text)
    explanation = text[:marker.start()] if marker else text
    explanation = explanation.strip()
    if len(explanation) > MAX_EXPLANATION_CHARS:
        explanation = explanation[:MAX_EXPLANATION_CHARS].rsplit(" ", 1)[0] + " ..."
    return explanation


MODEL_ID = "."
MAX_NEW_TOKENS_FULL = 2048   # normal pass -- room to show a full RULE TABLE
MAX_NEW_TOKENS_FAST = 768    # fallback pass when time is short -- less room to reason

print(f"[{time.time() - START_TIME:.0f}s] loading model...", flush=True)
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto",
).eval()
PAD_TOKEN_ID = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id

MAX_INPUT_TOKENS = 6000   # conservative guard against the model's context window -- leaves
                          # headroom for the system prompt, chat template overhead, and
                          # MAX_NEW_TOKENS_FULL generation on unusually long real IOL problems


def truncate_context(context, query):
    """Guard against exceeding the model's context window -- truncates `context` (never
    `query`, which holds the actual questions) from the end, keeping as much of the given
    data as fits within MAX_INPUT_TOKENS."""
    query_len = len(tok(query, add_special_tokens=False)["input_ids"])
    budget = MAX_INPUT_TOKENS - query_len
    if budget <= 0:
        return context
    context_ids = tok(context, add_special_tokens=False)["input_ids"]
    if len(context_ids) <= budget:
        return context
    return tok.decode(context_ids[:budget], skip_special_tokens=True)
print(f"[{time.time() - START_TIME:.0f}s] model loaded", flush=True)


def generate(system, user, max_new_tokens):
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]
    # tokenize=False + separate tok(...) call, rather than apply_chat_template(..., return_dict=True),
    # for compatibility with the eval sandbox's pinned transformers==4.44.1 (return_dict support on
    # apply_chat_template was added later).
    prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
    enc = tok(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **enc, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.5, top_p=0.95,
            pad_token_id=PAD_TOKEN_ID, repetition_penalty=1.15,
        )
    return tok.decode(out[0][enc["input_ids"].shape[-1]:], skip_special_tokens=True).strip()


def solve_row(problem, expected, avg_row_time):
    """One greedy pass, with a shorter token budget if the time budget is running low."""
    full_pass_estimate = avg_row_time if avg_row_time else 45.0   # seconds -- rough guess for row 1
    remaining = DEADLINE - time.time()
    if remaining > full_pass_estimate:
        text = generate(SYSTEM, problem, MAX_NEW_TOKENS_FULL)
    else:
        # Running low on time: fewer tokens, same prompt.
        text = generate(SYSTEM, problem, MAX_NEW_TOKENS_FAST)
    return parse_answers(text, expected_count=expected), extract_explanation(text)


df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
print(f"[{time.time() - START_TIME:.0f}s] {len(df)} rows to solve", flush=True)

rows_out = []
row_times = []
for i, r in df.iterrows():
    row_start = time.time()
    problem = f"{truncate_context(r['context'].strip(), r['query'].strip())}\n\n{r['query'].strip()}"
    expected = count_expected_items(r["query"])
    avg_row_time = sum(row_times) / len(row_times) if row_times else None

    if time.time() > DEADLINE:
        # Out of time: best-effort placeholder rather than risking the whole process
        # getting killed by the 30-minute hard limit with no submission.csv at all.
        answers = [""] * max(expected, 1)
        explanation = ""
    else:
        try:
            answers, explanation = solve_row(problem, expected, avg_row_time)
        except Exception as e:
            print(f"[{time.time() - START_TIME:.0f}s] row {r['id']} failed: {e}", flush=True)
            answers = [""] * max(expected, 1)
            explanation = ""

    rows_out.append({
        "id": r["id"],
        "pred": json.dumps(answers, ensure_ascii=False),
        "explanation": explanation,
    })
    row_times.append(time.time() - row_start)
    print(
        f"[{time.time() - START_TIME:.0f}s] row {i + 1}/{len(df)} done in {row_times[-1]:.0f}s, "
        f"{len(answers)} answers, {DEADLINE - time.time():.0f}s budget left",
        flush=True,
    )

pd.DataFrame(rows_out).to_csv("submission.csv", index=False)
print(f"[{time.time() - START_TIME:.0f}s] wrote submission.csv ({len(rows_out)} rows)", flush=True)