"""
IOL-AI Challenge 2026 — submission script (OFFLINE / Mode B).

Runtime facts (Space Submission tab):
  * T4 medium, 16 GB VRAM, Python 3.10, 30-min wall clock.
  * NO internet: cannot pip install or download anything. Model weights must be
    committed into THIS repo (the working dir) and loaded from ".". Only the
    pre-installed libraries/versions are available (torch 2.4.0, transformers
    4.44.1, accelerate 0.34.2, bitsandbytes 0.43.3, autoawq 0.2.7, pandas 2.2.2,
    numpy 2.1.3, ...). Do NOT pin different majors of torch/transformers/numpy.
  * Read hidden test set from /tmp/data/test.csv; write submission.csv here.
  * pred = JSON list, one entry per numbered item, in query order.

Ship the model in the repo with build_repo.py. This script loads it from "." with
bitsandbytes 4-bit by default (or auto-detected AWQ) so it fits 16 GB. T4 has no
bf16 -> use float16.

v3 adds: (1) task_type-conditioned prompts with a few-shot example per type,
(2) BATCHED generation to use the 30-min budget efficiently, (3) self-consistency
— SAMPLES sampled decodes per problem, majority-voted per item. Output is written
incrementally so a timeout still yields a valid partial submission.csv. Tunable via
IOL_SAMPLES / IOL_BATCH_SIZE / IOL_TEMPERATURE / IOL_TOP_P / IOL_MAX_NEW_TOKENS.

Local dev: set IOL_TEST_CSV to a mock file. Quantization auto-disables if there's
no CUDA so the plumbing can be exercised on CPU with a tiny model.
"""

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", ".")        # weights live in the repo
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", "1024"))
# "4bit" (bitsandbytes), "awq" (weights already AWQ-quantized), or "fp16".
QUANT = os.environ.get("IOL_QUANT", "4bit")

# --- Tier 1 (throughput) + Tier 2 (self-consistency) knobs ------------------
# SAMPLES>1 => draw that many sampled decodes per problem and MAJORITY-VOTE the
# answer per item (cancels one-off reasoning slips). SAMPLES=1 => greedy, no vote.
# BATCH_SIZE caps how many sequences share one generate() call; we pack whole
# problems (each expanded to SAMPLES decodes) into a batch so the T4 stays busy.
# Watch the 30-min wall: total decodes ~= n_rows * SAMPLES. Dial SAMPLES down (or
# MAX_NEW_TOKENS) if the Space logs show you nearing the cap.
SAMPLES = int(os.environ.get("IOL_SAMPLES", "5"))
# Keep one problem's worth of samples per generate() call by default: 5 sequences
# of KV cache sits safely under the T4's ~6 GB free after AWQ weights. Raise only
# if the Space logs show VRAM headroom.
BATCH_SIZE = int(os.environ.get("IOL_BATCH_SIZE", "5"))
TEMPERATURE = float(os.environ.get("IOL_TEMPERATURE", "0.7"))
TOP_P = float(os.environ.get("IOL_TOP_P", "0.9"))
# Safety valve for the 30-min wall: once this many seconds have elapsed, finish
# the remaining rows with ONE fast greedy decode instead of SAMPLES sampled ones,
# so every row still gets an answer rather than timing out mid-set.
TIME_BUDGET_S = float(os.environ.get("IOL_TIME_BUDGET_S", "1620"))  # 27 min

ANSWER_MARKER = "###ANSWERS###"

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 "
    "grammar and vocabulary using ONLY the data and hints in the problem. "
    "Work through it briefly, then give your final answers.\n\n"
    "You MUST end your reply with the answers in EXACTLY this format and write "
    "nothing after it:\n"
    f"{ANSWER_MARKER}\n"
    "1. <answer to item 1>\n"
    "2. <answer to item 2>\n"
    "...(one numbered line per item, in order)\n\n"
    "Each answer must contain ONLY the requested form and nothing else: a single "
    "word, phrase, number, or letter. Do NOT restate the question, explain, or add "
    "commentary after the answer. For letter-matching items give just the letter "
    "(e.g. B). For number items give the digits or written-out number as asked. "
    "Give exactly one answer for every numbered item — never leave one blank."
)

