#!/usr/bin/env python3
"""bank_eval.py (day-52) — does the artifact bank COMPOUND solve-rate, weights FROZEN?

Sealed meter = data/af/band_meter_v3.jsonl (320, DISJOINT from fuel_big). For each
meter problem we RETRIEVE the top-K most similar banked (statement, solution) pairs
by pure-python BM25 over tokenized statements, build a few-shot prompt (analogous
solved problems from the disjoint pool — never the answer), greedy-generate with the
frozen base, and grade on ALL the target's sampled tests.

COMPOUNDING TEST: sweep bank size S over BANK_SIZES (first S bank records = the
retrievable set; S=0 = no retrieval = frozen-base control). Compounding = solved/320
rises monotonically with S while the model never changes.

Env: BASE (/workspace/RSI/newbase), METER, BANK (/workspace/RSI/outputs/bank.jsonl),
     BANK_SIZES ("0,100,300,600,1000"), K (3), EVAL_TC (40), MAXTOK (900),
     MAXLEN (16384), GPU_UTIL (0.85), N (meter subset cap, for smoke tests),
     OUT (/workspace/RSI/outputs/bank_eval.json).
Prints: BANK_COMPOUNDING sizes=[...] solved=[...]
"""
import os, sys, json, math, re, tempfile, subprocess
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
sys.path.insert(0, "/workspace/RSI")
from src.utils.vllm_backend import VLLMModelLoader          # noqa: E402
from src.utils.external_benchmarks import _extract_code     # noqa: E402

BASE = os.environ.get("BASE", "/workspace/RSI/newbase")
METER = os.environ.get("METER", "/workspace/RSI/data/af/band_meter_v3.jsonl")
BANK = os.environ.get("BANK", "/workspace/RSI/outputs/bank.jsonl")
SIZES = [int(x) for x in os.environ.get("BANK_SIZES", "0,100,300,600,1000").split(",")]
K = int(os.environ.get("K", "3"))
EVAL_TC = int(os.environ.get("EVAL_TC", "40"))
MAXTOK = int(os.environ.get("MAXTOK", "900"))
MAXLEN = int(os.environ.get("MAXLEN", "16384"))
N = int(os.environ.get("N", "100000"))
OUT = os.environ.get("OUT", "/workspace/RSI/outputs/bank_eval.json")
BNB = {"load_in_4bit": True, "bnb_4bit_compute_dtype": "bfloat16"}

# ---- grading (identical semantics to band_ruler.py) ----
def _norm(s): return "\n".join(l.rstrip() for l in str(s).strip().splitlines())
def _run1(code, inp):
    d = tempfile.mkdtemp(); p = os.path.join(d, "s.py"); open(p, "w").write(code)
    try: return subprocess.run([sys.executable, p], input=str(inp), capture_output=True, text=True, timeout=6).stdout
    except Exception: return ""
    finally:
        try: os.remove(p); os.rmdir(d)
        except Exception: pass
