import os
import re
import json
import time
import unicodedata
from collections import Counter, defaultdict

T0 = time.time()
L1 = float(os.environ.get("IOL_TIME_LIMIT", "1800"))
S1 = float(os.environ.get("IOL_SAFETY", "150"))
D1 = T0 + L1 - S1

P1 = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
P2 = os.environ.get("IOL_OUT_CSV", "submission.csv")
M1 = os.environ.get("IOL_MODEL", ".")
E1 = os.environ.get("IOL_EXPLAIN", "1") == "1"

X1 = int(os.environ.get("IOL_MAXNEW", "512")) 
X2 = int(os.environ.get("IOL_MAXSAMPLES", "24"))
X3 = float(os.environ.get("IOL_TEMP", "0.5"))
X4 = int(os.environ.get("IOL_BATCH", "4"))

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

def lg(msg):
    print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)

def lf():
    return D1 - time.time()

_LN = re.compile(r"^[ \t]*(\d{1,3})[.)\]]", re.M)
_PN = re.compile(r"\((\d{1,3})\)")
_RG = re.compile(r"\(?(\d{1,3})\s*(?:[-–—]|to)\s*(\d{1,3})\)?")
_LL = re.compile(r"^[ \t]*([A-Z])[.)\]]\s", re.M)
_SP = re.compile(r"^\s*(?:\(?\d{1,3}\)?[.):\]]\s*|[-*•]\s+)")
_FC = re.compile(r"^```[a-zA-Z]*\s*$")
_CT = re.compile(
    r"^\s*(?:here (?:are|is)\b|answers?\s*:?\s*$|explanation\b|note\b|okay\b|"
    r"solution\b|reasoning\b|analysis\b|translations?\s*:?\s*$|the answers?\b|"
    r"let me\b|first,|so,|therefore\b|thus\b)", re.I)

def d1(q, t="", c=""):
    q = q or ""
    ln = [int(m) for m in _LN.findall(q)]
    pn = [int(m) for m in _PN.findall(q)]
    rn = 0
    for a, b in _RG.findall(q):
        a, b = int(a), int(b)
        if 0 < b - a < 60: rn = max(rn, b - a + 1)
    cand = max(len(set(ln)), len(set(pn)))
    if rn and cand and rn != cand: return cand
    cand = max(cand, len(set(_LL.findall(q))))
    n = max(rn, cand)
    if n > 1: return n
    lns = [l.strip() for l in q.splitlines() if l.strip()]
    if len(lns) > 1:
        h = lns[0]
        b = lns[1:] if h.endswith((":", ".")) else lns
        if b: return len(b)
    if c:
        cn = len(set(int(m) for m in _LN.findall(c)))
        if cn > 1: return cn
        cl = len(set(_LL.findall(c)))
        if cl > 1: return cl
    return max(n, 1)

def d2(q, n, t=""):
    if t.strip().lower() == "match_letters": return ["A"] * n
    q = q or ""
    out = []
    for ln in q.splitlines():
        s = ln.strip()
        if not s: continue
        m = re.match(r"^\(?(\d{1,3})\)?[.):\]]\s*(.+)$", s)
        if m: out.append(m.group(2).strip())
    if not out:
        lns = [l.strip() for l in q.splitlines() if l.strip()]
        if len(lns) > 1 and lns[0].endswith((":", ".")): out = lns[1:]
    out = [o.split("|")[0].strip() if "|" in o else o for o in out]
    out = [o for o in out if o]
    while len(out) < n: out.append(out[-1] if out else "?")
    return out[:n]

def d3(s):
    s = s.strip()
    s = _SP.sub("", s)
    s = s.strip().strip("`").strip()
    if len(s) >= 2 and s[0] == s[-1] and s[0] in "\"'“”": s = s[1:-1].strip()
    return s.strip()

def d4(items, n, fb=None):
    items = [i for i in items if i and i.strip()]
    if len(items) > n: items = items[-n:]
    while len(items) < n:
        if fb and len(items) < len(fb): items.append(fb[len(items)])
        else: items.append(items[-1] if items else "?")
    return items[:n]