# Per task_type: a one-line output constraint + a tiny worked example (shows the
# expected reasoning depth AND the exact answer shape). The CSV labels every row
# with `task_type`, so we tailor the instruction instead of one generic prompt.
TASK_GUIDE = {
    "translation": (
        "This is a TRANSLATION item: output ONLY the target-language word/phrase, "
        "no gloss or explanation.",
        "Data: kal = stay, kalar = they stay; git = go.\n"
        "Query: 1. they go\n"
        "Reason: the 'they' ending is -ar (kal->kalar), so git -> gitar... check "
        "vowel harmony with i -> giter.\n"
        f"{ANSWER_MARKER}\n1. giterler",
    ),
    "text_to_num": (
        "This is a TEXT->NUMBER item: output ONLY digits (e.g. 42).",
        "Data: dua=2, puluh=10, duapuluh=20, duapuluh lima=25, lima=5.\n"
        "Query: 1. limapuluh dua\n"
        "Reason: lima(5) before puluh -> 5*10=50, dua(2) after adds 2 -> 52.\n"
        f"{ANSWER_MARKER}\n1. 52",
    ),
    "num_to_text": (
        "This is a NUMBER->TEXT item: output ONLY the number written in the target "
        "language's words.",
        "Data: 2=dua, 10=puluh, 20=duapuluh, 5=lima.\n"
        "Query: 1. 25\n"
        "Reason: 25 = 2*10 + 5 = duapuluh + lima.\n"
        f"{ANSWER_MARKER}\n1. duapuluh lima",
    ),
    "match_letters": (
        "This is a MATCHING item: output ONLY the single option letter (A, B, C, ...).",
        "Data: root nimu='see'; prefix ka-='I', suffix -ka='they'.\n"
        "Forms: A. nimu  B. nimuka  C. kanimu\n"
        "Query: 1. I see\n"
        "Reason: 'I' is prefix ka- -> kanimu = form C.\n"
        f"{ANSWER_MARKER}\n1. C",
    ),
}


def build_messages(row):
    """Chat messages for one problem, tailored to its task_type with a few-shot."""
    context = (row.get("context") or "").strip()
    query = (row.get("query") or "").strip()
    ttype = (row.get("task_type") or "").strip().lower()

    system = SYSTEM_PROMPT
    guide = TASK_GUIDE.get(ttype)
    if guide:
        instruction, example = guide
        system = system + "\n\n" + instruction + "\n\nWorked example:\n" + example

    return [
        {"role": "system", "content": system},
        {"role": "user", "content": context + "\n\n" + query},
    ]


def count_items(query):
    """Number of numbered items in a query, e.g. '17. .. 18. ..' -> 2."""
    nums = re.findall(r"(?m)^\s*(\d+)[\.\)]", query)
    return len(nums) if nums else 1


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


def _numbered_map(segment, n_items):
    """Collect 'n. text' / 'n) text' lines into {index: answer}. Answers are the
    reliable anchor: even if reasoning is interleaved, the trailing numbered list
    is what we want, so a later line for the same index overwrites an earlier one."""
    out = {}
    for m in re.finditer(r"(?m)^\s*(\d+)[\.\)]\s*(.+?)\s*$", segment):
        idx = int(m.group(1))
        if 1 <= idx <= n_items:
            out[idx] = _clean_answer(m.group(2))
    return out


def parse_answers(text, n_items):
    """Extract exactly n_items answers. Prefer the marked block; anchor on the
    numbered list; fall back to the LAST n non-empty lines (answers come last)."""
    seg = text.rsplit(ANSWER_MARKER, 1)[1] if ANSWER_MARKER in text else text

    numbered = _numbered_map(seg, n_items)
    if len(numbered) >= n_items or (numbered and ANSWER_MARKER in text):
        answers = [numbered.get(i, "") for i in range(1, n_items + 1)]
    else:
        # No usable numbered list: take the last n_items non-empty lines.
        lines = [_clean_answer(ln) for ln in seg.splitlines() if ln.strip()]
        lines = [ln for ln in lines if ln]
        answers = lines[-n_items:] if len(lines) >= n_items else lines

    # Guarantee exactly n_items, never blank (fall back to last good answer).
    last_good = next((a for a in reversed(answers) if a), "")
    answers = [a if a else last_good for a in answers]
    if len(answers) < n_items:
        answers += [last_good] * (n_items - len(answers))
    return answers[:n_items]


def _norm(s):
    """Mirror the official scorer's normalization so voting groups answers the
    same way the metric will (ignore case, surrounding quotes, one trailing dot)."""
    s = " ".join((s or "").strip().split())
    s = s.strip("\"'“”‘’")
    if s.endswith("."):
        s = s[:-1]
    return s.strip().casefold()