_pool = ThreadPoolExecutor(max_workers=96)
def grade(code, ins, outs):
    if not code or not ins: return 0.0
    tc = list(zip(ins, outs)); step = max(1, len(tc)//EVAL_TC); tc = tc[::step][:EVAL_TC]
    return sum(_pool.map(lambda io: _norm(_run1(code, io[0])) == _norm(io[1]), tc)) / max(1, len(tc))

# ---- pure-python BM25 (no deps) ----
_tok_re = re.compile(r"[a-z0-9]+")
def toks(s): return _tok_re.findall(s.lower())

class BM25:
    def __init__(self, docs, k1=1.5, b=0.75):
        self.k1, self.b = k1, b
        self.docs = [Counter(toks(d)) for d in docs]
        self.dl = [sum(c.values()) for c in self.docs]
        self.avgdl = (sum(self.dl) / len(self.dl)) if self.dl else 1.0
        df = Counter()
        for c in self.docs: df.update(c.keys())
        n = len(self.docs)
        self.idf = {t: math.log(1 + (n - d + 0.5) / (d + 0.5)) for t, d in df.items()}
    def top(self, query, k):
        q = toks(query)
        scores = []
        for i, c in enumerate(self.docs):
            s = 0.0
            for t in q:
                f = c.get(t, 0)
                if not f: continue
                s += self.idf.get(t, 0.0) * f * (self.k1 + 1) / (f + self.k1 * (1 - self.b + self.b * self.dl[i] / self.avgdl))
            scores.append((s, i))
        scores.sort(reverse=True)
        return [i for s, i in scores[:k] if s > 0]

def split_chat(p):
    """(prefix incl user tag, user body, suffix from user <|im_end|>)."""
    a = p.find("<|im_start|>user\n")
    if a == -1: return None
    a += len("<|im_start|>user\n")
    b = p.find("<|im_end|>", a)
    if b == -1: return None
    return p[:a], p[a:b], p[b:]

def fewshot_prompt(target_prompt, exemplars):
    """exemplars = [(statement, solution)]; insert into the user turn, budget to MAXLEN."""
    sp = split_chat(target_prompt)
    if sp is None:
        prefix, body, suffix = "", target_prompt, ""
    else:
        prefix, body, suffix = sp
    budget = (MAXLEN - MAXTOK - 128) * 3  # ~3 chars/token, conservative
    exs = list(exemplars)
    while True:
        blocks = []
        for i, (st, sol) in enumerate(exs, 1):
            blocks.append(f"### Solved example {i}\nProblem:\n{st[:1500]}\n\nCorrect solution:\n```python\n{sol[:2500]}\n```\n")
        head = ("Here are previously solved problems similar to the one you must solve. "
                "Study their approaches, then solve the NEW problem.\n\n" + "\n".join(blocks) +
                "\n### NEW problem (solve THIS one)\n") if exs else ""
        full = prefix + head + body + suffix
        if len(full) <= budget or not exs: return full
        exs = exs[:-1]  # drop least-similar exemplar until it fits

meter = [json.loads(l) for l in open(METER) if l.strip()][:N]
for r in meter:
    assert r.get("prompt") and r.get("inputs") and r.get("outputs"), f"bad meter row {r.get('pid')}"
bank = [json.loads(l) for l in open(BANK) if l.strip()] if os.path.exists(BANK) else []
SIZES = sorted(set(min(s, len(bank)) for s in SIZES))
print(f"[bank_eval] meter={len(meter)} bank={len(bank)} sizes={SIZES} K={K} BASE={BASE}", flush=True)

loader = VLLMModelLoader(model_path=BASE, dtype="bfloat16", max_model_len=MAXLEN,
    gpu_memory_utilization=float(os.environ.get("GPU_UTIL", "0.85")),
    allow_remote_code=True, quantization_config=BNB, enforce_eager=True)
loader.load()
assert getattr(loader, "_llm", None) is not None, "FATAL: vLLM did not load (no cross-backend verdicts)"
print(f"[bank_eval] loaded {BASE} backend=vllm (FROZEN, no adapter)", flush=True)

results = {}   # size -> {"solved": int, "bits": [...], "mean_partial": float}
for S in SIZES:
    sub = bank[:S]
    if S > 0:
        bm = BM25([b["statement"] for b in sub])
    prompts = []
    n_with_ex = 0
    for r in meter:
        if S == 0:
            prompts.append(r["prompt"]); continue
        stmt = split_chat(r["prompt"])
        query = stmt[1] if stmt else r["prompt"]
        idxs = bm.top(query, K)
        exs = [(sub[i]["statement"], sub[i]["solution"]) for i in idxs]
        if exs: n_with_ex += 1
        prompts.append(fewshot_prompt(r["prompt"], exs))
    print(f"[bank_eval] S={S}: {n_with_ex}/{len(meter)} meter problems got exemplars; generating (greedy)...", flush=True)
    outs = loader.generate_batch(prompts, max_new_tokens=MAXTOK, temperature=0.0, top_p=1.0)
    bits, part = [], []
    for r, o in zip(meter, outs):
        g = grade(_extract_code(o) or o, r["inputs"], r["outputs"])
        part.append(g); bits.append(1 if g >= 0.999 else 0)
    results[S] = {"solved": sum(bits), "bits": bits, "mean_partial": sum(part) / len(part)}
    print(f"[bank_eval] S={S}: solved {sum(bits)}/{len(meter)} mean_partial={sum(part)/len(part):.4f}", flush=True)

os.makedirs(os.path.dirname(OUT), exist_ok=True)
json.dump({"base": BASE, "meter": METER, "bank": BANK, "k": K,
           "pids": [r["pid"] for r in meter],
           "sizes": SIZES, "results": {str(s): results[s] for s in SIZES}}, open(OUT, "w"))
print(f"[bank_eval] results -> {OUT}", flush=True)
print(f"BANK_COMPOUNDING sizes={SIZES} solved={[results[s]['solved'] for s in SIZES]}", flush=True)