"""IOL-AI 2026 — v12: RECOVERY + safe explanation.

WHAT HAPPENED: v7 (reasoning + ###ANSWERS### + count-fix, single greedy) scored 0.1233
(chrF 0.25). v9/v10/v11 all crashed to ~0.05 (chrF ~0.09 — below the broken v1). The
one structural thing v9-11 added that v7/v8 lacked, and that the WORKING third-party
submissions (v5, test-v4) deliberately avoided, is a MULTI-LINE explanation column.
Embedded newlines in a CSV cell corrupt naive row parsing → every later row's `pred`
is misread → chrF/EM collapse across the board. (v5/test-v4 collapsed their explanation
to ONE line and scored fine.)

THE FIX (this file): reproduce the EXACT v7 answer pipeline — same prompt, same parsing,
single greedy decode, count-fix — and add ONLY a CSV-SAFE single-line explanation derived
from the model's own reasoning. No ###WHY### marker, no prompt changes, no boosters, no
multi-sample pipeline. So it must reproduce v7's score, now with explanation_rate ~100.
Once confirmed, boosters/self-consistency get re-added ONE AT A TIME, each measured.
"""

import os
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

import re
import csv
import json

MODEL_DIR = os.environ.get("IOL_MODEL_DIR", ".")
TEST_CSV = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
OUT_CSV = os.environ.get("IOL_OUT_CSV", "submission.csv")
MAX_NEW_TOKENS = int(os.environ.get("IOL_MAX_NEW_TOKENS", "768"))
QUANT = os.environ.get("IOL_QUANT", "4bit")

ANSWER_MARKER = "###ANSWERS###"

# EXACTLY v7's system prompt (the version that scored 0.1233). Do not add sections.
SYSTEM_PROMPT = (
    "You are an expert competitor at the International Linguistics Olympiad. "
    "Each problem gives data from a language you have never seen; deduce its rules "
    "using ONLY the data and hints in the problem, then answer EVERY sub-question.\n\n"
    "A problem can have MANY sub-questions even when the query is one sentence: e.g. "
    "'give the correspondences' expects one answer for EACH numbered item in the data "
    "(often a dozen or more). Work out how many answers are required and give exactly "
    "that many, one per item, in the order the items appear.\n\n"
    "Reason briefly, then end your reply with the answers in EXACTLY this format, with "
    "nothing after it:\n"
    f"{ANSWER_MARKER}\n"
    "1. <answer to item 1>\n"
    "2. <answer to item 2>\n"
    "(one numbered line per sub-question, in order)\n\n"
    "Each answer line holds ONLY the requested form — a word, phrase, number, or "
    "letter — with no restating of the question and no commentary. Answer in the "
    "language and direction the query asks. For matching items give just the option "
    "letter; for number items give digits or the written-out number as asked. Never "
    "leave an item blank — always give your best guess."
)

TASK_HINT = {
    "translation": "This is a translation task: each answer is only the translated word/phrase.",
    "text_to_num": "This is a number task: each answer is only digits (e.g. 42).",
    "num_to_text": "This is a number task: each answer is only the number written in the target language's words.",
    "match_letters": "This is a matching task: each answer is only the option letter (A, B, C, ...); give one per item in the data.",
    "matching": "This is a matching task: each answer is only the option letter; give one per item in the data.",
    "fill_blank": "This is a fill-in-the-blank task: each answer is only the missing form.",
    "fill_blanks": "This is a fill-in-the-blank task: each answer is only the missing form.",
}


def build_messages(row):
    context = (row.get("context") or "").strip()
    query = (row.get("query") or "").strip()
    ttype = (row.get("task_type") or "").strip().lower()
    system = SYSTEM_PROMPT
    hint = TASK_HINT.get(ttype)
    if hint:
        system = system + "\n\n" + hint
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": context + "\n\n" + query},
    ]


def detect_count(context, query):
    """Sub-question count HINT + minimum pad, never a truncation. Matching queries
    number nothing → fall back to numbered items in the CONTEXT."""
    q = re.findall(r"(?m)^\s*(\d+)[\.\)]", query)
    if q:
        return len(q)
    par = re.findall(r"\((\d+)\)", query)
    if par:
        return len(set(par))
    c = re.findall(r"(?m)^\s*(\d+)[\.\)]", context)
    if c:
        return len(c)
    return 1