def vote_answers(sample_texts, n_items):
    """Self-consistency: parse each sampled decode, then per item pick the answer
    whose NORMALIZED form is most common across samples; return its surface form.
    Ties fall to the earliest-seen sample (insertion order in Counter)."""
    from collections import Counter

    counts = [Counter() for _ in range(n_items)]
    surface = [dict() for _ in range(n_items)]  # normalized -> first surface seen
    for text in sample_texts:
        for i, ans in enumerate(parse_answers(text, n_items)):
            key = _norm(ans)
            if not key:
                continue
            counts[i][key] += 1
            surface[i].setdefault(key, ans)

    out = []
    for i in range(n_items):
        if counts[i]:
            best = counts[i].most_common(1)[0][0]
            out.append(surface[i][best])
        else:
            out.append("")
    return out


def _already_quantized(model_dir):
    """True if the shipped weights are pre-quantized (e.g. AWQ) — then transformers
    auto-detects the config and we must NOT stack bitsandbytes on top."""
    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)
    # Decoder-only batched generation needs LEFT padding so every prompt's
    # continuation starts at the same column; fall back to eos as pad if unset.
    tok.padding_side = "left"
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    if not torch.cuda.is_available():
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_DIR, torch_dtype=torch.float32).eval()   # CPU dev fallback
        return tok, model

    kwargs = dict(torch_dtype=torch.float16, device_map="auto")  # T4 has no bf16
    if _already_quantized(MODEL_DIR):
        pass  # AWQ/pre-quant: transformers reads quantization_config from config.json
    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,
        )
    model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, **kwargs).eval()
    return tok, model


def generate_texts(tok, model, prompts, do_sample):
    """Batched decode: prompt strings -> generated continuations (one per prompt).
    Left-padded so we can slice the new tokens at a single shared offset."""
    import torch

    dev = model.device if hasattr(model, "device") else "cpu"
    enc = tok(prompts, return_tensors="pt", padding=True).to(dev)
    gkw = dict(max_new_tokens=MAX_NEW_TOKENS, pad_token_id=tok.pad_token_id)
    if do_sample:
        gkw.update(do_sample=True, temperature=TEMPERATURE, top_p=TOP_P)
    else:
        gkw.update(do_sample=False)
    with torch.no_grad():
        gen = model.generate(**enc, **gkw)
    new = gen[:, enc["input_ids"].shape[1]:]            # left pad => shared offset
    return [t.strip() for t in tok.batch_decode(new, skip_special_tokens=True)]


def main():
    import time
    tok, model = load_model()

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

    # Write incrementally so a 30-min timeout still leaves a valid partial file.
    fout = open(OUT_CSV, "w", newline="", encoding="utf-8")
    writer = csv.DictWriter(fout, fieldnames=["id", "pred"])
    writer.writeheader()
    fout.flush()

    start_t = time.time()
    rows_per_batch = max(1, BATCH_SIZE // max(1, SAMPLES))
    done = 0
    i = 0
    while i < len(rows):
        # Time guard: once past budget, drop to 1 fast greedy decode per row so
        # the remaining rows still get answered before the hard 30-min cut.
        over_budget = (time.time() - start_t) > TIME_BUDGET_S
        n_samp = 1 if (over_budget or SAMPLES <= 1) else SAMPLES
        step = rows_per_batch if n_samp > 1 else max(1, BATCH_SIZE)
        chunk = rows[i:i + step]

        prompts, meta = [], []                          # meta: n_items per row
        for r in chunk:
            prompt = tok.apply_chat_template(
                build_messages(r), tokenize=False, add_generation_prompt=True)
            prompts.extend([prompt] * n_samp)           # replicate for voting
            meta.append(count_items((r.get("query") or "").strip()))

        texts = generate_texts(tok, model, prompts, do_sample=(n_samp > 1))

        for j, r in enumerate(chunk):
            samples = texts[j * n_samp:(j + 1) * n_samp]
            answers = (vote_answers(samples, meta[j]) if n_samp > 1
                       else parse_answers(samples[0], meta[j]))
            writer.writerow({"id": r["id"],
                             "pred": json.dumps(answers, ensure_ascii=False)})
            done += 1
        fout.flush()                                    # survive a hard timeout
        print("%d/%d done%s" % (done, len(rows),
                                " [time-guard: greedy]" if n_samp == 1 and SAMPLES > 1 else ""),
              flush=True)
        i += step

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


if __name__ == "__main__":
    main()