import os
import subprocess
import sys


def _install_bundled_deps() -> None:
    """Install transformers from bundled wheels (eval sandbox has no PyPI access)."""
    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()

import re
import csv
import json
import shutil
import tempfile
import unicodedata
from collections import Counter

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# The repo is the working directory at run time, and there is no network.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
MODEL_ID = "."
MAX_NEW_TOKENS = 1200  # was 2000; cut to fit T4 wall clock
MAX_NEW_TOKENS_EXPLAIN = 600
TEMPERATURE = 0.8
TOP_P = 0.95
NUM_SAMPLES = 3  # independent rollouts per problem (was 5; timeout)
MAX_ATTEMPTS = 1  # no per-sample retry; majority vote absorbs misses


def load_tokenizer(model_id: str = "."):
    """Load tokenizer, converting tokenizer.json for older tokenizers if needed."""
    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)

    # Older tokenizers expect merge pairs as "a b" strings, not ["a", "b"] lists.
    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)


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 "
    "(beware: this could be many different things: a word, a 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"
    "As the first part of your answer, reason step by step about (1) the linguistic "
    "rules that can be deduced from the given examples in CONTEXT, and (2) "
    "how to apply them to the given problems in QUERY, and (3) in what format answers need to be returned (words, numbers, phonetic transcriptions, ...). \n"
    "Then write a draft of the final answer. "
    "Subsequently, compare it with the format requirements again, "
    "and verify it's compliant with the deduced rules, and it is complete, i.e. has an answer for each element in QUERY. "
    "If necessary, correct and refine."
    "Finally, write a line that says exactly `FINAL ANSWERS:` "
    "and, below it, write the answers to the items requested in QUERY (not those in CONTEXT),"
    "one answer per line (separated by \n) in the order the items are asked for in the QUERY -- the "
    "bare answer only, no numbering, no quotes, no extra text, according to the given TASK TYPE."
)

SYSTEM_POST_EXPLAIN = (
    "You are a helpful assistant that explains the reasoning behind the answers to the International Linguistics Olympiad problems.\n"
    "You are given the following information:\n"
    "- The context of the problem\n"
    "- The task type\n"
    "- The query\n"
    "- The full model output (reasoning draft and answers)\n"
    "You need to condense that output into a concise explanation focused on the key insights "
    "and linguistic rules deduced and applied.\n"
    "Do not include any other text, do not include the answer in the explanation, "
    "and do not invent any new information."
)

# Prefer a dedicated header line; also allow same-line answers after the colon.
FINAL_ANSWERS_LINE_RE = re.compile(
    r"(?im)^[^\w\n]*final answers?[^\w\n]*:?[ \t]*(?=\n|$)|"
    r"(?im)^[^\w\n]*final answers?\s*:\s*"
)
FINAL_ANSWERS_INLINE_RE = re.compile(
    r"(?is)\bfinal answers?\s*:\s*"
)


def extract_raw_final(text: str) -> str:
    """Return text after the last final-answers marker, or '' if none found."""
    line_matches = list(FINAL_ANSWERS_LINE_RE.finditer(text))
    if line_matches:
        return text[line_matches[-1].end() :]

    inline_matches = list(FINAL_ANSWERS_INLINE_RE.finditer(text))
    if inline_matches:
        return text[inline_matches[-1].end() :]

    return ""


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: str, expected: int, task_type: str) -> list[str]:
    text = text.strip()
    if expected <= 1:
        return [text]

    def try_split(pattern: str) -> list[str] | None:
        parts = [part.strip() for part in re.split(pattern, text) if part.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]

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

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