def d5(text, n, fb=None):
    if not text: return list(fb[:n]) if fb else ["?"] * n
    m = None
    for m2 in re.finditer(r"(?:^|\n)\s*(?:final\s+)?answers?\s*:\s*\n?", text, re.I): m = m2
    body = text[m.end():] if m else text
    numbered, raw = [], []
    for ln in body.splitlines():
        if _FC.match(ln): continue
        mm = re.match(r"^\s*\(?(\d{1,3})\)?[.):\]]\s*(.+)$", ln.strip())
        if mm:
            val = d3(mm.group(2))
            if val and not _CT.match(val): numbered.append((int(mm.group(1)), val))
        c = d3(ln)
        if c and not _CT.match(c): raw.append(c)
    if len(numbered) >= n:
        by_label = {}
        for lab, val in numbered: by_label[lab] = val
        labs = sorted(by_label)
        if len(labs) >= n: return [by_label[l] for l in labs[:n]]
    return d4(raw, n, fb)

def d6(s):
    s = unicodedata.normalize("NFC", (s or "").strip().lower())
    s = _SP.sub("", s)
    s = re.sub(r"\s+", " ", s)
    return s.strip(" .!?;:,")

def d7(cands, anchor=None):
    cands = [c for c in cands if c and c.strip()]
    if anchor is None: anchor = cands[0] if cands else "?"
    if len(cands) < 3: return anchor
    groups = defaultdict(list)
    for c in cands: groups[d6(c)].append(c)
    anchor_support = len(groups.get(d6(anchor), []))
    best_key, best_n = None, 0
    for k, v in groups.items():
        if len(v) > best_n: best_key, best_n = k, len(v)
    if best_key is not None and best_n >= 3 and best_n > anchor_support:
        return Counter(groups[best_key]).most_common(1)[0][0]
    return anchor

def d8(path, ids, preds, explanations=None):
    import pandas as pd
    rows = []
    for i in ids:
        rec = {"id": i, "pred": json.dumps(preds[i], ensure_ascii=False)}
        if explanations is not None: rec["explanation"] = explanations.get(i, "")
        rows.append(rec)
    pd.DataFrame(rows).to_csv(path, index=False)

