#!/usr/bin/env python3
"""IOL-AI 2026 submission entrypoint.

The evaluator runs this file from a public model repository whose root also
contains a quantized Qwen3.5-4B checkpoint.  It reads /tmp/data/test.csv and
writes submission.csv in the current working directory.

The implementation deliberately keeps all benchmark-facing logic in the
standard library. Heavy ML imports happen only inside load_model(), which also
repairs the old packages in the competition base image when necessary.
"""

from __future__ import annotations

import ast
import csv
import importlib.metadata
import json
import os
import re
import subprocess
import sys
import time
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable


TEST_PATH = Path(os.environ.get("IOL_TEST_PATH", "/tmp/data/test.csv"))
OUTPUT_PATH = Path(os.environ.get("IOL_OUTPUT_PATH", "submission.csv"))
MODEL_PATH = os.environ.get("IOL_MODEL_PATH", ".")
RUNTIME_BUDGET_SECONDS = int(os.environ.get("IOL_RUNTIME_BUDGET", "1500"))
MIN_TRANSFORMERS = (5, 8, 0)
MIN_COMPRESSED_TENSORS = (0, 15, 0)
PACKED_WEIGHT_BITS = 4

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


@dataclass
class ProblemRow:
    id: str
    context: str
    query: str
    work_lang: str = ""
    task_lang: str = ""
    task_type: str = ""
    eval_type: str = "single"

    @classmethod
    def from_dict(cls, row: dict[str, str]) -> "ProblemRow":
        return cls(
            id=str(row.get("id", "")).strip(),
            context=str(row.get("context", "")).strip(),
            query=str(row.get("query", "")).strip(),
            work_lang=str(row.get("work_lang", "")).strip(),
            task_lang=str(row.get("task_lang", "")).strip(),
            task_type=str(row.get("task_type", "")).strip(),
            eval_type=str(row.get("eval_type", "single")).strip() or "single",
        )


def version_tuple(value: str) -> tuple[int, int, int]:
    nums = [int(x) for x in re.findall(r"\d+", value)[:3]]
    return tuple((nums + [0, 0, 0])[:3])  # type: ignore[return-value]


def runtime_is_compatible() -> bool:
    try:
        transformers_v = version_tuple(importlib.metadata.version("transformers"))
        compressed_v = version_tuple(importlib.metadata.version("compressed-tensors"))
        return transformers_v >= MIN_TRANSFORMERS and compressed_v >= MIN_COMPRESSED_TENSORS
    except importlib.metadata.PackageNotFoundError:
        return False


