#!/usr/bin/env python
import os, re, json, time, unicodedata
from collections import Counter, defaultdict

T0 = time.time()
TIME_LIMIT = float(os.environ.get("IOL_TIME_LIMIT", "1800"))
SAFETY     = float(os.environ.get("IOL_SAFETY", "150"))
DEADLINE   = T0 + TIME_LIMIT - SAFETY
TEST_CSV   = os.environ.get("IOL_TEST_CSV", "/tmp/data/test.csv")
OUT_CSV    = os.environ.get("IOL_OUT_CSV", "submission.csv")
MODEL_ID   = os.environ.get("IOL_MODEL", ".")
WANT_EXPL  = os.environ.get("IOL_EXPLAIN", "1") == "1"
TOK_PER_S  = float(os.environ.get("IOL_TOKS", "30"))
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")

MAX_NEW    = int(os.environ.get("IOL_MAXNEW", "600"))    # was 900 — faster passes, answers are short
MAX_SAMPLES= int(os.environ.get("IOL_MAXSAMPLES", "40")) # was 12/24 — more votes
BATCH_SIZE = int(os.environ.get("IOL_BATCH", "4"))       # REVERT from 8 (this caused 0.1485)
SAMPLE_TEMP= float(os.environ.get("IOL_TEMP", "0.5"))    # keep — value that gave 0.2039
COT        = os.environ.get("IOL_COT", "0") == "1"       # keep off

def log(m): print(f"[{time.time()-T0:7.1f}s] {m}", flush=True)
def left(): return DEADLINE - time.time()

# ---- item count + source fallback ------------------------------------------
_LINE_NUM=re.compile(r"^[ \t]*(\d{1,3})[.)\]]",re.M); _PAREN_NUM=re.compile(r"\((\d{1,3})\)")
_RANGE=re.compile(r"\(?(\d{1,3})\s*(?:[-–—]|to)\s*(\d{1,3})\)?"); _LINE_LET=re.compile(r"^[ \t]*([A-Z])[.)\]]\s",re.M)
def detect_n_items(query, task_type="", context=""):
    q=query or ""
    line=[int(m) for m in _LINE_NUM.findall(q)]; par=[int(m) for m in _PAREN_NUM.findall(q)]
    rng=0
    for a,b in _RANGE.findall(q):
        a,b=int(a),int(b)
        if 0<b-a<60: rng=max(rng,b-a+1)
    cand=max(len(set(line)),len(set(par)))
    if rng and cand and rng!=cand: return cand
    cand=max(cand,len(set(_LINE_LET.findall(q)))); n=max(rng,cand)
    if n>1: return n
    lines=[l.strip() for l in q.splitlines() if l.strip()]
    if len(lines)>1:
        body=lines[1:] if lines[0].endswith((":",".")) else lines
        if body: return len(body)
    if context:
        cn=len(set(int(m) for m in _LINE_NUM.findall(context)))
        if cn>1: return cn
        cl=len(set(_LINE_LET.findall(context)))
        if cl>1: return cl
    return max(n,1)
def extract_item_sources(query,n):
    q=query 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:
        lines=[l.strip() for l in q.splitlines() if l.strip()]
        if len(lines)>1 and lines[0].endswith((":",".")): out=lines[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]

