"""IOL-AI 2026 — v13: BETTER-MODEL experiment (DeepSeek-R1-Distill-Qwen-14B-AWQ).

A reasoning ("thinking") model instead of Qwen2.5-14B-Instruct. It loads natively on
transformers 4.44.1 (model_type=qwen2, AWQ 4-bit) — no wheels hack. It emits a long
<think>…</think> chain, then the answer; the model's own reasoning REPLACES the
self-consistency/refine pipeline, so we do ONE sampled decode per row.

Same proven scaffolding as v12: count-fix (answers from model/context), CSV-SAFE
single-line explanation, incremental writes. Plus R1 specifics:
  * DeepSeek guidance: no system prompt, temperature ~0.6, all instructions in the user turn.
  * parse the answer AFTER </think>; fall back to the whole text if thinking got truncated.
  * a STRICT adaptive time-guard (per-row max_time + soft/hard wall phases) because R1's
    long generations are the real risk on a T4 (this is likely why a thinking model scored
    EM 0 before). If a row runs out of time it degrades to a placeholder — the file stays whole.
"""

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

import re
import csv
import json
import time

SCRIPT_START = time.time()

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", "2048"))  # room to think + answer
QUANT = os.environ.get("IOL_QUANT", "4bit")
TEMPERATURE = float(os.environ.get("IOL_TEMPERATURE", "0.6"))       # DeepSeek R1 recommendation
TOP_P = float(os.environ.get("IOL_TOP_P", "0.95"))
MIN_DECODE_S = float(os.environ.get("IOL_MIN_DECODE_S", "25"))
SOFT_BUDGET_S = float(os.environ.get("IOL_SOFT_BUDGET_S", "1560"))  # 26.0m
HARD_BUDGET_S = float(os.environ.get("IOL_HARD_BUDGET_S", "1710"))  # 28.5m

ANSWER_MARKER = "###ANSWERS###"
THINK_END = "</think>"

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


def build_messages(row):
    """R1: single user turn (no system prompt), instructions first, then the problem."""
    context = (row.get("context") or "").strip()
    query = (row.get("query") or "").strip()
    ttype = (row.get("task_type") or "").strip().lower()
    hint = TASK_HINT.get(ttype, "")
    user = (
        "Solve this International Linguistics Olympiad problem using ONLY the data given "
        "(no outside knowledge of any language). A problem may have MANY sub-questions — "
        "answer EVERY one, in order; if the query is not numbered (e.g. matching), give one "
        "answer for EACH item in the data. Keep your reasoning concise. "
        + (hint + " " if hint else "")
        + "Copy exact characters and diacritics from the data. End your response with the "
        "answers in EXACTLY this format and nothing after it:\n"
        f"{ANSWER_MARKER}\n1. <answer 1>\n2. <answer 2>\n(one numbered line per sub-question)\n\n"
        "PROBLEM:\n" + context + "\n\n" + query
    )
    return [{"role": "user", "content": user}]


def detect_count(context, query):
    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 _after_think(text):
    """The answer lives AFTER </think>. If thinking was truncated (no close tag), use the
    whole text as a best-effort fallback."""
    return text.rsplit(THINK_END, 1)[1] if THINK_END in text else text


def parse_answers(text, min_count=1):
    seg = _after_think(text)
    seg = seg.rsplit(ANSWER_MARKER, 1)[1] if ANSWER_MARKER in seg else seg
    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(text, row):
    """CSV-SAFE single-line summary: prefer the post-think final text (minus the answer
    block); else the tail of the reasoning; else a generic line. Newlines collapsed."""
    if THINK_END in text:
        think, after = text.split(THINK_END, 1)
        src = after.split(ANSWER_MARKER, 1)[0].strip() or think[-400:]
    else:
        src = text.split(ANSWER_MARKER, 1)[0][-400:]
    src = " ".join(src.split())[:300].strip()
    if src:
        return src
    t = (row.get("task_type") or "linguistic").replace("_", " ")
    return f"Reasoned from the given examples to infer the {t} rule for 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, max_time):
    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)
    gkw = dict(max_new_tokens=MAX_NEW_TOKENS, do_sample=True,
               temperature=TEMPERATURE, top_p=TOP_P, pad_token_id=tok.pad_token_id)
    if max_time and max_time > 0:
        gkw["max_time"] = float(max_time)
    with torch.no_grad():
        gen = model.generate(ids, **gkw)
    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()

    n = len(rows)
    for k, r in enumerate(rows):
        context = (r.get("context") or "").strip()
        query = (r.get("query") or "").strip()
        min_count = detect_count(context, query)
        elapsed = time.time() - SCRIPT_START
        rows_left = n - k
        # per-row share of the remaining soft budget bounds this decode's wall time
        row_budget = max(0.0, SOFT_BUDGET_S - elapsed) / max(1, rows_left)
        try:
            if elapsed > HARD_BUDGET_S:
                answers = ["?"] * min_count
                explanation = make_explanation("", r)
                mode = "placeholder"
            else:
                raw = generate_one(tok, model, build_messages(r),
                                   max_time=max(MIN_DECODE_S, row_budget))
                answers = parse_answers(raw, min_count)
                explanation = make_explanation(raw, r)
                mode = "think" if THINK_END in raw else "no-close"
        except Exception as e:
            print("row %s fallback: %r" % (r.get("id"), e), flush=True)
            answers = ["?"] * min_count
            explanation = make_explanation("", r)
            mode = "error"
        writer.writerow({"id": r["id"],
                         "pred": json.dumps(answers, ensure_ascii=False),
                         "explanation": explanation})
        fout.flush()
        print("%d/%d [%s] elapsed=%ds" % (k + 1, n, mode, time.time() - SCRIPT_START),
              flush=True)

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


if __name__ == "__main__":
    main()