def unpacked_weight_shape(packed_shape: Iterable[int], num_bits: int = PACKED_WEIGHT_BITS) -> tuple[int, ...]:
    """Recover the dense shape used by compressed-tensors bit packing."""
    shape = tuple(int(value) for value in packed_shape)
    if not shape or num_bits <= 0 or 32 % num_bits:
        raise ValueError(f"Invalid packed shape or bit width: {shape}, {num_bits}")
    return (*shape[:-1], shape[-1] * (32 // num_bits))


def repair_packed_weight_shapes(model: Any, torch: Any) -> int:
    """Repair integer shape metadata lost by the evaluator's old Torch loader.

    The submitted checkpoint uses symmetric 4-bit pack-quantized weights with
    dimensions divisible by the eight-values-per-int32 packing factor. Its
    safetensors header independently confirms this for every packed tensor.
    """
    repaired = 0
    for module in model.modules():
        packed = getattr(module, "weight_packed", None)
        stored_shape = getattr(module, "weight_shape", None)
        if packed is None or stored_shape is None:
            continue
        expected = unpacked_weight_shape(packed.shape)
        current = tuple(int(value) for value in stored_shape.detach().cpu().tolist())
        if current == expected:
            continue
        replacement = torch.tensor(expected, dtype=stored_shape.dtype, device=stored_shape.device)
        with torch.no_grad():
            stored_shape.copy_(replacement)
        repaired += 1
    return repaired


def ensure_runtime() -> None:
    """Install loaders required by the Qwen3.5 compressed-tensors checkpoint."""
    if runtime_is_compatible():
        return
    wheel_dir = Path(__file__).resolve().parent / "wheels"
    wheels = sorted(wheel_dir.glob("*.whl"))
    if not wheels:
        raise RuntimeError(
            "Offline runtime wheels are missing. The evaluation sandbox cannot "
            "download a modern Transformers release."
        )
    print(f"Installing Qwen3.5 runtime from {len(wheels)} bundled wheels...", flush=True)
    subprocess.run(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "-q",
            "--disable-pip-version-check",
            "--no-cache-dir",
            "--no-index",
            "--no-deps",
            *[str(path) for path in wheels],
        ],
        check=True,
    )


def load_rows(path: Path = TEST_PATH) -> list[ProblemRow]:
    with path.open(newline="", encoding="utf-8") as f:
        rows = [ProblemRow.from_dict(r) for r in csv.DictReader(f)]
    if not rows or any(not r.id for r in rows):
        raise ValueError("test.csv is empty or contains a blank id")
    return rows


def normalize_context(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip()


def group_rows(rows: Iterable[ProblemRow]) -> list[list[ProblemRow]]:
    """Group rows from the same IOL puzzle without changing first-seen order."""
    groups: OrderedDict[str, list[ProblemRow]] = OrderedDict()
    for row in rows:
        groups.setdefault(normalize_context(row.context), []).append(row)
    return list(groups.values())


def _numbered_lines(text: str) -> list[str]:
    return re.findall(r"(?m)^\s*\d+\s*[.)]\s*\S.*$", text)


def count_items(row: ProblemRow) -> int:
    """Infer expected list length from query shape and task metadata."""
    numbered = _numbered_lines(row.query)
    if numbered:
        return len(numbered)

    ranges = re.findall(r"\((\d+)\s*[-–—]\s*(\d+)\)", row.query)
    if ranges:
        # IOL prompts commonly say "Fill the blanks (1–14)" without repeating
        # the blank-bearing table in the Linguini query field. Explicit query
        # lines take precedence because a range can also refer back to context.
        return sum(abs(int(end) - int(start)) + 1 for start, end in ranges)

    placeholders = sorted({int(n) for n in re.findall(r"\((\d+)\)", row.query)})
    if placeholders:
        return len(placeholders)

    task = row.task_type.lower()
    # A historical number-system format names several forms in prose and then
    # gives lettered equations. Keep support for that shape without splitting
    # ordinary translation phrases on commas.
    named_numbers = re.search(
        r"(?is)write\s+the\s+numbers?\s+(.+?)\s+and\s+the\s+equalit(?:y|ies)",
        row.query,
    )
    lettered = re.findall(r"(?m)^\s*[A-Z]\s*[.)]\s*\S.*$", row.query)
    if named_numbers and lettered:
        named_count = named_numbers.group(1).count(",") + 1
        return named_count + len(lettered)

    if task.startswith("match"):
        source_items = _numbered_lines(row.context)
        if source_items:
            return len(source_items)

    # Number tasks in Linguini often put one unnumbered form on each line after
    # a one-line instruction. This also provides a safe fallback for rare task
    # types with the same presentation.
    lines = [ln.strip() for ln in row.query.splitlines() if ln.strip()]
    if len(lines) > 1:
        body = [ln for ln in lines[1:] if not re.fullmatch(r"[-–—]+", ln)]
        if body:
            return len(body)
    return 1


def task_guidance(task_types: Iterable[str]) -> str:
    tasks = {t.lower().strip() for t in task_types}
    notes: list[str] = []
    if any(t.startswith("translat") for t in tasks):
        notes.append(
            "For translation, align examples into a morpheme/word table; track "
            "person, number, tense, case, polarity and word order; copy every "
            "task-language character and diacritic exactly."
        )
    if any(t.startswith("match") for t in tasks):
        notes.append(
            "For matching, solve the correspondence as one global bijection. "
            "Use repeated roots and contrasts, and return only option labels."
        )
    if any("blank" in t for t in tasks):
        notes.append(
            "For blanks, infer the smallest transformation system consistent "
            "with every example, including phonological alternations."
        )
    if any("num" in t for t in tasks):
        notes.append(
            "For numerals, derive the arithmetic base and composition order, "
            "verify the rule on every provided equation, then calculate exactly."
        )
    return "\n".join(f"- {note}" for note in notes)


SYSTEM_PROMPT = """You solve International Linguistics Olympiad problems.
Infer the grammar, lexicon, morphology, sound rules or number system strictly
from the supplied data. Check the rule against every example, but do not print
your reasoning. Accuracy and exact spelling matter: preserve Unicode,
diacritics, punctuation and requested language.

Return only one valid JSON array of answer strings, in item order. Do not add
numbering, markdown, labels, explanations or alternative nested lists. Never
abstain; give the best answer for every requested item."""


def build_row_prompt(row: ProblemRow) -> str:
    guidance = task_guidance([row.task_type])
    return "\n\n".join(
        [
            f"TASK TYPE: {row.task_type or 'unknown'}",
            f"REQUIRED ANSWER COUNT: {count_items(row)}",
            guidance,
            "CONTEXT:\n" + row.context,
            "QUERY:\n" + row.query,
            f"Return exactly {count_items(row)} strings as one JSON array.",
        ]
    )


def answer_token_budget(expected: int) -> int:
    """Bound direct-answer decoding so every puzzle gets GPU time."""
    return min(768, max(128, 64 + 48 * expected))


def _json_candidates(text: str) -> Iterable[Any]:
    cleaned = text.strip()
    if "</think>" in cleaned:
        cleaned = cleaned.rsplit("</think>", 1)[-1].strip()
    cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned, flags=re.I)
    cleaned = re.sub(r"\s*```$", "", cleaned)

    decoder = json.JSONDecoder()
    for index, char in enumerate(cleaned):
        if char not in "[{":
            continue
        try:
            value, _ = decoder.raw_decode(cleaned[index:])
            yield value
        except json.JSONDecodeError:
            continue

    # Models occasionally emit Python list syntax despite the JSON contract.
    try:
        yield ast.literal_eval(cleaned)
    except (ValueError, SyntaxError):
        pass


def _coerce_answer_list(value: Any) -> list[str]:
    if isinstance(value, list):
        out: list[str] = []
        for item in value:
            if isinstance(item, (list, tuple)):
                # Multi-reference rows still require one selected answer.
                item = item[0] if item else ""
            out.append(str(item).strip())
        return out
    if value is None:
        return []
    if isinstance(value, str):
        text = value.strip()
        for parser in (json.loads, ast.literal_eval):
            try:
                parsed = parser(text)
                if isinstance(parsed, list):
                    return _coerce_answer_list(parsed)
            except Exception:
                pass
        numbered = re.findall(r"(?m)^\s*\d+\s*[.)-]\s*(.+?)\s*$", text)
        if numbered:
            return [x.strip() for x in numbered]
        lines = [x.strip(" -\t") for x in text.splitlines() if x.strip()]
        return lines or [text]
    return [str(value).strip()]


def _rows_from_object(obj: Any) -> list[dict[str, Any]]:
    if isinstance(obj, dict) and isinstance(obj.get("rows"), list):
        return [x for x in obj["rows"] if isinstance(x, dict)]
    if isinstance(obj, list):
        return [x for x in obj if isinstance(x, dict)]
    if isinstance(obj, dict):
        rows: list[dict[str, Any]] = []
        for key, value in obj.items():
            if isinstance(value, dict):
                rows.append({"id": key, **value})
            elif isinstance(value, list):
                rows.append({"id": key, "answers": value})
        return rows
    return []


def fit_answers(values: list[str], expected: int) -> list[str]:
    values = [str(v).strip().strip('"').strip("'") for v in values]
    if len(values) >= expected:
        return values[:expected]
    return values + [""] * (expected - len(values))


def parse_model_output(text: str, group: list[ProblemRow]) -> tuple[dict[str, list[str]], dict[str, str]]:
    expected_ids = {r.id for r in group}
    answers: dict[str, list[str]] = {}
    explanations: dict[str, str] = {}
    for candidate in _json_candidates(text):
        for item in _rows_from_object(candidate):
            row_id = str(item.get("id", "")).strip()
            if row_id not in expected_ids:
                continue
            raw_answers = item.get("answers", item.get("pred", item.get("answer")))
            answers[row_id] = _coerce_answer_list(raw_answers)
            explanations[row_id] = str(item.get("explanation", "")).strip()
        if answers:
            break

    # Single-row fallback: accept a bare JSON array.
    if not answers and len(group) == 1:
        for candidate in _json_candidates(text):
            if isinstance(candidate, list) and not any(isinstance(x, dict) for x in candidate):
                answers[group[0].id] = _coerce_answer_list(candidate)
                break
    return answers, explanations


def parse_direct_answers(text: str) -> list[str]:
    """Salvage a direct JSON array or one-answer-per-line response."""
    for candidate in _json_candidates(text):
        if isinstance(candidate, list) and not any(isinstance(x, dict) for x in candidate):
            return _coerce_answer_list(candidate)

    cleaned = text.strip()
    if "</think>" in cleaned:
        cleaned = cleaned.rsplit("</think>", 1)[-1].strip()
    cleaned = re.sub(r"^```(?:json)?\s*", "", cleaned, flags=re.I)
    cleaned = re.sub(r"\s*```$", "", cleaned)
    lines: list[str] = []
    for raw in cleaned.splitlines():
        line = raw.strip()
        if not line or re.fullmatch(r"(?i)answers?:?", line):
            continue
        line = re.sub(r"^\s*(?:[-*•]|\d+\s*[.)-])\s*", "", line).strip()
        if line:
            lines.append(line)
    return lines


def load_model() -> tuple[Any, Any, Any]:
    ensure_runtime()
    import torch

    # compressed-tensors references the torch.nn.Buffer convenience class added
    # after the evaluator's Torch 2.4 image. It only needs a type for isinstance.
    if not hasattr(torch.nn, "Buffer"):
        class _CompatBuffer(torch.Tensor):
            pass

        torch.nn.Buffer = _CompatBuffer

    from transformers import AutoModelForCausalLM, AutoTokenizer

    if not torch.cuda.is_available():
        raise RuntimeError("The competition T4 GPU is not visible")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        local_files_only=True,
        dtype=torch.float16,
        device_map="cuda:0",
        low_cpu_mem_usage=True,
        attn_implementation="sdpa",
    ).eval()
    repaired = repair_packed_weight_shapes(model, torch)
    if repaired:
        print(f"Repaired {repaired} packed weight-shape tensors for Torch 2.4", flush=True)
    return tokenizer, model, torch