def _clean_answer(s):
    s = re.sub(r"^\s*(?:\d+[\.\):]|[-*•])\s*", "", s).strip()
    s = re.sub(r"^(?:answer|ans|translation|result)s?\s*[:\-]\s*", "", s, flags=re.I).strip()
    return s.strip("\"'“”‘’` ").strip()


def parse_answers(text, min_count=1):
    """Full answer list from the model (count from the model, never truncated)."""
    seg = text.rsplit(ANSWER_MARKER, 1)[1] if ANSWER_MARKER in text else text
    numbered = {}
    for m in re.finditer(r"(?m)^\s*(\d+)[\.\)]\s*(.+?)\s*$", seg):
        numbered[int(m.group(1))] = _clean_answer(m.group(2))
    if numbered:
        answers = [numbered.get(i, "") for i in range(1, max(numbered) + 1)]
    else:
        lines = [ln.strip() for ln in seg.splitlines() if ln.strip()]
        comma_line = next((ln for ln in reversed(lines) if "," in ln), "")
        if comma_line:
            answers = [_clean_answer(x) for x in comma_line.split(",")]
        else:
            answers = [_clean_answer(ln) for ln in lines]
    answers = [a if a else "?" for a in answers]
    if len(answers) < min_count:
        answers += ["?"] * (min_count - len(answers))
    return answers if answers else ["?"]


def make_explanation(raw, row):
    """CSV-SAFE single-line explanation from the model's reasoning (the text before the
    answer block). ALL whitespace (incl. newlines) collapsed to single spaces — this is
    the whole point: no embedded newlines to corrupt the submission CSV."""
    head = raw.split(ANSWER_MARKER, 1)[0] if ANSWER_MARKER in raw else raw
    head = " ".join(head.split())            # collapse newlines/tabs/runs -> one line
    head = head[:300].strip()
    if head:
        return head
    t = (row.get("task_type") or "linguistic").replace("_", " ")
    return f"Inferred the {t} rule from the given examples and applied it to each item."


def _already_quantized(model_dir):
    cfg = os.path.join(model_dir, "config.json")
    try:
        with open(cfg, encoding="utf-8") as f:
            return "quantization_config" in json.load(f)
    except Exception:
        return False


def load_model():
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM

    tok = AutoTokenizer.from_pretrained(MODEL_DIR)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    if not torch.cuda.is_available():
        return tok, AutoModelForCausalLM.from_pretrained(
            MODEL_DIR, torch_dtype=torch.float32).eval()
    kwargs = dict(torch_dtype=torch.float16, device_map="auto")
    if _already_quantized(MODEL_DIR):
        pass
    elif QUANT == "4bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True)
    return tok, AutoModelForCausalLM.from_pretrained(MODEL_DIR, **kwargs).eval()


def generate_one(tok, model, messages):
    import torch
    dev = model.device if hasattr(model, "device") else "cpu"
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt").to(dev)
    with torch.no_grad():
        gen = model.generate(ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=False,
                             pad_token_id=tok.pad_token_id)
    return tok.decode(gen[0][ids.shape[-1]:], skip_special_tokens=True).strip()


def main():
    tok, model = load_model()
    with open(TEST_CSV, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))

    fout = open(OUT_CSV, "w", newline="", encoding="utf-8")
    writer = csv.DictWriter(fout, fieldnames=["id", "pred", "explanation"])
    writer.writeheader()
    fout.flush()

    for k, r in enumerate(rows):
        context = (r.get("context") or "").strip()
        query = (r.get("query") or "").strip()
        min_count = detect_count(context, query)
        try:
            raw = generate_one(tok, model, build_messages(r))
            answers = parse_answers(raw, min_count)
            explanation = make_explanation(raw, r)
        except Exception as e:
            print("row %s fallback: %r" % (r.get("id"), e), flush=True)
            answers = ["?"] * min_count
            explanation = "Answer derived from the patterns in the examples."
        writer.writerow({"id": r["id"],
                         "pred": json.dumps(answers, ensure_ascii=False),
                         "explanation": explanation})
        fout.flush()
        print("%d/%d done" % (k + 1, len(rows)), flush=True)

    fout.close()
    print("wrote %s (%d rows)" % (OUT_CSV, len(rows)), flush=True)


if __name__ == "__main__":
    main()