"""IOL-AI 2026 submission entrypoint — PLAIN baseline harness.

Faithful replica of the official How_to_submit.md baseline (the config the
organizers posted as "Baseline-Qwen2.5-14B-AWQ", public score 0.1227): a minimal
system prompt, the raw context+query, greedy decoding, 512 new tokens, and a
naive newline split of the output with NO forced item count. The only additions
are timeout-safe incremental writes (an all-blank submission.csv up front + a
periodic flush), which never change the output of a run that finishes in time —
they only guarantee a valid partial file if we were ever killed.

Runs on the platform's T4 (16 GB), offline, within the 30-minute limit. Reads
/tmp/data/test.csv, writes submission.csv (id,pred) to the working directory.
Model weights (Qwen2.5-14B-Instruct-AWQ) ship in this repo and load from ".".
"""

import os
import time

# Must be set BEFORE importing transformers: no network at run time.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import json

import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "."  # weights are shipped in this repo (offline)
TEST_CSV = "/tmp/data/test.csv"
OUT_CSV = "submission.csv"

MAX_NEW_TOKENS = 512          # same as the official baseline
SAFETY_S = 25 * 60            # stop generating past this; buffer before the 30-min cap
WRITE_EVERY = 3               # flush submission.csv every N problems (timeout safety)
_START = time.monotonic()

# The official baseline's exact system prompt.
SYSTEM = (
    "You solve International Linguistics Olympiad problems. Answer every numbered "
    "item. Put each answer on its own line, in order, with no numbering and no "
    "extra text."
)


def elapsed():
    return time.monotonic() - _START


def load_model():
    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,  # T4 has no bfloat16
        device_map="auto",
    ).eval()
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    return tok, model


def main():
    df = pd.read_csv(TEST_CSV, dtype=str).fillna("")
    total = len(df)
    print(f"loaded {total} problems from {TEST_CSV}", flush=True)

    # Write a COMPLETE all-blank submission.csv up front so a valid file with
    # every id always exists, even if we are killed during load or generation.
    records = [
        {"id": df.iloc[i]["id"], "pred": json.dumps([], ensure_ascii=False)}
        for i in range(total)
    ]

    def flush():
        pd.DataFrame(records, columns=["id", "pred"]).to_csv(OUT_CSV, index=False)

    flush()
    print(f"wrote blank {OUT_CSV} ({total} rows) at {elapsed():.0f}s", flush=True)

    tok, model = load_model()
    print(f"model loaded at {elapsed():.0f}s", flush=True)

    for i in range(total):
        if elapsed() > SAFETY_S:
            print(f"[warn] time budget hit at problem {i}; leaving the rest blank", flush=True)
            break
        r = df.iloc[i]
        messages = [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
        ]
        try:
            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,
                    do_sample=False, pad_token_id=tok.pad_token_id,
                )
            text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
            answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
        except Exception as exc:
            print(f"[warn] id={r['id']} failed: {exc!r}", flush=True)
            answers = []
        records[i]["pred"] = json.dumps(answers, ensure_ascii=False)
        if (i + 1) % WRITE_EVERY == 0:
            flush()
        print(f"{i + 1}/{total} done at {elapsed():.0f}s", flush=True)

    flush()
    print(f"final {OUT_CSV} ({total} rows) at {elapsed():.0f}s", flush=True)


if __name__ == "__main__":
    main()