def generate(tokenizer: Any, model: Any, torch: Any, prompt: str, max_new_tokens: int, max_time: float) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]
    encoded = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        enable_thinking=False,
        tokenize=True,
        return_tensors="pt",
        return_dict=True,
    )
    encoded = {k: v.to(model.device) for k, v in encoded.items()}
    input_len = encoded["input_ids"].shape[-1]
    with torch.inference_mode():
        output = model.generate(
            **encoded,
            max_new_tokens=max_new_tokens,
            max_time=max(30.0, max_time),
            do_sample=False,
            repetition_penalty=1.0,
            use_cache=True,
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(output[0][input_len:], skip_special_tokens=True).strip()


def repair_prompt(row: ProblemRow, previous: str) -> str:
    return f"""Repair the answer format for this one row. Solve from the context if
needed. Return ONLY a JSON array of exactly {count_items(row)} strings.

CONTEXT:
{row.context}

QUERY:
{row.query}

PREVIOUS ATTEMPT:
{previous[-3000:]}
"""


def solve(rows: list[ProblemRow]) -> tuple[dict[str, list[str]], dict[str, str]]:
    started = time.monotonic()
    deadline = started + RUNTIME_BUDGET_SECONDS
    tokenizer, model, torch = load_model()
    all_answers: dict[str, list[str]] = {}
    all_explanations: dict[str, str] = {}

    for row_index, row in enumerate(rows, 1):
        remaining = deadline - time.monotonic()
        remaining_rows = len(rows) - row_index + 1
        if remaining < 60:
            print("Runtime reserve reached; emitting safe placeholders", flush=True)
            break
        expected = count_items(row)
        token_budget = answer_token_budget(expected)
        per_row_time = min(180.0, max(45.0, remaining / remaining_rows - 10.0))
        print(
            f"Solving puzzle {row_index}/{len(rows)}: "
            f"{expected} answers, budget={token_budget}",
            flush=True,
        )
        text = generate(
            tokenizer,
            model,
            torch,
            build_row_prompt(row),
            max_new_tokens=token_budget,
            max_time=per_row_time,
        )
        parsed, _ = parse_model_output(text, [row])
        values = parsed.get(row.id, []) or parse_direct_answers(text)
        # Spend one short repair only for malformed or incomplete output.
        if len(values) < expected and deadline - time.monotonic() > 90:
            repaired_text = generate(
                tokenizer,
                model,
                torch,
                repair_prompt(row, text),
                max_new_tokens=min(384, 96 + 32 * expected),
                max_time=min(60.0, deadline - time.monotonic() - 45.0),
            )
            repaired, _ = parse_model_output(repaired_text, [row])
            repaired_values = repaired.get(row.id, []) or parse_direct_answers(repaired_text)
            if repaired_values:
                values = repaired_values
        all_answers[row.id] = fit_answers(values, expected)
        all_explanations[row.id] = (
            f"Inferred the {row.task_type or 'linguistic'} pattern from the "
            "provided examples and applied it to each requested item."
        )

    for row in rows:
        all_answers.setdefault(row.id, [""] * count_items(row))
        all_explanations.setdefault(
            row.id,
            f"Attempted to infer the {row.task_type or 'linguistic'} pattern from the supplied context.",
        )
    return all_answers, all_explanations


def write_submission(
    rows: list[ProblemRow],
    answers: dict[str, list[str]],
    explanations: dict[str, str],
    path: Path = OUTPUT_PATH,
) -> None:
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["id", "pred", "explanation"])
        writer.writeheader()
        for row in rows:
            writer.writerow(
                {
                    "id": row.id,
                    "pred": json.dumps(answers[row.id], ensure_ascii=False),
                    "explanation": explanations.get(row.id, ""),
                }
            )


def main() -> None:
    rows = load_rows()
    print(f"Loaded {len(rows)} rows in {len(group_rows(rows))} puzzles", flush=True)
    answers, explanations = solve(rows)
    write_submission(rows, answers, explanations)
    print(f"Wrote {OUTPUT_PATH} with {len(rows)} rows", flush=True)


if __name__ == "__main__":
    main()