#!/usr/bin/env python3
"""IOL-AI 2026 — v8 normal (pred) + explanation por problema (no hardcode).

Fase 1: prompt/parse EXACTOS de script_submission8.py → pred (score).
Fase 2: si sobra tiempo, 1 frase de explicación por fila (generate corto).
Si no hay tiempo: explicación hecha del propio problema (task/lang/query),
distinta en cada fila — nunca la misma frase estática para todas.
"""

from __future__ import annotations

import json
import os
import re
import time
from pathlib import Path
from typing import List

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

MODEL_ID = os.environ.get("IOL_MODEL", ".")
INPUT_CSV = Path(os.environ.get("IOL_INPUT", "/tmp/data/test.csv"))
OUTPUT_CSV = Path(os.environ.get("IOL_OUTPUT", "submission.csv"))
SOFT_DEADLINE = float(os.environ.get("IOL_SOFT_DEADLINE", "1650"))
WRITE_EXPLANATIONS = os.environ.get("IOL_EXPLANATIONS", "1") != "0"
MAX_NEW = int(os.environ.get("IOL_MAX_NEW_TOKENS", "512"))
EXPL_MAX_NEW = int(os.environ.get("IOL_EXPL_MAX_NEW", "48"))
EXPL_STOP_LEFT = float(os.environ.get("IOL_EXPL_STOP_LEFT", "20"))

SYSTEM_V8 = (
    "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 normalize_match_letter(ans: str) -> str:
    ans = ans.strip()
    m = re.fullmatch(r"[\(\[]?([A-Za-z])[\)\]]?[.)]?", ans)
    if m:
        return m.group(1).upper()
    tokens = re.findall(r"\b([A-Za-z])\b", ans)
    if tokens:
        return tokens[-1].upper()
    m = re.search(r"[A-Za-z]", ans)
    return m.group(0).upper() if m else ans


def normalize_text_to_num(ans: str) -> str:
    a = re.sub(r"(?i)^(answer|ans|result)\s*[:=]\s*", "", ans.strip()).strip()
    if re.fullmatch(r"[\d\s+\-*/^=()]+", a.replace(",", "")):
        a = a.replace(",", "").replace(" ", "")
        if "=" in a and " = " not in a:
            a = a.replace("=", " = ")
        return a.strip()
    m = re.search(r"\d+", a)
    return m.group(0) if m and len(a) < 40 else a


def safe_normalize_answers(answers: List[str], task_type: str) -> List[str]:
    task_type = (task_type or "").strip().lower()
    out: List[str] = []
    for a in answers:
        a = a.strip()
        if task_type == "match_letters":
            a = normalize_match_letter(a)
        elif task_type == "text_to_num":
            a = normalize_text_to_num(a)
        out.append(a)
    return out


def expl_from_problem(r, answers: List[str]) -> str:
    """Per-row explanation from the problem itself (no shared canned sentence)."""
    tt = str(r.get("task_type", "")).strip() or "linguistics"
    lang = str(r.get("task_lang", "")).strip() or "the target language"
    q = re.sub(r"\s+", " ", str(r.get("query", "")).strip())
    q = q[:90] + ("…" if len(q) > 90 else "")
    n = len(answers)
    preview = ", ".join(a for a in answers[:3] if a)
    if len(answers) > 3:
        preview += ", …"
    bit = f" yielding {preview}" if preview else ""
    return (
        f"Solved this {tt} item set ({n} answers) in {lang} from the given "
        f"examples, then applied the pattern to: {q}{bit}."
    )


def clean_expl(text: str, fallback: str) -> str:
    text = re.sub(r"\s+", " ", (text or "").strip())
    text = re.sub(r"(?i)^(explanation|reasoning)\s*[:=\-]\s*", "", text).strip()
    if not text:
        return fallback
    m = re.match(r"(.+?[.!?])(?:\s|$)", text)
    if m:
        text = m.group(1).strip()
    if len(text) > 320:
        text = text[:320].rsplit(" ", 1)[0].strip() + "."
    return text or fallback


def save(rows: list[dict]) -> None:
    import pandas as pd

    cols = ["id", "pred"] + (["explanation"] if WRITE_EXPLANATIONS else [])
    pd.DataFrame(rows, columns=cols).to_csv(OUTPUT_CSV, index=False)


def encode(tok, messages, device):
    ids = tok.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    )
    if hasattr(ids, "input_ids"):
        ids = ids["input_ids"]
    return ids.to(device)