def parse_answer_lines(text_after_marker: str, query: str, task_type: str) -> list[str]:
    """Parse cleaned answer lines from the raw final-answers section."""
    answers = []
    for line in text_after_marker.splitlines():
        stripped_line = line.strip("`").strip()
        if stripped_line == "":
            continue

        match_numbered_prefix = re.match(r"^\s*\d+[.)]\s+(.*)", stripped_line)
        if match_numbered_prefix:
            cleaned_line = match_numbered_prefix.group(1).strip()
        else:
            cleaned_line = stripped_line

        cleaned_line = re.sub(r"\*\*", "", cleaned_line).strip()

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

        if cleaned_line:
            answers.append(cleaned_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


def postprocess_answer(text, query, task_type):
    """Keep only the content after the last 'FINAL ANSWERS' marker."""
    text_after_marker = extract_raw_final(text)
    if not text_after_marker.strip():
        return []
    return parse_answer_lines(text_after_marker, query, task_type)


def normalize_for_vote(text: str, task_type: str) -> str:
    text = unicodedata.normalize("NFC", text.strip())
    if task_type == "match_letters":
        return text.upper()
    return " ".join(text.split())


def majority_vote(
    rollouts: list[list[str]],
    expected: int,
    task_type: str,
) -> list[str]:
    """Per-item majority vote; ties break toward the first sample."""
    if expected <= 0:
        return []

    # Prefer rollouts whose length matches the expected answer count.
    eligible_idx = [
        i
        for i, rollout in enumerate(rollouts)
        if len(rollout) == expected and any(a.strip() for a in rollout)
    ]
    if not eligible_idx:
        eligible_idx = [
            i
            for i, rollout in enumerate(rollouts)
            if any(a.strip() for a in rollout)
        ]
    if not eligible_idx:
        return []

    final: list[str] = []
    for slot in range(expected):
        tagged = [
            (i, rollouts[i][slot])
            for i in eligible_idx
            if slot < len(rollouts[i]) and rollouts[i][slot].strip()
        ]
        if not tagged:
            final.append("")
            continue

        pairs = [
            (i, normalize_for_vote(ans, task_type), ans) for i, ans in tagged
        ]
        counter = Counter(norm for _, norm, _ in pairs)
        top_count = max(counter.values())
        tied_norms = {norm for norm, count in counter.items() if count == top_count}

        # Tie-break: earliest sample whose answer is among the tied winners.
        for i, norm, ans in sorted(pairs, key=lambda p: p[0]):
            if norm in tied_norms:
                final.append(ans)
                break
    return final


def select_source_sample(
    rollouts: list[list[str]],
    voted: list[str],
    task_type: str,
) -> int:
    """Pick the sample closest to the voted answers (tie: earliest index)."""
    if not rollouts:
        return 0

    best_i = 0
    best_score = -1
    for i, rollout in enumerate(rollouts):
        score = 0
        for slot, ans in enumerate(voted):
            if (
                slot < len(rollout)
                and rollout[slot].strip()
                and normalize_for_vote(rollout[slot], task_type)
                == normalize_for_vote(ans, task_type)
            ):
                score += 1
        if score > best_score:
            best_score = score
            best_i = i
    return best_i


def generate_sample(tok, model, ids) -> str:
    """Sample one completion, retrying if FINAL ANSWERS cannot be extracted."""
    text = ""
    for attempt in range(1, MAX_ATTEMPTS + 1):
        with torch.no_grad():
            out = model.generate(
                ids,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=True,
                temperature=TEMPERATURE,
                top_p=TOP_P,
            )
        text = tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip()
        if extract_raw_final(text).strip():
            return text
        print(f"  retry {attempt}/{MAX_ATTEMPTS}: no FINAL ANSWERS", flush=True)
    return text


def generate_explanation(
    tok, model, context: str, task_type: str, query: str, source_text: str,
) -> str:
    """Condense one selected sample's full output into a short explanation."""
    messages = [
        {"role": "system", "content": SYSTEM_POST_EXPLAIN},
        {
            "role": "user",
            "content": (
                f"CONTEXT:{context.strip()}\n"
                f"TASK TYPE:`{task_type}`\n\n"
                f"QUERY:{query.strip()}\n\n"
                f"MODEL OUTPUT:{source_text.strip()}"
            ),
        },
    ]
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt",
    ).to(model.device)
    with torch.no_grad():
        out = model.generate(
            ids,
            max_new_tokens=MAX_NEW_TOKENS_EXPLAIN,
            do_sample=False,
        )
    return tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip()


tok = load_tokenizer(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, device_map="auto"
).eval()

with open("/tmp/data/test.csv", encoding="utf-8", newline="") as f:
    test_rows = list(csv.DictReader(f))

# Write incrementally so a wall-clock kill still leaves a partial submission.csv.
with open("submission.csv", "w", encoding="utf-8", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
    writer.writeheader()
    f.flush()

    for idx, r in enumerate(test_rows, start=1):
        messages = [
            {"role": "system", "content": SYSTEM},
            {
                "role": "user",
                "content": (
                    f"CONTEXT:{r['context'].strip()}\n"
                    f"TASK TYPE:`{r['task_type']}`\n\n"
                    f"QUERY:{r['query'].strip()}"
                ),
            },
        ]
        ids = tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt",
        ).to(model.device)

        print(f"{idx}/{len(test_rows)} sampling {NUM_SAMPLES}x", flush=True)
        raw_texts: list[str] = []
        rollouts: list[list[str]] = []
        for s in range(1, NUM_SAMPLES + 1):
            text = generate_sample(tok, model, ids)
            answers = postprocess_answer(text, r["query"], r["task_type"])
            raw_texts.append(text)
            rollouts.append(answers)
            print(f"  sample {s}/{NUM_SAMPLES} parsed={answers!r}", flush=True)

        expected = expected_answer_count(r["query"], r["task_type"])
        voted = majority_vote(rollouts, expected, r["task_type"])
        source_i = select_source_sample(rollouts, voted, r["task_type"])
        print(f"  vote -> {voted!r} (explain from sample {source_i + 1})", flush=True)

        explanation = generate_explanation(
            tok, model, r["context"], r["task_type"], r["query"], raw_texts[source_i],
        )

        writer.writerow(
            {
                "id": r["id"],
                "pred": json.dumps(voted, ensure_ascii=False),
                "explanation": explanation,
            }
        )
        f.flush()

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