"""IOL-AI submission script — RUN 1: pipeline smoke test."""
import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import re
import json

MODEL_ID = "."   

SYSTEM = (
    "You solve International Linguistics Olympiad problems. Everything needed is in the "
    "problem itself. Answer every numbered item. Put each answer on its own line, in the "
    "same order the items appear, with no numbering and no extra commentary."
)


def count_items(query: str) -> int:
    line_items = re.findall(r"(?m)^\s*[\(\[]?(\d+)[\.\)\]]", query or "")
    if line_items:
        return len(line_items)
    paren_items = re.findall(r"\((\d+)\)", query or "")
    if paren_items:
        return len(set(paren_items))
    return 1


def _clean(line: str) -> str:
    line = re.sub(r"^\s*[\(\[]?\d+[\.\)\]]?\s*", "", line)
    return line.strip().strip('"').strip("'").strip()


def parse_answers(text: str, n: int):
    lines = [_clean(ln) for ln in (text or "").splitlines()]
    lines = [ln for ln in lines if ln]
    if len(lines) >= n:
        return lines[:n]
    return lines + [""] * (n - len(lines))


def main() -> None:
    import pandas as pd
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM

    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, torch_dtype=torch.float16, device_map="auto"
    ).eval()

    df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")
    out_rows = []
    for _, r in df.iterrows():
        n = count_items(r["query"])
        messages = [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},
        ]
        ids = tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt"
        ).to(model.device)
        with torch.no_grad():
            gen = model.generate(ids, max_new_tokens=512, do_sample=False)
        text = tok.decode(gen[0][ids.shape[-1]:], skip_special_tokens=True).strip()
        answers = parse_answers(text, n)
        out_rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})
        print(f"{len(out_rows)}/{len(df)} done  id={r['id']}  items={n}", flush=True)

    pd.DataFrame(out_rows).to_csv("submission.csv", index=False)
    print("wrote submission.csv", flush=True)


if __name__ == "__main__":
    main()