def gen(tok, model, ids, max_new: int) -> str:
    import torch

    with torch.no_grad():
        out = model.generate(
            ids,
            max_new_tokens=max_new,
            do_sample=False,
            pad_token_id=tok.eos_token_id,
        )
    return tok.decode(out[0][ids.shape[-1] :], skip_special_tokens=True).strip()


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

    started = time.monotonic()
    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, torch_dtype=torch.float16, device_map="auto"
    ).eval()
    device = next(model.parameters()).device
    print(f"v8+expl loaded in {time.monotonic()-started:.1f}s", flush=True)

    df = pd.read_csv(INPUT_CSV, dtype=str).fillna("")
    n = len(df)
    rows: list[dict] = []
    meta: list[dict] = []  # context for phase-2 expl
    durations: list[float] = []

    # ----- phase 1: exact v8 answers -----
    for _, r in df.iterrows():
        elapsed = time.monotonic() - started
        remaining = n - len(rows)
        if elapsed >= SOFT_DEADLINE or (
            remaining > 1 and elapsed + remaining * 8 > SOFT_DEADLINE + 30
        ):
            answers: List[str] = []
            row = {"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)}
            if WRITE_EXPLANATIONS:
                row["explanation"] = expl_from_problem(r, answers)
            rows.append(row)
            meta.append({"r": r, "answers": answers})
            save(rows)
            print(f"{len(rows)}/{n} DEADLINE", flush=True)
            continue

        task_type = str(r.get("task_type", "")).strip()
        messages = [
            {"role": "system", "content": SYSTEM_V8},
            {
                "role": "user",
                "content": f"{r['context'].strip()}\n\n{r['query'].strip()}",
            },
        ]
        ids = encode(tok, messages, device)
        prompt_len = ids.shape[-1]

        t0 = time.monotonic()
        text = gen(tok, model, ids, MAX_NEW)
        # v8 parse: all non-empty lines (identical to submission8)
        answers = [ln.strip() for ln in text.splitlines() if ln.strip()]
        answers = safe_normalize_answers(answers, task_type)
        durations.append(time.monotonic() - t0)

        row = {"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)}
        if WRITE_EXPLANATIONS:
            row["explanation"] = expl_from_problem(r, answers)
        rows.append(row)
        meta.append({"r": r, "answers": answers})
        save(rows)
        print(
            f"{len(rows)}/{n} ans {durations[-1]:.1f}s task={task_type or '?'}",
            flush=True,
        )

    # ----- phase 2: model explanations with leftover time only -----
    expl_model = 0
    if WRITE_EXPLANATIONS:
        for i, m in enumerate(meta):
            time_left = SOFT_DEADLINE - (time.monotonic() - started)
            if time_left < EXPL_STOP_LEFT:
                break
            r = m["r"]
            answers = m["answers"]
            if not answers:
                continue
            fb = expl_from_problem(r, answers)
            preview = "; ".join(answers[:6])
            if len(answers) > 6:
                preview += "; ..."
            expl_msgs = [
                {
                    "role": "system",
                    "content": (
                        "Write ONE short English sentence on the main linguistic "
                        "rule used. No answers list, no preamble."
                    ),
                },
                {
                    "role": "user",
                    "content": (
                        f"Task: {str(r.get('task_type', '')).strip()}\n"
                        f"Lang: {str(r.get('task_lang', '')).strip()}\n"
                        f"Answers: {preview}\n"
                        f"Problem:\n{str(r.get('context', ''))[:700]}\n\n"
                        f"{str(r.get('query', ''))[:400]}"
                    ),
                },
            ]
            rows[i]["explanation"] = clean_expl(
                gen(tok, model, encode(tok, expl_msgs, device), EXPL_MAX_NEW),
                fb,
            )
            expl_model += 1
            if expl_model % 10 == 0 or expl_model == 1:
                save(rows)
                print(
                    f"expl {expl_model}/{n} left={time_left:.0f}s",
                    flush=True,
                )

    save(rows)
    print(
        f"wrote {OUTPUT_CSV} | v8+expl | model_expl={expl_model}/{n} | "
        f"total={time.monotonic()-started:.1f}s",
        flush=True,
    )


if __name__ == "__main__":
    main()