"""IOL-AI 2026 submission entrypoint.

Runs on the platform's T4 (16 GB), offline, within a 30-minute limit. Reads the
hidden test set from /tmp/data/test.csv and writes submission.csv (id,pred) to
the working directory. Model weights ship inside this repo and load from ".".

Baseline model: Qwen/Qwen2.5-1.5B-Instruct. The score lever here is the harness
(task-aware prompts + robust parsing) in iol_harness.py, not the model size.
"""

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

import iol_harness as H

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

TIME_LIMIT_S = 30 * 60
SAFETY_S = 25 * 60  # stop generating past this; leaves buffer for a final gen+write
WRITE_EVERY = 3     # flush submission.csv every N problems (crash/timeout safety)
_START = time.monotonic()

# Optional per-repo harness config (reasoning-model repos ship one to turn on CoT).
CFG = {}
if os.path.exists("harness_cfg.json"):
    try:
        CFG = json.load(open("harness_cfg.json"))
    except Exception:
        CFG = {}
COT = bool(CFG.get("cot", False))
MAX_CAP = int(CFG.get("max_new_cap", 1024))
MAX_PASSES = int(CFG.get("sc_max_passes", 1))   # >1 turns on self-consistency
SC_TEMP = float(CFG.get("sc_temperature", 0.7))
FEWSHOT = bool(CFG.get("fewshot", False))        # prepend synthetic worked example


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


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


def generate(tok, model, messages, max_new_tokens, sample=False, temperature=0.7):
    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)
    kwargs = dict(max_new_tokens=max_new_tokens, pad_token_id=tok.pad_token_id)
    if sample:
        kwargs.update(do_sample=True, temperature=temperature, top_p=0.95)
    else:
        kwargs.update(do_sample=False, num_beams=1)  # greedy
    with torch.no_grad():
        out = model.generate(inputs, **kwargs)
    return tok.decode(out[0][inputs.shape[-1] :], skip_special_tokens=True).strip()


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

    # Pre-count items (no model needed) and write a COMPLETE all-blank
    # submission.csv up front. A valid file with every id then always exists —
    # even if we are killed during model load or a long generation — so a
    # timeout yields partial credit instead of total failure.
    ns = [
        H.count_items(r.get("context", ""), r.get("query", ""), r.get("task_type", ""))
        for _, r in df.iterrows()
    ]
    records = [
        {"id": df.iloc[i]["id"], "pred": json.dumps([""] * ns[i], 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)

    # Precompute per-problem prompt + token budget once.
    prompts, budgets = [], []
    for i in range(total):
        row = df.iloc[i]
        prompts.append(H.build_messages(
            row.get("context", ""), row.get("query", ""), row.get("task_type", ""), ns[i],
            cot=COT, fewshot=FEWSHOT))
        budgets.append(H.max_new_tokens_for(row.get("task_type", ""), ns[i], cot=COT, cap=MAX_CAP))

    # Self-consistency: pass 0 is greedy (a floor); passes 1+ are sampled. After
    # each pass we majority-vote across all passes and save — so we never do worse
    # than greedy, and extra time only helps. Time-budgeted to the 30-min limit.
    passes = []  # passes[k][i] = answer list for problem i in pass k
    est_pass = 0.0
    for k in range(max(1, MAX_PASSES)):
        if k > 0 and elapsed() + est_pass > SAFETY_S:
            print(f"[info] stopping before pass {k+1}: not enough time budget", flush=True)
            break
        t0 = elapsed()
        torch.manual_seed(1000 + k)  # reproducible sampling, diverse across passes
        this = []
        for i in range(total):
            if elapsed() > SAFETY_S:  # ran out mid-pass: blank the rest of this pass
                this += [[""] * ns[j] for j in range(i, total)]
                print(f"[warn] time budget hit mid-pass {k+1} at problem {i}", flush=True)
                break
            try:
                text = generate(tok, model, prompts[i], budgets[i],
                                sample=(k > 0), temperature=SC_TEMP)
                this.append(H.parse_answers(text, ns[i]))
            except Exception as exc:
                print(f"[warn] pass {k+1} id={df.iloc[i]['id']} failed: {exc!r}", flush=True)
                this.append([""] * ns[i])
        passes.append(this)
        est_pass = max(est_pass, elapsed() - t0)

        # vote across all passes so far and save
        for i in range(total):
            cands = [passes[p][i] for p in range(len(passes)) if i < len(passes[p])]
            records[i]["pred"] = json.dumps(H.majority_vote(cands, ns[i]), ensure_ascii=False)
        flush()
        print(f"pass {k+1}/{MAX_PASSES} done, voted+saved ({elapsed():.0f}s, ~{est_pass:.0f}s/pass)", flush=True)

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


if __name__ == "__main__":
    main()