# ---- parsing ----------------------------------------------------------------
_STRIP_PREFIX=re.compile(r"^\s*(?:\(?\d{1,3}\)?[.):\]]\s*|[-*•]\s+)"); _FENCE=re.compile(r"^```[a-zA-Z]*\s*$")
_CHATTY=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|let me\b|first,|so,|therefore\b|thus\b)",re.I)
def clean_line(s):
    s=s.strip(); s=_STRIP_PREFIX.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 fit_to_n(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 raw_lines(text,n,fb=None):
    """Champion base parse: every non-empty line, forced to N, never blank."""
    lines=[ln.strip() for ln in (text or "").splitlines() if ln.strip()]
    return fit_to_n(lines,n,fb)
def parse_answers(text,n,fb=None):
    """Robust parse for CoT / samples: slice after ANSWERS:, else salvage tail."""
    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 _FENCE.match(ln): continue
        mm=re.match(r"^\s*\(?(\d{1,3})\)?[.):\]]\s*(.+)$",ln.strip())
        if mm:
            v=clean_line(mm.group(2))
            if v and not _CHATTY.match(v): numbered.append((int(mm.group(1)),v))
        c=clean_line(ln)
        if c and not _CHATTY.match(c): raw.append(c)
    if len(numbered)>=n:
        by={};  [by.__setitem__(l,v) for l,v in numbered]
        labs=sorted(by)
        if len(labs)>=n: return [by[l] for l in labs[:n]]
    return fit_to_n(raw,n,fb)
def norm(s):
    s=unicodedata.normalize("NFC",(s or "").strip().lower()); s=re.sub(r"\s+"," ",s)
    return s.strip(" .!?;:,")
def vote(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[norm(c)].append(c)
    a_sup=len(groups.get(norm(anchor),[])); bk,bn=None,0
    for k,v in groups.items():
        if len(v)>bn: bk,bn=k,len(v)
    if bk is not None and bn>=2 and bn>a_sup: return Counter(groups[bk]).most_common(1)[0][0]
    return anchor

# ---- prompts ----------------------------------------------------------------
SYS_TRIVIAL=("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.")
SYS_COT=("You solve International Linguistics Olympiad problems. Everything needed is in the "
    "problem. Briefly analyse: segment the words, find the recurring morphemes and the rules "
    "that order them, and check them against every given example. Keep the analysis short. "
    "Then write a line containing exactly ANSWERS: and, below it, one answer per item in the "
    "order asked — no numbering, no commentary, no blank lines. Never leave an item blank.")

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

def main():
    import pandas as pd, torch
    from transformers import (AutoTokenizer,AutoModelForCausalLM,StoppingCriteria,StoppingCriteriaList)
    torch.backends.cuda.matmul.allow_tf32=True; torch.backends.cudnn.allow_tf32=True
    log(f"MODE: {'CoT' if COT else 'trivial'} | rep=1.0 | max_samples={MAX_SAMPLES} temp={SAMPLE_TEMP}")

    df=pd.read_csv(TEST_CSV,dtype=str).fillna("")
    ids=[str(x) for x in df["id"].tolist()]
    ns=[detect_n_items(r.get("query",""),r.get("task_type",""),r.get("context","")) for _,r in df.iterrows()]
    log(f"loaded {len(df)} problems, {sum(ns)} items")
    srcs={i:extract_item_sources(r.get("query",""),n) for i,(_,r),n in zip(ids,df.iterrows(),ns)}

    preds={i:list(srcs[i]) for i in ids}; expl={i:"" for i in ids} if WANT_EXPL else None
    write_submission(OUT_CSV,ids,preds,expl); log(f"placeholder written ({len(ids)} rows)")

    class Deadline(StoppingCriteria):
        def __init__(self,t): self.t=t
        def __call__(self,i,s,**k): return time.time()>self.t
    log("loading model ...")
    tok=AutoTokenizer.from_pretrained(MODEL_ID,trust_remote_code=True)
    if tok.pad_token is None: tok.pad_token=tok.eos_token
    tok.padding_side="left"
    def _load(dm):
        try: return AutoModelForCausalLM.from_pretrained(MODEL_ID,torch_dtype=torch.float16,device_map=dm,trust_remote_code=True).eval()
        except TypeError: return AutoModelForCausalLM.from_pretrained(MODEL_ID,dtype=torch.float16,device_map=dm,trust_remote_code=True).eval()
    try: model=_load({"":0} if torch.cuda.is_available() else "auto")
    except Exception as e: log(f"pinned load failed ({e}); auto"); model=_load("auto")
    log(f"model ready ({left():.0f}s left)")

    sysmsg=SYS_COT if COT else SYS_TRIVIAL
    base_parse=parse_answers if COT else raw_lines      # <-- the champion distinction
    prompts=[tok.apply_chat_template(
        [{"role":"system","content":sysmsg},
         {"role":"user","content":f"{r['context'].strip()}\n\n{r['query'].strip()}"}],
        tokenize=False,add_generation_prompt=True) for _,r in df.iterrows()]

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

    cap=max(MAX_NEW,1200) if COT else MAX_NEW           # reasoning needs headroom
    adaptive=int(0.40*max(1.0,left())*TOK_PER_S/max(1,len(df)))
    max_new=max(192,min(cap,adaptive))
    log(f"Pass 1 greedy, {max_new} tok/item")
    t=time.time(); texts=generate(prompts,max_new=max_new,sample=False); c1=time.time()-t
    samples={i:[] for i in ids}
    for i,n,txt in zip(ids,ns,texts):
        a=base_parse(txt,n,srcs[i]); preds[i]=a; samples[i].append(a)
    write_submission(OUT_CSV,ids,preds,expl); log(f"Pass 1 done in {c1:.0f}s")

    reserve=min(300.0,0.25*c1+60) if WANT_EXPL else 30.0
    ne=0
    while left()-reserve>c1*1.25 and ne<MAX_SAMPLES:
        ne+=1; log(f"sample pass {ne} ({left():.0f}s left)")
        texts=generate(prompts,max_new=max_new,sample=True,temp=SAMPLE_TEMP)
        for i,n,txt in zip(ids,ns,texts):
            if txt: samples[i].append(parse_answers(txt,n,srcs[i]))
        for i,n in zip(ids,ns):
            if len(samples[i])>=3:
                g=samples[i][0]
                preds[i]=fit_to_n([vote([s[k] for s in samples[i] if k<len(s)],
                                        anchor=g[k] if k<len(g) else None) for k in range(n)],n,srcs[i])
        write_submission(OUT_CSV,ids,preds,expl); log(f"voted over {ne+1} samples")
    log(f"self-consistency: {ne} sample pass(es) completed")

    if WANT_EXPL and left()>60:
        log(f"explanations ({left():.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. Concise (2-4 sentences).")
        ep=[tok.apply_chat_template([{"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 preds[str(r['id'])])+"\n\nBriefly explain the linguistic rules."}],
            tokenize=False,add_generation_prompt=True) for _,r in df.iterrows()]
        for i,e in zip(ids,generate(ep,max_new=200,sample=False)):
            e=re.sub(r"\s+"," ",(e or "").strip())
            if e: expl[i]=e[:1200]
        write_submission(OUT_CSV,ids,preds,expl); log("explanations written")

    bad=[i for i,n in zip(ids,ns) if len(preds[i])!=n or any(not str(x).strip() for x in preds[i])]
    if bad:
        for i,n in zip(ids,ns): preds[i]=fit_to_n([x for x in preds[i] if str(x).strip()],n,srcs[i])
        write_submission(OUT_CSV,ids,preds,expl)
    log(f"DONE. {len(ids)} rows, {time.time()-T0:.0f}s elapsed.")

if __name__=="__main__":
    main()