def main():
    import pandas as pd
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    df = pd.read_csv(P1, dtype=str).fillna("")
    I1 = [str(x) for x in df["id"].tolist()]
    N1 = [d1(r.get("query", ""), r.get("task_type", ""), r.get("context", "")) for _, r in df.iterrows()]
    total_items = sum(N1)
    lg(f"loaded {len(df)} problems, {total_items} items")

    S2 = {i: d2(r.get("query", ""), n, r.get("task_type", "")) for i, (_, r), n in zip(I1, df.iterrows(), N1)}

    R1 = {i: list(S2[i]) for i in I1}
    E2 = {i: "" for i in I1} if E1 else None
    d8(P2, I1, R1, E2)
    lg(f"wrote placeholder {P2} ({len(I1)} rows)")

    lg("loading tokenizer/model ...")
    tk = AutoTokenizer.from_pretrained(M1, trust_remote_code=True)
    if tk.pad_token is None: tk.pad_token = tk.eos_token
    tk.padding_side = "left"

    def _ld(dm):
        try:
            return AutoModelForCausalLM.from_pretrained(M1, torch_dtype=torch.float16, device_map=dm, trust_remote_code=True).eval()
        except TypeError:
            return AutoModelForCausalLM.from_pretrained(M1, dtype=torch.float16, device_map=dm, trust_remote_code=True).eval()

    try:
        ml = _ld({"": 0} if torch.cuda.is_available() else "auto")
    except Exception as e:
        lg(f"pinned load failed ({e}); falling back to auto")
        ml = _ld("auto")

    lg(f"model ready ({lf():.0f}s left)")

    P3 = []
    for _, r in df.iterrows():
        msgs = [
            {"role": "system", "content": "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."},
            {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"}
        ]
        P3.append(tk.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))

    B1 = X4
    
    class Deadline(StoppingCriteria):
        def __init__(self, stop_at): self.stop_at = stop_at
        def __call__(self, input_ids, scores, **kw): return time.time() > self.stop_at

    def d9(texts, max_new, sample, temp=0.7):
        nonlocal B1
        out = [""] * len(texts)
        order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
        i = 0
        while i < len(order):
            if lf() < 25: break
            idx = order[i:i + B1]
            chunk = [texts[j] for j in idx]
            try:
                enc = tk(chunk, return_tensors="pt", padding=True, truncation=True, max_length=6144).to(ml.device)
                kw = dict(max_new_tokens=max_new, pad_token_id=tk.pad_token_id, repetition_penalty=1.0, stopping_criteria=StoppingCriteriaList([Deadline(D1 - 10)]))
                if sample:
                    kw.update(do_sample=True, temperature=temp, top_p=0.95)
                else:
                    kw.update(do_sample=False)
                with torch.no_grad():
                    o = ml.generate(**enc, **kw)
                for k, j in enumerate(idx):
                    out[j] = tk.decode(o[k][enc["input_ids"].shape[1]:], skip_special_tokens=True)
                i += B1
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache()
                if B1 == 1: i += 1
                else: B1 = max(1, B1 // 2)
            except Exception:
                i += B1
        return out

    lg(f"Pass 1 (greedy) starting... budget: {X1} tokens/item")
    t = time.time()
    texts = d9(P3, max_new=X1, sample=False)
    c1 = time.time() - t
    
    V1 = {i: [] for i in I1}
    for i, n, txt in zip(I1, N1, texts):
        p1 = [d3(ln) for ln in (txt or "").splitlines() if ln.strip()]
        R1[i] = d4(p1, n, S2[i])
        V1[i].append(R1[i])
        
    d8(P2, I1, R1, E2)
    lg(f"Pass 1 done in {c1:.0f}s. Written to disk.")

    reserve = min(300.0, 0.25 * c1 + 60) if E1 else 30.0
    n_extra = 0
    while lf() - reserve > c1 * 1.25 and n_extra < X2:
        n_extra += 1
        lg(f"Self-consistency pass {n_extra} starting... ({lf():.0f}s left)")
        texts = d9(P3, max_new=X1, sample=True, temp=X3)
        for i, n, txt in zip(I1, N1, texts):
            if txt:
                V1[i].append(d5(txt, n, S2[i]))
        
        for i, n in zip(I1, N1):
            if len(V1[i]) >= 3:
                greedy = V1[i][0]
                voted = [d7([s[k] for s in V1[i] if k < len(s)], anchor=greedy[k] if k < len(greedy) else None) for k in range(n)]
                R1[i] = d4(voted, n, S2[i])
                
        d8(P2, I1, R1, E2)
        lg(f"Pass {n_extra+1} voted and written.")

    if E1 and lf() > 60:
        lg(f"Generating explanations ({lf():.0f}s left)...")
        ex_sys = "You explain International Linguistics Olympiad solutions to a human judge. State the key rules of the language: morphemes, word order, sound changes. Be concise (2-4 sentences)."
        ex_prompts = []
        for _, r in df.iterrows():
            i = str(r["id"])
            msgs = [
                {"role": "system", "content": ex_sys},
                {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}\n\nAnswers given:\n" + "\n".join(f"- {a}" for a in R1[i]) + "\n\nBriefly explain the linguistic rules behind these answers."}
            ]
            ex_prompts.append(tk.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True))
        
        ex_texts = d9(ex_prompts, max_new=200, sample=False)
        for i, e in zip(I1, ex_texts):
            e = re.sub(r"\s+", " ", (e or "").strip())
            if e: E2[i] = e[:1200]
        d8(P2, I1, R1, E2)
        lg("Explanations written.")

    bad = [i for i, n in zip(I1, N1) if len(R1[i]) != n or any(not str(x).strip() for x in R1[i])]
    if bad:
        lg(f"Repairing {len(bad)} malformed rows")
        for i, n in zip(I1, N1):
            R1[i] = d4([x for x in R1[i] if str(x).strip()], n, S2[i])
        d8(P2, I1, R1, E2)

    lg(f"DONE. {len(I1)} rows, {time.time() - T0:.0f}s elapsed.")

if __name__ == "__main__":
    main()