                    results.append(("reject", code, {"reason": "policy", "policy_reason": _policy_reason(code_for_policy2), "score": 0.0}))
                    try:
                        mem.log({
                            "step": step,
                            "mode": "arena",
                            "cand": int(idx),
                            "status": "reject",
                            "policy_reason": _policy_reason(code_for_policy2),
                            "prompt_head": prompt_head,
                            "code_head": (code or ""),
                        })
                    except Exception:
                        pass
                    try:
                        if bool(cfg.get("log_console", False)):
                            print(f"[selfplay] cand[{idx}] reject reason=policy:{_policy_reason(code_for_policy2)}")
                    except Exception:
                        pass
                    continue
                # Prefilter: extract solution; accept raw code with def solve or fenced block
                def _extract_fenced(txt: str) -> str | None:
                    if txt.count("```") < 2:
                        return None
                    try:
                        s = txt
                        first_nl = s.find("\n")
                        next_fence = s.find("```", first_nl + 1) if first_nl != -1 else -1
                        fence_end = next_fence if next_fence != -1 else s.rfind("```")
                        if first_nl == -1 or fence_end == -1 or fence_end <= first_nl:
                            return None
                        inner = s[first_nl + 1:fence_end]
                        lines = inner.splitlines()
                        if lines and re.match(r"^[a-zA-Z0-9_+-]+$", lines[0].strip()):
                            inner = "\n".join(lines[1:])
                        inner = inner.strip()
                        return inner if inner else None
                    except Exception:
                        return None
                # Try to extract first fenced code block from raw text before sanitizing
                raw_for_extract = code.lstrip("\ufeff")
                body = extract_first_fenced_block(raw_for_extract) or _extract_fenced(raw_for_extract)
                # If no fenced block, fall back to sanitizing raw text and heuristics
                # Allow a leading language header line in raw outputs (e.g., "python")
                code_sanitized = code.lstrip("\ufeff")
                try:
                    first_nonempty = next((ln for ln in code_sanitized.splitlines() if ln.strip()), "")
                except Exception:
                    first_nonempty = ""
                if first_nonempty.strip().lower() in ("python", "python3"):
                    try:
                        lines = code_sanitized.splitlines()
                        first_idx = next((i for i, ln in enumerate(lines) if ln.strip()), 0)
                        code_sanitized = "\n".join(lines[first_idx+1:])
                    except Exception:
                        pass
                # Strip diagnostic noise that some models echo into completions
                try:
                    cleaned_lines: List[str] = []


                    except Exception:
                        explore_p = 0.3
                    # Apply replay bins even in fallback path
                    spec = None
                    try:
                        bank = task_bank_load()
                        spec = _pick_from_bank_with_bins(bank)
                    except Exception:
                        spec = None
                    if not spec:
                        spec = task_bank_pick(explore_p)
                    if not spec:
                        # Target a learning band 40%-70% to keep challenge adaptive
                        spec = model_generate_task(tok, model, curr.level, target_sr=(0.4, 0.7), log_fn=lambda o: mem.log(o))
            else:
                # Apply replay bins in static sampler path
                spec = None
                try:
                    bank = task_bank_load()
                    # Reuse helper
                    def _pick_from_bank_with_bins(bank: Dict[str, Any]) -> Dict[str, Any] | None:
                        try:
                            wb = str(os.environ.get("SELFPLAY_REPLAY_BINS", "0.4,0.4,0.2")).split(",")
                            w_fresh, w_recent, w_long = [float(wb[i]) if i < len(wb) else 0.0 for i in range(3)]
                        except Exception:
                            w_fresh, w_recent, w_long = 0.4, 0.4, 0.2
                        try:
                            topk = int(os.environ.get("SELFPLAY_REPLAY_TOPK", "5") or 5)
                        except Exception:
                            topk = 5
                        import random as _r  # type: ignore
                        r = _r.random()
                        cum_fresh = max(0.0, w_fresh)
                        cum_recent = cum_fresh + max(0.0, w_recent)
                        picked = "fresh"
                        if r < cum_fresh:
                            picked = "fresh"
                        elif r < cum_recent:
                            picked = "recent"
                        else:
                            picked = "longtail"
                        if (not bank) or picked == "fresh":
                            return None
                        entries = [v for v in bank.values() if isinstance(v.get("spec"), dict)]
                        if not entries:
                            return None
                        if picked == "recent":
                            entries.sort(key=lambda x: (-(int(x.get("n", 0))), float(x.get("wins", 0)) / float(max(1, int(x.get("n", 0))))) )
                        else:
                            entries.sort(key=lambda x: (int(x.get("n", 0)), int(x.get("wins", 0))))
                        pool = entries[: max(1, topk)]
                        cooled = [e for e in pool if str(e.get("name", "")) not in set(recent_tasks)]
                        if not cooled:
                            return None
                        pick = _r.choice(cooled)
                        return pick.get("spec") if isinstance(pick.get("spec"), dict) else None
                    spec = _pick_from_bank_with_bins(bank)
                except Exception:
                    spec = None
                if spec is None:
                    spec = sample_task(curr.level)
            base_prompt = render_prompt(spec)
            task_name = str(spec.get("name", ""))
            try:
                recent_tasks.append(task_name)
            except Exception:
                pass
            # Telemetry: summarize selected task
            try:
                mem.log({
                    "task_select": {
                        "name": task_name,
                        "hash": task_hash(spec),
                        "bank": task_bank_stats(task_name),
                        "summary": summarize_tests(spec),
                    }


                    min_steps = int(cfg.get("min_steps", 1))
                    patience = int(cfg.get("patience", 1))
                    # Curriculum-based scheduling
                    try:
                        L_curr = int(level_before)
                        if L_curr <= 1:
                            # Easier levels: allow more patience/steps
                            patience = max(patience, 2)
                        elif L_curr >= 4:
                            # Harder levels: cap halting tighter
                            patience = max(1, min(patience, 2))
                    except Exception:
                        pass
                    patience_ctr = 0
                    # TRM-inspired exploration and halting budget
                    try:
                        p_explore = float(os.environ.get("SELFPLAY_REFINE_EXPLORATION_P", "0.1") or 0.1)
                    except Exception:
                        p_explore = 0.1
                    # Anneal exploration over refinement steps k
                    def _anneal(p0: float, step_k: int, total_K: int) -> float:
                        try:
                            return max(0.01, p0 * (1.0 - float(step_k) / float(max(1, total_K))))
                        except Exception:
                            return p0
                    try:
                        halt_max = int(os.environ.get("SELFPLAY_HALT_MAX_STEPS", "0") or 0)
                    except Exception:
                        halt_max = 0
                    try:
                        sp_margin = float(os.environ.get("SELFPLAY_VERIFIER_SP_MARGIN", "0.02") or 0.02)
                    except Exception:
                        sp_margin = 0.02
                    try:
                        v_every = int(os.environ.get("SELFPLAY_VERIFIER_EVERY", "2") or 2)
                    except Exception:
                        v_every = 2
                    halt_reason = "budget_exhausted"
                    k_taken = 0
                    aborted = False
                    resets = 0
                    try:
                        max_resets = int(os.environ.get("SELFPLAY_REFINE_MAX_RESETS", "1") or 1)
                    except Exception:
                        max_resets = 1


                        except Exception:
                            sp_final = str(scratchpad)
                        ptxt_b = base_inner + "\n# Use the scratchpad to produce the final solution.\n"
                        ptxt_b += "\n# SCRATCHPAD FINAL\n" + sp_final + "\n"
                        cands_b = sample(tok, model, ptxt_b, n=k2, max_new_tokens=int(cfg.get("max_tokens", 256)), stop=["```"])
                        try:
                            used_tokens_approx += int(sum(len(cb) for cb in cands_b) // 4)
                        except Exception:
                            pass
                        if token_budget > 0 and used_tokens_approx >= token_budget:
                            halt_reason = "token_budget"
                            aborted = True
                            # No candidates to evaluate due to budget
                            results_b = []
                            break
                        results_b: List[Tuple[str, str, Dict[str, Any]]] = []
                        for code_b in cands_b:
                            # Precheck Phase-B candidate: extract body then syntax + subset tests + quick verifier
                            code_b_s = code_b.lstrip("\ufeff")
                            # Try to extract fenced body similar to Phase-A
                            def _extract_fenced_inner(txt: str) -> str | None:
                                if txt.count("```") < 2:
                                    return None
                                try:
                                    s = txt
                                    first_nl = s.find("\n")
                                    next_fence = s.find("```", first_nl + 1) if first_nl != -1 else -1
                                    fence_end = next_fence if next_fence != -1 else s.rfind("```")
                                    if first_nl == -1 or fence_end == -1 or fence_end <= first_nl:
                                        return None
                                    inner = s[first_nl + 1:fence_end]
                                    lines = inner.splitlines()
                                    if lines and re.match(r"^[a-zA-Z0-9_+-]+$", lines[0].strip()):
                                        inner = "\n".join(lines[1:])
                                    inner = inner.strip()
                                    return inner if inner else None
                                except Exception:
                                    return None
                            body_b = _extract_fenced_inner(code_b_s)
                            if body_b is None:
                                if re.search(r"^\s*def\s+solve\s*\(.*\):\s*$", code_b_s, flags=re.MULTILINE):
                                    body_b = code_b_s
                                else:
                                    results_b.append(("fail", code_b, {"score": 0.0, "reason": "no_def_solve"}))
                                    continue
                            try:
                                _ = ast.parse(body_b)
                                _syntax_ok_b = True
                            except Exception:
                                _syntax_ok_b = False
                            if not _syntax_ok_b:
                                results_b.append(("fail", code_b, {"score": 0.0, "reason": "syntax_error"}))
                                continue
                            # HITL import approval gate (GF Phase-B path)
                            try:
                                _mods_union_b = sorted(set((_extract_imports(code_b) or []) + (_extract_imports(body_b) or [])))
                                if not _maybe_gate_imports(_mods_union_b, bool(cfg.get("require_import_approval", False)), int(cfg.get("approval_timeout", 60))):
                                    results_b.append(("fail", code_b, {"score": 0.0, "reason": "imports_unapproved"}))


from __future__ import annotations

from .common_imports import *
from .generation import generate_text
from .tasks_foundry import _REGISTRY as _FOUNDRY_REG, _mk_tests as _FOUNDRY_MK, TaskTemplate as _FT  # type: ignore

def _list_templates(level: int, want_tags: list[str]) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []
    for t in _FOUNDRY_REG:
        try:
            name = getattr(t, "name", "")
            sig = getattr(t, "signature", "")
            tags = list(getattr(t, "tags", []) or [])
            diff = int(getattr(t, "difficulty", 0))
            ok_diff = diff <= min(5, int(level) + 2)
            ok_tags = (not want_tags) or any((tg in tags) for tg in want_tags)
            if ok_diff and ok_tags:
                items.append({"name": name, "signature": sig, "tags": tags, "difficulty": diff})
        except Exception:
            continue
    if not items:
        # Fallback to all templates if filter was too strict
        for t in _FOUNDRY_REG:
            try:
                items.append({
                    "name": getattr(t, "name", ""),
                    "signature": getattr(t, "signature", ""),
                    "tags": list(getattr(t, "tags", []) or []),
                    "difficulty": int(getattr(t, "difficulty", 0)),
                })
            except Exception:
                continue
    return items

def llm_sample_task(tok, model, level: int, tags: list[str] | None = None, log_prompts: bool = False) -> Dict[str, Any]:
    """Use the LLM to choose a foundry template, then instantiate tests via sampler/metamorphics.

    Returns a TaskSpec compatible with render_prompt and run_and_test.
    """
    want_tags = list(tags or [])
    items = _list_templates(level, want_tags)
    # Build a concise selection prompt
    options_text = "\n".join(
        f"- name: {it['name']} | diff: {it['difficulty']} | tags: {','.join(it['tags'])} | sig: {it['signature']}" for it in items[:40]
    )
    system = (
        "You are selecting a task template for a code selfplay benchmark. "
        "Return ONLY a compact JSON object with the chosen template name."
    )
    user = (
        "Templates (subset):\n" + options_text +
        f"\n\nLevel: {int(level)}. Prefer difficulty <= level+2. "
        + (f"Prefer tags: {','.join(want_tags)}. " if want_tags else "") +
        "Respond with JSON of the form {\"name\": \"<template_name>\"} and nothing else."
    )
    logger = get_prompt_logger("foundry", bool(log_prompts))  # type: ignore
    if logger:
        try:
            logger.log_text("prompt", system + "\n\n" + user)
        except Exception:
            pass
    # Try chat first; fallback to generate
    try:
        messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
        text = str(model.chat(messages, max_new_tokens=128, temperature=0.0))
    except Exception:
        text = str(generate_text(tok, model, user, max_new_tokens=128))
    if logger:
        try:
            logger.log_text("response", text)
        except Exception:
            pass
    # Parse JSON {"name": "..."}
    sel_name = None
    try:
        import json as _json  # type: ignore
        sel_name = str((_json.loads(text) or {}).get("name", "")).strip() or None
    except Exception:
        sel_name = None
    # Fallback: try to extract a name token from free text
    if not sel_name:
        try:
            for it in items:
                if it["name"] in text:
                    sel_name = it["name"]
                    break
        except Exception:
            sel_name = None
    # Choose template object
    T: _FT | None = None  # type: ignore
    try:
        for t in _FOUNDRY_REG:
            if getattr(t, "name", None) == sel_name:
                T = t
                break
    except Exception:
        T = None
    if T is None:
        # Fallback: choose first filtered item or any
        try:
            target = items[0] if items else None
            if target is not None:
                for t in _FOUNDRY_REG:
                    if getattr(t, "name", None) == target.get("name"):
                        T = t
                        break
        except Exception:
            T = None
    if T is None and _FOUNDRY_REG:
        T = _FOUNDRY_REG[0]
    if T is None:
        raise ValueError("no_foundry_templates")

    # Seeded RNG similar to tasks.sample_task
    import random as _rnd  # type: ignore
    try:
        seed_env = os.environ.get("SELFPLAY_SEED_EP")
        if seed_env is None or str(seed_env).strip() == "":
            seed_env = os.environ.get("SELFPLAY_SEED")
        seed_val = int(seed_env) if seed_env is not None and str(seed_env).strip() != "" else None
    except Exception:
        seed_val = None
    rng = _rnd.Random(seed_val) if seed_val is not None else _rnd.Random()

    # Instantiate tests
    inputs, oracle = T.sampler(rng)
    tests = _FOUNDRY_MK(inputs, oracle)
    try:
        for mr in getattr(T, "metamorphics", []) or []:
            try:
                for (inp, out) in mr(inputs, oracle):
                    tests.append({"inp": list(inp), "out": out})
            except Exception:
                continue
    except Exception:
        pass
    if len(tests) < 8:
        base = list(tests)
        while len(tests) < 8 and base:
            tests.append(base[len(tests) % len(base)])
    tests = tests[:32]
    return {"name": getattr(T, "name", "foundry"), "signature": getattr(T, "signature", "def solve(n:int)->int:"), "tests": tests}


from __future__ import annotations

from .common_imports import *
from .generation import generate_text
from .tasks_foundry import _REGISTRY as _FOUNDRY_REG, _mk_tests as _FOUNDRY_MK, TaskTemplate as _FT  # type: ignore

def _list_templates(level: int, want_tags: list[str]) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []
    for t in _FOUNDRY_REG:
        try:
            name = getattr(t, "name", "")
            sig = getattr(t, "signature", "")
            tags = list(getattr(t, "tags", []) or [])
            diff = int(getattr(t, "difficulty", 0))
            ok_diff = diff <= min(5, int(level) + 2)
            ok_tags = (not want_tags) or any((tg in tags) for tg in want_tags)
            if ok_diff and ok_tags:
                items.append({"name": name, "signature": sig, "tags": tags, "difficulty": diff})
        except Exception:
            continue
    if not items:
        # Fallback to all templates if filter was too strict
        for t in _FOUNDRY_REG:
            try:
                items.append({
                    "name": getattr(t, "name", ""),
                    "signature": getattr(t, "signature", ""),
                    "tags": list(getattr(t, "tags", []) or []),
                    "difficulty": int(getattr(t, "difficulty", 0)),
                })
            except Exception:
                continue
    return items

def llm_sample_task(tok, model, level: int, tags: list[str] | None = None, log_prompts: bool = False) -> Dict[str, Any]:
    """Use the LLM to choose a foundry template, then instantiate tests via sampler/metamorphics.

    Returns a TaskSpec compatible with render_prompt and run_and_test.
    """
    want_tags = list(tags or [])
    items = _list_templates(level, want_tags)
    # Build a concise selection prompt
    options_text = "\n".join(
        f"- name: {it['name']} | diff: {it['difficulty']} | tags: {','.join(it['tags'])} | sig: {it['signature']}" for it in items[:40]
    )
    system = (
        "You are selecting a task template for a code selfplay benchmark. "
        "Return ONLY a compact JSON object with the chosen template name."
    )
    user = (
        "Templates (subset):\n" + options_text +
        f"\n\nLevel: {int(level)}. Prefer difficulty <= level+2. "


            ok_tags = (not want_tags) or any((tg in tags) for tg in want_tags)
            if ok_diff and ok_tags:
                items.append({"name": name, "signature": sig, "tags": tags, "difficulty": diff})
        except Exception:
            continue
    if not items:
        # Fallback to all templates if filter was too strict
        for t in _FOUNDRY_REG:
            try:
                items.append({
                    "name": getattr(t, "name", ""),
                    "signature": getattr(t, "signature", ""),
                    "tags": list(getattr(t, "tags", []) or []),
                    "difficulty": int(getattr(t, "difficulty", 0)),
                })
            except Exception:
                continue
    return items

def llm_sample_task(tok, model, level: int, tags: list[str] | None = None, log_prompts: bool = False) -> Dict[str, Any]:
    """Use the LLM to choose a foundry template, then instantiate tests via sampler/metamorphics.

    Returns a TaskSpec compatible with render_prompt and run_and_test.
    """
    want_tags = list(tags or [])
    items = _list_templates(level, want_tags)
    # Build a concise selection prompt
    options_text = "\n".join(
        f"- name: {it['name']} | diff: {it['difficulty']} | tags: {','.join(it['tags'])} | sig: {it['signature']}" for it in items[:40]
    )
    system = (
        "You are selecting a task template for a code selfplay benchmark. "
        "Return ONLY a compact JSON object with the chosen template name."
    )
    user = (
        "Templates (subset):\n" + options_text +
        f"\n\nLevel: {int(level)}. Prefer difficulty <= level+2. "
        + (f"Prefer tags: {','.join(want_tags)}. " if want_tags else "") +
        "Respond with JSON of the form {\"name\": \"<template_name>\"} and nothing else."
    )
    logger = get_prompt_logger("foundry", bool(log_prompts))  # type: ignore
    if logger:
        try:
            logger.log_text("prompt", system + "\n\n" + user)
        except Exception:
            pass
    # Try chat first; fallback to generate
    try:
        messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
        text = str(model.chat(messages, max_new_tokens=128, temperature=0.0))
    except Exception:
        text = str(generate_text(tok, model, user, max_new_tokens=128))
    if logger:
        try:
            logger.log_text("response", text)
        except Exception:
            pass
    # Parse JSON {"name": "..."}
    sel_name = None
    try:
        import json as _json  # type: ignore
        sel_name = str((_json.loads(text) or {}).get("name", "")).strip() or None
    except Exception:
        sel_name = None
    # Fallback: try to extract a name token from free text
    if not sel_name:
        try:
            for it in items:
                if it["name"] in text:
                    sel_name = it["name"]
                    break
        except Exception:
            sel_name = None
    # Choose template object
    T: _FT | None = None  # type: ignore
    try:
        for t in _FOUNDRY_REG:
            if getattr(t, "name", None) == sel_name:
                T = t
                break
    except Exception:
        T = None
    if T is None:
        # Fallback: choose first filtered item or any
        try:
            target = items[0] if items else None
            if target is not None:
                for t in _FOUNDRY_REG:
                    if getattr(t, "name", None) == target.get("name"):
                        T = t
                        break
        except Exception:
            T = None
    if T is None and _FOUNDRY_REG:
        T = _FOUNDRY_REG[0]
    if T is None:
        raise ValueError("no_foundry_templates")

    # Seeded RNG similar to tasks.sample_task
    import random as _rnd  # type: ignore
    try:
        seed_env = os.environ.get("SELFPLAY_SEED_EP")
        if seed_env is None or str(seed_env).strip() == "":
            seed_env = os.environ.get("SELFPLAY_SEED")
        seed_val = int(seed_env) if seed_env is not None and str(seed_env).strip() != "" else None
    except Exception:
        seed_val = None
    rng = _rnd.Random(seed_val) if seed_val is not None else _rnd.Random()

    # Instantiate tests
    inputs, oracle = T.sampler(rng)
    tests = _FOUNDRY_MK(inputs, oracle)
    try:
        for mr in getattr(T, "metamorphics", []) or []:
            try:
                for (inp, out) in mr(inputs, oracle):
                    tests.append({"inp": list(inp), "out": out})
            except Exception:
                continue
    except Exception:
        pass
    if len(tests) < 8:
        base = list(tests)
        while len(tests) < 8 and base:
            tests.append(base[len(tests) % len(base)])
    tests = tests[:32]
    return {"name": getattr(T, "name", "foundry"), "signature": getattr(T, "signature", "def solve(n:int)->int:"), "tests": tests}


from __future__ import annotations

from .common_imports import *
from .paths import data_path

def task_bank_path() -> Path:
    return data_path("selfplay", "task_bank.jsonl")

def parse_denylist() -> set[str]:
    try:
        dl = str(os.environ.get("SELFPLAY_TASK_DENYLIST", "")).strip()
        if not dl:
            return set()
        return set(x.strip() for x in dl.split(",") if x.strip())
    except Exception:
        return set()

def task_bank_load() -> Dict[str, Any]:
    bank: Dict[str, Any] = {}
    try:
        p = task_bank_path()
        if p.exists():
            for ln in p.read_text(encoding="utf-8").splitlines():
                ln = ln.strip()
                if not ln:
                    continue
                try:
                    obj = json.loads(ln)
                    name = str(obj.get("name", "")).strip()
                    if name:
                        bank[name] = obj
                except Exception:
                    continue
        # Purge denylisted
        try:
            deny = parse_denylist()
            if deny:
                changed = False
                for name in list(bank.keys()):
                    if name in deny:
                        bank.pop(name, None)
                        changed = True
                if changed:
                    task_bank_save(bank)
        except Exception:
            pass
    except Exception:
        bank = {}
    return bank

def task_bank_save(bank: Dict[str, Any]) -> None:
    try:
        p = task_bank_path()
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w", encoding="utf-8") as f:
            for name, obj in bank.items():
                try:
                    f.write(json.dumps(obj, ensure_ascii=False) + "\n")
                except Exception:
                    continue
    except Exception:
        pass

def task_bank_add(spec: Dict[str, Any]) -> None:
    try:
        bank = task_bank_load()
        name = str(spec.get("name", "")).strip() or "gen_task"
        if name in parse_denylist():
            return
        if name not in bank:
            bank[name] = {"name": name, "spec": spec, "n": 0, "wins": 0, "last_ts": int(__import__("time").time())}
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update(name: str, passed: bool) -> None:
    try:
        bank = task_bank_load()
        if name in bank:
            bank[name]["n"] = int(bank[name].get("n", 0)) + 1
            if passed:
                bank[name]["wins"] = int(bank[name].get("wins", 0)) + 1
            bank[name]["last_ts"] = int(__import__("time").time())
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update_assim(name: str, gain: float) -> None:
    """Update assimilation statistics for a task: simple EMA of gain and count.

    gain is expected in [0, 1]. We keep an EMA with alpha from env (default 0.2).
    """
    try:
        bank = task_bank_load()
        if name not in bank:
            return
        try:
            alpha = float(os.environ.get("SELFPLAY_ASSIM_ALPHA", "0.2") or 0.2)
        except Exception:
            alpha = 0.2
        prev = float(bank[name].get("assim_ema", 0.0) or 0.0)
        new_ema = (1.0 - alpha) * prev + alpha * float(max(0.0, min(1.0, float(gain))))
        bank[name]["assim_ema"] = float(new_ema)
        bank[name]["assim_n"] = int(bank[name].get("assim_n", 0)) + 1
        task_bank_save(bank)
    except Exception:
        pass

def task_bank_pick(explore_p: float) -> Dict[str, Any] | None:
    try:
        import random as _r  # type: ignore
        bank = task_bank_load()
        if (not bank) or (_r.random() < max(0.0, min(1.0, float(explore_p)))):
            return None
        # UCB1 over tasks
        T = sum(int(v.get("n", 0)) for v in bank.values())
        T = max(1, int(T))
        def _ucb(obj: Dict[str, Any]) -> float:
            n = max(1, int(obj.get("n", 0)))
            w = float(obj.get("wins", 0)) / float(n)
            return w + math.sqrt(2.0 * math.log(T) / float(n))
        best = max(bank.values(), key=_ucb)
        spec = best.get("spec")
        return spec if isinstance(spec, dict) else None
    except Exception:
        return None

def task_bank_stats(name: str) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    try:
        bank = task_bank_load()
        if name in bank:
            entry = bank[name]
            n = int(entry.get("n", 0))
            wins = int(entry.get("wins", 0))
            T = sum(int(v.get("n", 0)) for v in bank.values())
            T = max(1, int(T))
            n_eff = max(1, n)
            ucb = (float(wins) / float(n_eff)) + math.sqrt(2.0 * math.log(T) / float(n_eff))
            out = {"n": n, "wins": wins, "ucb": float(ucb)}
            try:
                if "assim_ema" in entry:
                    out["assim_ema"] = float(entry.get("assim_ema", 0.0) or 0.0)
                    out["assim_n"] = int(entry.get("assim_n", 0) or 0)
            except Exception:
                pass
    except Exception:
        out = {}
    return out

def task_hash(spec: Dict[str, Any]) -> str:
    try:
        import hashlib as _hl  # type: ignore
        sig = str(spec.get("signature", ""))
        name = str(spec.get("name", ""))
        tests = spec.get("tests", [])
        key = {"name": name, "signature": sig, "tests": tests}
        blob = json.dumps(key, sort_keys=True, ensure_ascii=False).encode("utf-8")
        return _hl.sha1(blob).hexdigest()
    except Exception:
        return ""

def summarize_tests(spec: Dict[str, Any]) -> Dict[str, Any]:
    try:
        tests = spec.get("tests", [])
        ints: List[int] = []
        list_lens: List[int] = []
        str_lens: List[int] = []
        for t in tests[:64]:
            for a in (t.get("inp", []) if isinstance(t, dict) else []):
                if isinstance(a, bool):
                    continue
                if isinstance(a, int):
                    ints.append(int(a))
                elif isinstance(a, str):
                    str_lens.append(len(a))
                elif isinstance(a, list):
                    list_lens.append(len(a))
        out: Dict[str, Any] = {"tests": int(len(tests))}
        if ints:
            out["int_min"] = int(min(ints)); out["int_max"] = int(max(ints))
        if list_lens:
            out["list_len_min"] = int(min(list_lens)); out["list_len_max"] = int(max(list_lens))
        if str_lens:
            out["str_len_min"] = int(min(str_lens)); out["str_len_max"] = int(max(str_lens))
        return out
    except Exception:
        return {"tests": None}


from __future__ import annotations

from .common_imports import *
from .paths import data_path

def task_bank_path() -> Path:
    return data_path("selfplay", "task_bank.jsonl")

def parse_denylist() -> set[str]:
    try:
        dl = str(os.environ.get("SELFPLAY_TASK_DENYLIST", "")).strip()
        if not dl:
            return set()
        return set(x.strip() for x in dl.split(",") if x.strip())
    except Exception:
        return set()

def task_bank_load() -> Dict[str, Any]:
    bank: Dict[str, Any] = {}
    try:
        p = task_bank_path()
        if p.exists():
            for ln in p.read_text(encoding="utf-8").splitlines():
                ln = ln.strip()
                if not ln:


from __future__ import annotations

from .common_imports import *
from .paths import data_path

def task_bank_path() -> Path:
    return data_path("selfplay", "task_bank.jsonl")

def parse_denylist() -> set[str]:
    try:
        dl = str(os.environ.get("SELFPLAY_TASK_DENYLIST", "")).strip()
        if not dl:
            return set()
        return set(x.strip() for x in dl.split(",") if x.strip())
    except Exception:
        return set()

def task_bank_load() -> Dict[str, Any]:
    bank: Dict[str, Any] = {}
    try:
        p = task_bank_path()
        if p.exists():
            for ln in p.read_text(encoding="utf-8").splitlines():
                ln = ln.strip()
                if not ln:
                    continue
                try:
                    obj = json.loads(ln)
                    name = str(obj.get("name", "")).strip()
                    if name:
                        bank[name] = obj
                except Exception:
                    continue
        # Purge denylisted
        try:


from __future__ import annotations

from .common_imports import *
from .paths import data_path

def task_bank_path() -> Path:
    return data_path("selfplay", "task_bank.jsonl")

def parse_denylist() -> set[str]:
    try:
        dl = str(os.environ.get("SELFPLAY_TASK_DENYLIST", "")).strip()
        if not dl:
            return set()
        return set(x.strip() for x in dl.split(",") if x.strip())
    except Exception:
        return set()

def task_bank_load() -> Dict[str, Any]:
    bank: Dict[str, Any] = {}
    try:
        p = task_bank_path()
        if p.exists():
            for ln in p.read_text(encoding="utf-8").splitlines():
                ln = ln.strip()
                if not ln:
                    continue
                try:
                    obj = json.loads(ln)
                    name = str(obj.get("name", "")).strip()
                    if name:
                        bank[name] = obj
                except Exception:
                    continue
        # Purge denylisted
        try:
            deny = parse_denylist()
            if deny:
                changed = False
                for name in list(bank.keys()):
                    if name in deny:
                        bank.pop(name, None)
                        changed = True
                if changed:
                    task_bank_save(bank)
        except Exception:
            pass
    except Exception:
        bank = {}
    return bank

def task_bank_save(bank: Dict[str, Any]) -> None:
    try:
        p = task_bank_path()
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w", encoding="utf-8") as f:
            for name, obj in bank.items():
                try:
                    f.write(json.dumps(obj, ensure_ascii=False) + "\n")
                except Exception:
                    continue
    except Exception:
        pass

def task_bank_add(spec: Dict[str, Any]) -> None:
    try:
        bank = task_bank_load()
        name = str(spec.get("name", "")).strip() or "gen_task"


                except Exception:
                    continue
        # Purge denylisted
        try:
            deny = parse_denylist()
            if deny:
                changed = False
                for name in list(bank.keys()):
                    if name in deny:
                        bank.pop(name, None)
                        changed = True
                if changed:
                    task_bank_save(bank)
        except Exception:
            pass
    except Exception:
        bank = {}
    return bank

def task_bank_save(bank: Dict[str, Any]) -> None:
    try:
        p = task_bank_path()
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w", encoding="utf-8") as f:
            for name, obj in bank.items():
                try:
                    f.write(json.dumps(obj, ensure_ascii=False) + "\n")
                except Exception:
                    continue
    except Exception:
        pass

def task_bank_add(spec: Dict[str, Any]) -> None:
    try:
        bank = task_bank_load()
        name = str(spec.get("name", "")).strip() or "gen_task"
        if name in parse_denylist():
            return
        if name not in bank:
            bank[name] = {"name": name, "spec": spec, "n": 0, "wins": 0, "last_ts": int(__import__("time").time())}
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update(name: str, passed: bool) -> None:
    try:
        bank = task_bank_load()
        if name in bank:
            bank[name]["n"] = int(bank[name].get("n", 0)) + 1


            pass
    except Exception:
        bank = {}
    return bank

def task_bank_save(bank: Dict[str, Any]) -> None:
    try:
        p = task_bank_path()
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, "w", encoding="utf-8") as f:
            for name, obj in bank.items():
                try:
                    f.write(json.dumps(obj, ensure_ascii=False) + "\n")
                except Exception:
                    continue
    except Exception:
        pass

def task_bank_add(spec: Dict[str, Any]) -> None:
    try:
        bank = task_bank_load()
        name = str(spec.get("name", "")).strip() or "gen_task"
        if name in parse_denylist():
            return
        if name not in bank:
            bank[name] = {"name": name, "spec": spec, "n": 0, "wins": 0, "last_ts": int(__import__("time").time())}
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update(name: str, passed: bool) -> None:
    try:
        bank = task_bank_load()
        if name in bank:
            bank[name]["n"] = int(bank[name].get("n", 0)) + 1
            if passed:
                bank[name]["wins"] = int(bank[name].get("wins", 0)) + 1
            bank[name]["last_ts"] = int(__import__("time").time())
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update_assim(name: str, gain: float) -> None:
    """Update assimilation statistics for a task: simple EMA of gain and count.

    gain is expected in [0, 1]. We keep an EMA with alpha from env (default 0.2).
    """


                    f.write(json.dumps(obj, ensure_ascii=False) + "\n")
                except Exception:
                    continue
    except Exception:
        pass

def task_bank_add(spec: Dict[str, Any]) -> None:
    try:
        bank = task_bank_load()
        name = str(spec.get("name", "")).strip() or "gen_task"
        if name in parse_denylist():
            return
        if name not in bank:
            bank[name] = {"name": name, "spec": spec, "n": 0, "wins": 0, "last_ts": int(__import__("time").time())}
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update(name: str, passed: bool) -> None:
    try:
        bank = task_bank_load()
        if name in bank:
            bank[name]["n"] = int(bank[name].get("n", 0)) + 1
            if passed:
                bank[name]["wins"] = int(bank[name].get("wins", 0)) + 1
            bank[name]["last_ts"] = int(__import__("time").time())
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update_assim(name: str, gain: float) -> None:
    """Update assimilation statistics for a task: simple EMA of gain and count.

    gain is expected in [0, 1]. We keep an EMA with alpha from env (default 0.2).
    """
    try:
        bank = task_bank_load()
        if name not in bank:
            return
        try:
            alpha = float(os.environ.get("SELFPLAY_ASSIM_ALPHA", "0.2") or 0.2)
        except Exception:
            alpha = 0.2
        prev = float(bank[name].get("assim_ema", 0.0) or 0.0)
        new_ema = (1.0 - alpha) * prev + alpha * float(max(0.0, min(1.0, float(gain))))
        bank[name]["assim_ema"] = float(new_ema)
        bank[name]["assim_n"] = int(bank[name].get("assim_n", 0)) + 1
        task_bank_save(bank)


        if name not in bank:
            bank[name] = {"name": name, "spec": spec, "n": 0, "wins": 0, "last_ts": int(__import__("time").time())}
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update(name: str, passed: bool) -> None:
    try:
        bank = task_bank_load()
        if name in bank:
            bank[name]["n"] = int(bank[name].get("n", 0)) + 1
            if passed:
                bank[name]["wins"] = int(bank[name].get("wins", 0)) + 1
            bank[name]["last_ts"] = int(__import__("time").time())
            task_bank_save(bank)
    except Exception:
        pass

def task_bank_update_assim(name: str, gain: float) -> None:
    """Update assimilation statistics for a task: simple EMA of gain and count.

    gain is expected in [0, 1]. We keep an EMA with alpha from env (default 0.2).
    """
    try:
        bank = task_bank_load()
        if name not in bank:
            return
        try:
            alpha = float(os.environ.get("SELFPLAY_ASSIM_ALPHA", "0.2") or 0.2)
        except Exception:
            alpha = 0.2
        prev = float(bank[name].get("assim_ema", 0.0) or 0.0)
        new_ema = (1.0 - alpha) * prev + alpha * float(max(0.0, min(1.0, float(gain))))
        bank[name]["assim_ema"] = float(new_ema)
        bank[name]["assim_n"] = int(bank[name].get("assim_n", 0)) + 1
        task_bank_save(bank)
    except Exception:
        pass

def task_bank_pick(explore_p: float) -> Dict[str, Any] | None:
    try:
        import random as _r  # type: ignore
        bank = task_bank_load()
        if (not bank) or (_r.random() < max(0.0, min(1.0, float(explore_p)))):
            return None
        # UCB1 over tasks
        T = sum(int(v.get("n", 0)) for v in bank.values())
        T = max(1, int(T))
        def _ucb(obj: Dict[str, Any]) -> float:
            n = max(1, int(obj.get("n", 0)))
            w = float(obj.get("wins", 0)) / float(n)
            return w + math.sqrt(2.0 * math.log(T) / float(n))
        best = max(bank.values(), key=_ucb)
        spec = best.get("spec")
        return spec if isinstance(spec, dict) else None
    except Exception:
        return None


    gain is expected in [0, 1]. We keep an EMA with alpha from env (default 0.2).
    """
    try:
        bank = task_bank_load()
        if name not in bank:
            return
        try:
            alpha = float(os.environ.get("SELFPLAY_ASSIM_ALPHA", "0.2") or 0.2)
        except Exception:
            alpha = 0.2
        prev = float(bank[name].get("assim_ema", 0.0) or 0.0)
        new_ema = (1.0 - alpha) * prev + alpha * float(max(0.0, min(1.0, float(gain))))
        bank[name]["assim_ema"] = float(new_ema)
        bank[name]["assim_n"] = int(bank[name].get("assim_n", 0)) + 1
        task_bank_save(bank)
    except Exception:
        pass

def task_bank_pick(explore_p: float) -> Dict[str, Any] | None:
    try:
        import random as _r  # type: ignore
        bank = task_bank_load()
        if (not bank) or (_r.random() < max(0.0, min(1.0, float(explore_p)))):
            return None
        # UCB1 over tasks
        T = sum(int(v.get("n", 0)) for v in bank.values())
        T = max(1, int(T))
        def _ucb(obj: Dict[str, Any]) -> float:
            n = max(1, int(obj.get("n", 0)))
            w = float(obj.get("wins", 0)) / float(n)
            return w + math.sqrt(2.0 * math.log(T) / float(n))
        best = max(bank.values(), key=_ucb)
        spec = best.get("spec")
        return spec if isinstance(spec, dict) else None
    except Exception:
        return None

def task_bank_stats(name: str) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    try:
        bank = task_bank_load()
        if name in bank:
            entry = bank[name]
            n = int(entry.get("n", 0))
            wins = int(entry.get("wins", 0))
            T = sum(int(v.get("n", 0)) for v in bank.values())
            T = max(1, int(T))
            n_eff = max(1, n)
            ucb = (float(wins) / float(n_eff)) + math.sqrt(2.0 * math.log(T) / float(n_eff))
            out = {"n": n, "wins": wins, "ucb": float(ucb)}
            try:
                if "assim_ema" in entry:
                    out["assim_ema"] = float(entry.get("assim_ema", 0.0) or 0.0)
                    out["assim_n"] = int(entry.get("assim_n", 0) or 0)
            except Exception:


def task_bank_pick(explore_p: float) -> Dict[str, Any] | None:
    try:
        import random as _r  # type: ignore
        bank = task_bank_load()
        if (not bank) or (_r.random() < max(0.0, min(1.0, float(explore_p)))):
            return None
        # UCB1 over tasks
        T = sum(int(v.get("n", 0)) for v in bank.values())
        T = max(1, int(T))
        def _ucb(obj: Dict[str, Any]) -> float:
            n = max(1, int(obj.get("n", 0)))
            w = float(obj.get("wins", 0)) / float(n)
            return w + math.sqrt(2.0 * math.log(T) / float(n))
        best = max(bank.values(), key=_ucb)
        spec = best.get("spec")
        return spec if isinstance(spec, dict) else None
    except Exception:
        return None

def task_bank_stats(name: str) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    try:
        bank = task_bank_load()
        if name in bank:
            entry = bank[name]
            n = int(entry.get("n", 0))
            wins = int(entry.get("wins", 0))
            T = sum(int(v.get("n", 0)) for v in bank.values())
            T = max(1, int(T))
            n_eff = max(1, n)
            ucb = (float(wins) / float(n_eff)) + math.sqrt(2.0 * math.log(T) / float(n_eff))
            out = {"n": n, "wins": wins, "ucb": float(ucb)}
            try:
                if "assim_ema" in entry:
                    out["assim_ema"] = float(entry.get("assim_ema", 0.0) or 0.0)
                    out["assim_n"] = int(entry.get("assim_n", 0) or 0)
            except Exception:
                pass
    except Exception:
        out = {}
    return out

def task_hash(spec: Dict[str, Any]) -> str:
    try:
        import hashlib as _hl  # type: ignore
        sig = str(spec.get("signature", ""))
        name = str(spec.get("name", ""))
        tests = spec.get("tests", [])
        key = {"name": name, "signature": sig, "tests": tests}
        blob = json.dumps(key, sort_keys=True, ensure_ascii=False).encode("utf-8")
        return _hl.sha1(blob).hexdigest()
    except Exception:
        return ""

def summarize_tests(spec: Dict[str, Any]) -> Dict[str, Any]:
    try:
        tests = spec.get("tests", [])
        ints: List[int] = []
        list_lens: List[int] = []


        if name in bank:
            entry = bank[name]
            n = int(entry.get("n", 0))
            wins = int(entry.get("wins", 0))
            T = sum(int(v.get("n", 0)) for v in bank.values())
            T = max(1, int(T))
            n_eff = max(1, n)
            ucb = (float(wins) / float(n_eff)) + math.sqrt(2.0 * math.log(T) / float(n_eff))
            out = {"n": n, "wins": wins, "ucb": float(ucb)}
            try:
                if "assim_ema" in entry:
                    out["assim_ema"] = float(entry.get("assim_ema", 0.0) or 0.0)
                    out["assim_n"] = int(entry.get("assim_n", 0) or 0)
            except Exception:
                pass
    except Exception:
        out = {}
    return out

def task_hash(spec: Dict[str, Any]) -> str:
    try:
        import hashlib as _hl  # type: ignore
        sig = str(spec.get("signature", ""))
        name = str(spec.get("name", ""))
        tests = spec.get("tests", [])
        key = {"name": name, "signature": sig, "tests": tests}
        blob = json.dumps(key, sort_keys=True, ensure_ascii=False).encode("utf-8")
        return _hl.sha1(blob).hexdigest()
    except Exception:
        return ""

def summarize_tests(spec: Dict[str, Any]) -> Dict[str, Any]:
    try:
        tests = spec.get("tests", [])
        ints: List[int] = []
        list_lens: List[int] = []
        str_lens: List[int] = []
        for t in tests[:64]:
            for a in (t.get("inp", []) if isinstance(t, dict) else []):
                if isinstance(a, bool):
                    continue
                if isinstance(a, int):
                    ints.append(int(a))
                elif isinstance(a, str):
                    str_lens.append(len(a))
                elif isinstance(a, list):
                    list_lens.append(len(a))
        out: Dict[str, Any] = {"tests": int(len(tests))}
        if ints:


            except Exception:
                pass
    except Exception:
        out = {}
    return out

def task_hash(spec: Dict[str, Any]) -> str:
    try:
        import hashlib as _hl  # type: ignore
        sig = str(spec.get("signature", ""))
        name = str(spec.get("name", ""))
        tests = spec.get("tests", [])
        key = {"name": name, "signature": sig, "tests": tests}
        blob = json.dumps(key, sort_keys=True, ensure_ascii=False).encode("utf-8")
        return _hl.sha1(blob).hexdigest()
    except Exception:
        return ""

def summarize_tests(spec: Dict[str, Any]) -> Dict[str, Any]:
    try:
        tests = spec.get("tests", [])
        ints: List[int] = []
        list_lens: List[int] = []
        str_lens: List[int] = []
        for t in tests[:64]:
            for a in (t.get("inp", []) if isinstance(t, dict) else []):
                if isinstance(a, bool):
                    continue
                if isinstance(a, int):
                    ints.append(int(a))
                elif isinstance(a, str):
                    str_lens.append(len(a))
                elif isinstance(a, list):
                    list_lens.append(len(a))
        out: Dict[str, Any] = {"tests": int(len(tests))}
        if ints:
            out["int_min"] = int(min(ints)); out["int_max"] = int(max(ints))
        if list_lens:
            out["list_len_min"] = int(min(list_lens)); out["list_len_max"] = int(max(list_lens))
        if str_lens:
            out["str_len_min"] = int(min(str_lens)); out["str_len_max"] = int(max(str_lens))
        return out
    except Exception:
        return {"tests": None}


        except Exception:
            alpha = 0.2
        prev = float(bank[name].get("assim_ema", 0.0) or 0.0)
        new_ema = (1.0 - alpha) * prev + alpha * float(max(0.0, min(1.0, float(gain))))
        bank[name]["assim_ema"] = float(new_ema)
        bank[name]["assim_n"] = int(bank[name].get("assim_n", 0)) + 1
        task_bank_save(bank)
    except Exception:
        pass

def task_bank_pick(explore_p: float) -> Dict[str, Any] | None:
    try:
        import random as _r  # type: ignore
        bank = task_bank_load()
        if (not bank) or (_r.random() < max(0.0, min(1.0, float(explore_p)))):
            return None
        # UCB1 over tasks
        T = sum(int(v.get("n", 0)) for v in bank.values())
        T = max(1, int(T))
        def _ucb(obj: Dict[str, Any]) -> float:
            n = max(1, int(obj.get("n", 0)))
            w = float(obj.get("wins", 0)) / float(n)
            return w + math.sqrt(2.0 * math.log(T) / float(n))
        best = max(bank.values(), key=_ucb)
        spec = best.get("spec")
        return spec if isinstance(spec, dict) else None
    except Exception:
        return None

def task_bank_stats(name: str) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    try:
        bank = task_bank_load()
        if name in bank:
            entry = bank[name]
            n = int(entry.get("n", 0))
            wins = int(entry.get("wins", 0))
            T = sum(int(v.get("n", 0)) for v in bank.values())
            T = max(1, int(T))
            n_eff = max(1, n)
            ucb = (float(wins) / float(n_eff)) + math.sqrt(2.0 * math.log(T) / float(n_eff))
            out = {"n": n, "wins": wins, "ucb": float(ucb)}


from .common_imports import *
from .lora import _inject_lora, _patch_forwards

def _build_device_map_and_max_memory() -> Tuple[Any, Optional[Dict[Any, str]]]:
    """Build (device_map, max_memory) using env overrides, preferring larger GPUs.

    Env controls:
    - HF_DEVICE_MAP: explicit device_map string (e.g., "auto", "balanced_low_0")
    - HF_MAX_MEMORY_JSON: JSON dict like {"0":"23GiB","1":"23GiB","2":"6GiB","cpu":"64GiB"}
    - SELFPLAY_AVOID_GPU_IDS: comma-separated GPU ids to de-prioritize; they get "1GiB"
    - SELFPLAY_MAX_MEM_HEADROOM_GIB: integer GiB to leave free per preferred GPU (default 1)
    """
    # device_map
    dm_env = os.environ.get("HF_DEVICE_MAP", None)
    device_map: Any = (dm_env if (dm_env is not None and str(dm_env).strip() != "") else ("auto" if torch.cuda.is_available() else None))
    # max_memory
    mm_env = os.environ.get("HF_MAX_MEMORY_JSON", None)
    if mm_env:
        try:
            mm = json.loads(mm_env)
            # normalize keys to int where possible
            norm: Dict[Any, str] = {}
            for k, v in mm.items():
                try:
                    kk: Any = int(k)
                except Exception:
                    kk = k
                norm[kk] = str(v)
            return device_map, norm
        except Exception:
            pass
    # Build from GPU properties with avoid list
    max_memory: Dict[Any, str] = {}
    try:
        avoid_env = os.environ.get("SELFPLAY_AVOID_GPU_IDS", "")
        avoid = {int(x) for x in avoid_env.split(",") if x.strip().isdigit()}
    except Exception:
        avoid = set()
    try:
        headroom = int(os.environ.get("SELFPLAY_MAX_MEM_HEADROOM_GIB", "1") or 1)
    except Exception:
        headroom = 1
    if torch.cuda.is_available():
        try:
            n = torch.cuda.device_count()
            for i in range(n):
                props = torch.cuda.get_device_properties(i)
                total_gib = int(max(1, props.total_memory // (1024**3)))
                if i in avoid:
                    cap = 1  # keep 1GiB to avoid placement
                else:
                    cap = max(1, total_gib - headroom)
                max_memory[i] = f"{cap}GiB"
            # allow CPU fallback if needed
            max_memory["cpu"] = os.environ.get("HF_MAX_MEMORY_CPU", "64GiB")
        except Exception:
            max_memory = {}
    return device_map, (max_memory if max_memory else None)

def load_seed(cfg: Dict[str, Any]) -> Tuple[Any, Any]:
    # Resolve best adapter/model defaults
    root = Path(__file__).resolve().parents[1]
    # Prefer explicit adapter_dir in cfg; else try common best paths
    adapter_dir = str(cfg.get("adapter_dir", "") or "").strip()
    if not adapter_dir:
        try:
            # Try verifier_qlora and coder_qlora best/default dirs
            for rel in ("models/coder_qlora", "models/verifier_qlora", "models/coder_qlora.humaneval"):
                p = (root / rel)
                if p.exists():
                    adapter_dir = str(p)
                    break
        except Exception:
            adapter_dir = ""
    # Optional direct HF load if requested
    if bool(cfg.get("direct_model", False)):
        tok = AutoTokenizer.from_pretrained(cfg["model_name"])  # type: ignore
        try:
            if getattr(tok, "pad_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
                tok.pad_token_id = tok.eos_token_id  # type: ignore[attr-defined]
        except Exception:
            pass
        # Prefer GPU with sharded placement and BF16/FP16 when available
        device_map, max_memory = _build_device_map_and_max_memory()
        # dtype heuristic: BF16 if supported, else FP16 on CUDA
        torch_dtype = None
        if torch.cuda.is_available():
            try:
                if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
                    torch_dtype = torch.bfloat16
                else:
                    torch_dtype = torch.float16
            except Exception:
                torch_dtype = torch.float16
        model = AutoModelForCausalLM.from_pretrained(  # type: ignore
            cfg["model_name"],
            device_map=device_map,
            max_memory=max_memory,
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True,
        )
        # If adapter present, attach via AdapterCache for consistency
        try:
            if adapter_dir:
                tok, model = AdapterCache.get(cfg["model_name"], adapter_dir)
        except Exception:
            pass
    else:
        client = HFClient.get_cached(cfg["model_name"])  # HF device/dtype defaults
        tok = client.tokenizer
        try:
            if getattr(tok, "pad_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
                tok.pad_token_id = tok.eos_token_id  # type: ignore[attr-defined]
        except Exception:
            pass
        # Attach adapter if available
        if adapter_dir:
            try:
                tok, model = AdapterCache.get(cfg["model_name"], adapter_dir)
            except Exception:
                model = client.model
        else:
            model = client.model
    model.eval()
    # Optional gradient checkpointing to reduce activation memory
    try:
        if str(os.environ.get("SELFPLAY_GRAD_CKPT", "0")).strip() not in ("", "0", "false", "False"):
            if hasattr(model, "gradient_checkpointing_enable"):
                model.gradient_checkpointing_enable()
    except Exception:
        pass
    # Enable TF32 for speed on supported GPUs
    try:
        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
            torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Ensure generation config has proper special token ids and disable unsupported SWA
    try:
        if getattr(model, "generation_config", None) is not None:
            if getattr(model.generation_config, "pad_token_id", None) is None and getattr(tok, "pad_token_id", None) is not None:
                model.generation_config.pad_token_id = tok.pad_token_id  # type: ignore[attr-defined]
            if getattr(model.generation_config, "eos_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
                model.generation_config.eos_token_id = tok.eos_token_id  # type: ignore[attr-defined]
    except Exception:
        pass
    try:
        # Some backends log warnings for sliding window attention with SDPA; disable if present
        if getattr(getattr(model, "config", object()), "sliding_window", None) is not None:
            model.config.sliding_window = None  # type: ignore[attr-defined]
    except Exception:
        pass
    # Skip custom LoRA injection if a PEFT adapter is attached unless explicitly forced
    should_inject = bool(cfg.get("force_lora_inject", False))
    is_peft_model = False
    try:
        from peft import PeftModel  # type: ignore
        if isinstance(model, PeftModel):
            is_peft_model = True
            # Only inject custom LoRA if explicitly forced
            should_inject = bool(cfg.get("force_lora_inject", False))
    except Exception:
        is_peft_model = False
    if should_inject:
        model = _inject_lora(model, r=int(cfg.get("adapter_r", 8)), alpha=int(cfg.get("adapter_alpha", 16)))
        model = _patch_forwards(model)
    # Configure trainable parameters:
    # - If we injected custom LoRA: train only A/B; freeze base
    # - If using PEFT adapter: respect PEFT's own trainable flags
    # - Otherwise: freeze all (inference-only)
    if should_inject:
        for n, p in model.named_parameters():
            if ".A" in n or ".B" in n:
                p.requires_grad_(True)
            else:
                p.requires_grad_(False)
    elif is_peft_model:
        pass  # leave PEFT-configured trainable params as-is
    else:
        for _, p in model.named_parameters():
            p.requires_grad_(False)
    # Auto-inject LoRA if no trainable params (gated)
    try:
        has_trainable = any(p.requires_grad for _, p in model.named_parameters())
        if bool(cfg.get("auto_lora_if_frozen", False)) and (not has_trainable):
            model = _inject_lora(model, r=int(cfg.get("adapter_r", 8)), alpha=int(cfg.get("adapter_alpha", 16)))
            model = _patch_forwards(model)
            for n, p in model.named_parameters():
                if ".A" in n or ".B" in n:
                    p.requires_grad_(True)
                else:
                    p.requires_grad_(False)
    except Exception:
        pass
    return tok, model

def save_adapters(model: torch.nn.Module, path: str = "adapters.pt") -> None:
    # Ensure target directory exists for checkpoints
    try:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    adap = {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad}
    torch.save(adap, path)
    # Sidecar JSON with telemetry
    try:
        sidecar = str(path) + ".json"
        trainable_params = int(sum(int(p.numel()) for p in model.parameters() if p.requires_grad))
        # Try to infer adapter rank from any LoRA A matrix
        adapter_r = None
        try:
            for _n, _m in model.named_modules():
                if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                    adapter_r = int(getattr(_m, "A").shape[0])  # type: ignore[attr-defined]
                    break
        except Exception:
            adapter_r = None
        with open(sidecar, "w", encoding="utf-8") as f:
            json.dump({"trainable_params": trainable_params, "adapter_r": adapter_r}, f, ensure_ascii=False)
    except Exception:
        pass

def load_adapters(model: torch.nn.Module, path: str) -> int:
    """Load adapter weights saved by save_adapters into matching trainable params.

    Returns number of parameters loaded.
    """
    # Load adapters onto model device when possible; fallback to CPU
    try:
        dev = next(model.parameters()).device
    except Exception:
        dev = "cpu"
    try:
        state = torch.load(path, map_location=dev, weights_only=True)
    except Exception:
        try:
            state = torch.load(path, map_location="cpu", weights_only=True)
        except Exception:
            return 0
    if not isinstance(state, dict):
        return 0
    loaded = 0
    for n, p in model.named_parameters():
        try:
            if n in state:
                src = state[n]
                if isinstance(src, torch.Tensor) and tuple(src.shape) == tuple(p.shape):
                    with torch.no_grad():
                        p.data.copy_(src.to(p.device, dtype=p.dtype))
                    loaded += 1
        except Exception:
            continue
    return loaded


from .common_imports import *
from .lora import _inject_lora, _patch_forwards

def _build_device_map_and_max_memory() -> Tuple[Any, Optional[Dict[Any, str]]]:
    """Build (device_map, max_memory) using env overrides, preferring larger GPUs.

    Env controls:
    - HF_DEVICE_MAP: explicit device_map string (e.g., "auto", "balanced_low_0")
    - HF_MAX_MEMORY_JSON: JSON dict like {"0":"23GiB","1":"23GiB","2":"6GiB","cpu":"64GiB"}
    - SELFPLAY_AVOID_GPU_IDS: comma-separated GPU ids to de-prioritize; they get "1GiB"
    - SELFPLAY_MAX_MEM_HEADROOM_GIB: integer GiB to leave free per preferred GPU (default 1)
    """
    # device_map
    dm_env = os.environ.get("HF_DEVICE_MAP", None)
    device_map: Any = (dm_env if (dm_env is not None and str(dm_env).strip() != "") else ("auto" if torch.cuda.is_available() else None))
    # max_memory
    mm_env = os.environ.get("HF_MAX_MEMORY_JSON", None)
    if mm_env:
        try:
            mm = json.loads(mm_env)
            # normalize keys to int where possible
            norm: Dict[Any, str] = {}
            for k, v in mm.items():
                try:
                    kk: Any = int(k)
                except Exception:
                    kk = k
                norm[kk] = str(v)
            return device_map, norm
        except Exception:
            pass
    # Build from GPU properties with avoid list
    max_memory: Dict[Any, str] = {}
    try:
        avoid_env = os.environ.get("SELFPLAY_AVOID_GPU_IDS", "")
        avoid = {int(x) for x in avoid_env.split(",") if x.strip().isdigit()}
    except Exception:
        avoid = set()
    try:
        headroom = int(os.environ.get("SELFPLAY_MAX_MEM_HEADROOM_GIB", "1") or 1)
    except Exception:
        headroom = 1
    if torch.cuda.is_available():
        try:
            n = torch.cuda.device_count()
            for i in range(n):
                props = torch.cuda.get_device_properties(i)
                total_gib = int(max(1, props.total_memory // (1024**3)))
                if i in avoid:
                    cap = 1  # keep 1GiB to avoid placement
                else:
                    cap = max(1, total_gib - headroom)
                max_memory[i] = f"{cap}GiB"
            # allow CPU fallback if needed
            max_memory["cpu"] = os.environ.get("HF_MAX_MEMORY_CPU", "64GiB")
        except Exception:
            max_memory = {}
    return device_map, (max_memory if max_memory else None)

def load_seed(cfg: Dict[str, Any]) -> Tuple[Any, Any]:
    # Resolve best adapter/model defaults
    root = Path(__file__).resolve().parents[1]
    # Prefer explicit adapter_dir in cfg; else try common best paths
    adapter_dir = str(cfg.get("adapter_dir", "") or "").strip()
    if not adapter_dir:
        try:
            # Try verifier_qlora and coder_qlora best/default dirs
            for rel in ("models/coder_qlora", "models/verifier_qlora", "models/coder_qlora.humaneval"):
                p = (root / rel)
                if p.exists():
                    adapter_dir = str(p)
                    break
        except Exception:
            adapter_dir = ""
    # Optional direct HF load if requested
    if bool(cfg.get("direct_model", False)):
        tok = AutoTokenizer.from_pretrained(cfg["model_name"])  # type: ignore
        try:


        headroom = int(os.environ.get("SELFPLAY_MAX_MEM_HEADROOM_GIB", "1") or 1)
    except Exception:
        headroom = 1
    if torch.cuda.is_available():
        try:
            n = torch.cuda.device_count()
            for i in range(n):
                props = torch.cuda.get_device_properties(i)
                total_gib = int(max(1, props.total_memory // (1024**3)))
                if i in avoid:
                    cap = 1  # keep 1GiB to avoid placement
                else:
                    cap = max(1, total_gib - headroom)
                max_memory[i] = f"{cap}GiB"
            # allow CPU fallback if needed
            max_memory["cpu"] = os.environ.get("HF_MAX_MEMORY_CPU", "64GiB")
        except Exception:
            max_memory = {}
    return device_map, (max_memory if max_memory else None)

def load_seed(cfg: Dict[str, Any]) -> Tuple[Any, Any]:
    # Resolve best adapter/model defaults
    root = Path(__file__).resolve().parents[1]
    # Prefer explicit adapter_dir in cfg; else try common best paths
    adapter_dir = str(cfg.get("adapter_dir", "") or "").strip()
    if not adapter_dir:
        try:
            # Try verifier_qlora and coder_qlora best/default dirs
            for rel in ("models/coder_qlora", "models/verifier_qlora", "models/coder_qlora.humaneval"):
                p = (root / rel)
                if p.exists():
                    adapter_dir = str(p)
                    break
        except Exception:
            adapter_dir = ""
    # Optional direct HF load if requested
    if bool(cfg.get("direct_model", False)):
        tok = AutoTokenizer.from_pretrained(cfg["model_name"])  # type: ignore
        try:
            if getattr(tok, "pad_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
                tok.pad_token_id = tok.eos_token_id  # type: ignore[attr-defined]
        except Exception:
            pass
        # Prefer GPU with sharded placement and BF16/FP16 when available
        device_map, max_memory = _build_device_map_and_max_memory()
        # dtype heuristic: BF16 if supported, else FP16 on CUDA
        torch_dtype = None
        if torch.cuda.is_available():
            try:
                if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
                    torch_dtype = torch.bfloat16
                else:
                    torch_dtype = torch.float16
            except Exception:
                torch_dtype = torch.float16
        model = AutoModelForCausalLM.from_pretrained(  # type: ignore
            cfg["model_name"],
            device_map=device_map,
            max_memory=max_memory,
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True,
        )
        # If adapter present, attach via AdapterCache for consistency
        try:
            if adapter_dir:
                tok, model = AdapterCache.get(cfg["model_name"], adapter_dir)
        except Exception:
            pass
    else:
        client = HFClient.get_cached(cfg["model_name"])  # HF device/dtype defaults
        tok = client.tokenizer
        try:
            if getattr(tok, "pad_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
                tok.pad_token_id = tok.eos_token_id  # type: ignore[attr-defined]
        except Exception:
            pass
        # Attach adapter if available
        if adapter_dir:
            try:
                tok, model = AdapterCache.get(cfg["model_name"], adapter_dir)
            except Exception:
                model = client.model
        else:
            model = client.model
    model.eval()
    # Optional gradient checkpointing to reduce activation memory
    try:
        if str(os.environ.get("SELFPLAY_GRAD_CKPT", "0")).strip() not in ("", "0", "false", "False"):
            if hasattr(model, "gradient_checkpointing_enable"):
                model.gradient_checkpointing_enable()
    except Exception:
        pass
    # Enable TF32 for speed on supported GPUs
    try:
        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
            torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
    except Exception:
        pass
    # Ensure generation config has proper special token ids and disable unsupported SWA
    try:
        if getattr(model, "generation_config", None) is not None:
            if getattr(model.generation_config, "pad_token_id", None) is None and getattr(tok, "pad_token_id", None) is not None:
                model.generation_config.pad_token_id = tok.pad_token_id  # type: ignore[attr-defined]
            if getattr(model.generation_config, "eos_token_id", None) is None and getattr(tok, "eos_token_id", None) is not None:
                model.generation_config.eos_token_id = tok.eos_token_id  # type: ignore[attr-defined]
    except Exception:
        pass
    try:
        # Some backends log warnings for sliding window attention with SDPA; disable if present
        if getattr(getattr(model, "config", object()), "sliding_window", None) is not None:
            model.config.sliding_window = None  # type: ignore[attr-defined]
    except Exception:
        pass
    # Skip custom LoRA injection if a PEFT adapter is attached unless explicitly forced
    should_inject = bool(cfg.get("force_lora_inject", False))
    is_peft_model = False
    try:
        from peft import PeftModel  # type: ignore
        if isinstance(model, PeftModel):
            is_peft_model = True
            # Only inject custom LoRA if explicitly forced
            should_inject = bool(cfg.get("force_lora_inject", False))
    except Exception:
        is_peft_model = False
    if should_inject:
        model = _inject_lora(model, r=int(cfg.get("adapter_r", 8)), alpha=int(cfg.get("adapter_alpha", 16)))
        model = _patch_forwards(model)
    # Configure trainable parameters:
    # - If we injected custom LoRA: train only A/B; freeze base
    # - If using PEFT adapter: respect PEFT's own trainable flags
    # - Otherwise: freeze all (inference-only)
    if should_inject:
        for n, p in model.named_parameters():
            if ".A" in n or ".B" in n:
                p.requires_grad_(True)
            else:
                p.requires_grad_(False)
    elif is_peft_model:
        pass  # leave PEFT-configured trainable params as-is
    else:
        for _, p in model.named_parameters():
            p.requires_grad_(False)
    # Auto-inject LoRA if no trainable params (gated)
    try:
        has_trainable = any(p.requires_grad for _, p in model.named_parameters())
        if bool(cfg.get("auto_lora_if_frozen", False)) and (not has_trainable):
            model = _inject_lora(model, r=int(cfg.get("adapter_r", 8)), alpha=int(cfg.get("adapter_alpha", 16)))
            model = _patch_forwards(model)
            for n, p in model.named_parameters():
                if ".A" in n or ".B" in n:
                    p.requires_grad_(True)
                else:
                    p.requires_grad_(False)
    except Exception:
        pass
    return tok, model

def save_adapters(model: torch.nn.Module, path: str = "adapters.pt") -> None:
    # Ensure target directory exists for checkpoints
    try:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    adap = {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad}
    torch.save(adap, path)
    # Sidecar JSON with telemetry
    try:
        sidecar = str(path) + ".json"
        trainable_params = int(sum(int(p.numel()) for p in model.parameters() if p.requires_grad))
        # Try to infer adapter rank from any LoRA A matrix
        adapter_r = None
        try:
            for _n, _m in model.named_modules():
                if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                    adapter_r = int(getattr(_m, "A").shape[0])  # type: ignore[attr-defined]
                    break


    elif is_peft_model:
        pass  # leave PEFT-configured trainable params as-is
    else:
        for _, p in model.named_parameters():
            p.requires_grad_(False)
    # Auto-inject LoRA if no trainable params (gated)
    try:
        has_trainable = any(p.requires_grad for _, p in model.named_parameters())
        if bool(cfg.get("auto_lora_if_frozen", False)) and (not has_trainable):
            model = _inject_lora(model, r=int(cfg.get("adapter_r", 8)), alpha=int(cfg.get("adapter_alpha", 16)))
            model = _patch_forwards(model)
            for n, p in model.named_parameters():
                if ".A" in n or ".B" in n:
                    p.requires_grad_(True)
                else:
                    p.requires_grad_(False)
    except Exception:
        pass
    return tok, model

def save_adapters(model: torch.nn.Module, path: str = "adapters.pt") -> None:
    # Ensure target directory exists for checkpoints
    try:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    adap = {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad}
    torch.save(adap, path)
    # Sidecar JSON with telemetry
    try:
        sidecar = str(path) + ".json"
        trainable_params = int(sum(int(p.numel()) for p in model.parameters() if p.requires_grad))
        # Try to infer adapter rank from any LoRA A matrix
        adapter_r = None
        try:
            for _n, _m in model.named_modules():
                if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                    adapter_r = int(getattr(_m, "A").shape[0])  # type: ignore[attr-defined]
                    break
        except Exception:
            adapter_r = None
        with open(sidecar, "w", encoding="utf-8") as f:
            json.dump({"trainable_params": trainable_params, "adapter_r": adapter_r}, f, ensure_ascii=False)
    except Exception:
        pass

def load_adapters(model: torch.nn.Module, path: str) -> int:
    """Load adapter weights saved by save_adapters into matching trainable params.

    Returns number of parameters loaded.
    """
    # Load adapters onto model device when possible; fallback to CPU
    try:
        dev = next(model.parameters()).device
    except Exception:
        dev = "cpu"
    try:
        state = torch.load(path, map_location=dev, weights_only=True)
    except Exception:
        try:
            state = torch.load(path, map_location="cpu", weights_only=True)
        except Exception:
            return 0
    if not isinstance(state, dict):


    torch.save(adap, path)
    # Sidecar JSON with telemetry
    try:
        sidecar = str(path) + ".json"
        trainable_params = int(sum(int(p.numel()) for p in model.parameters() if p.requires_grad))
        # Try to infer adapter rank from any LoRA A matrix
        adapter_r = None
        try:
            for _n, _m in model.named_modules():
                if isinstance(_m, torch.nn.Linear) and hasattr(_m, "A"):
                    adapter_r = int(getattr(_m, "A").shape[0])  # type: ignore[attr-defined]
                    break
        except Exception:
            adapter_r = None
        with open(sidecar, "w", encoding="utf-8") as f:
            json.dump({"trainable_params": trainable_params, "adapter_r": adapter_r}, f, ensure_ascii=False)
    except Exception:
        pass

def load_adapters(model: torch.nn.Module, path: str) -> int:
    """Load adapter weights saved by save_adapters into matching trainable params.

    Returns number of parameters loaded.
    """
    # Load adapters onto model device when possible; fallback to CPU
    try:
        dev = next(model.parameters()).device
    except Exception:
        dev = "cpu"
    try:
        state = torch.load(path, map_location=dev, weights_only=True)
    except Exception:
        try:
            state = torch.load(path, map_location="cpu", weights_only=True)
        except Exception:
            return 0
    if not isinstance(state, dict):
        return 0
    loaded = 0
    for n, p in model.named_parameters():
        try:
            if n in state:
                src = state[n]
                if isinstance(src, torch.Tensor) and tuple(src.shape) == tuple(p.shape):
                    with torch.no_grad():
                        p.data.copy_(src.to(p.device, dtype=p.dtype))
                    loaded += 1
        except Exception:
            continue
    return loaded


from .common_imports import *

class EarlyExitHead(torch.nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, 1)

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        # hidden_state: [B, H]
        return torch.sigmoid(self.classifier(hidden_state))

def generate_with_early_exit(
    tok,
    model,
    prompt: str,
    head: EarlyExitHead,
    *,
    layer_index: int = -1,
    threshold: float = 0.9,
    max_new_tokens: int = 128,
    temperature: float = 0.2,
    top_p: float = 0.9,
) -> str:
    device = next(model.parameters()).device
    enc = tok(prompt, return_tensors="pt")
    input_ids = enc["input_ids"].to(device)
    attn = enc.get("attention_mask", torch.ones_like(input_ids)).to(device)
    out_ids = [input_ids]
    steps = 0
    head = head.to(device).eval()
    while steps < int(max_new_tokens):
        with torch.inference_mode():
            out = model(input_ids=input_ids, attention_mask=attn, output_hidden_states=True)
            logits = out.logits[:, -1, :]
            # sample next token
            probs = torch.softmax(logits / max(1e-6, float(temperature)), dim=-1)
            if top_p is not None:
                sorted_probs, sorted_indices = torch.sort(probs, descending=True)
                cumsum = torch.cumsum(sorted_probs, dim=-1)
                cutoff = (cumsum > float(top_p)).float().argmax(dim=-1)
                mask = torch.zeros_like(probs).scatter_(1, sorted_indices, 1.0)
                # retain tokens up to cutoff
                retain = torch.zeros_like(probs)
                for i in range(probs.size(0)):
                    k = int(cutoff[i].item())
                    retain[i, sorted_indices[i, : max(1, k + 1)]] = 1.0
                probs = probs * retain
                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-8)
            next_id = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_id], dim=1)
            attn = torch.cat([attn, torch.ones_like(next_id)], dim=1)
            out_ids.append(next_id)
            steps += 1
            # early-exit check on chosen layer (pooled last token state)
            hs = out.hidden_states[layer_index][:, -1, :]
            p_halt = head(hs).squeeze(-1)
            if float(p_halt[0].item()) >= float(threshold):
                break
    gen = torch.cat(out_ids, dim=1)
    text = tok.decode(gen[0][enc["input_ids"].shape[1] :], skip_special_tokens=True)
    return text


EarlyExitHead

from .common_imports import *

class EarlyExitHead(torch.nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, 1)

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        # hidden_state: [B, H]
        return torch.sigmoid(self.classifier(hidden_state))

def generate_with_early_exit(
    tok,
    model,
    prompt: str,
    head: EarlyExitHead,
    *,
    layer_index: int = -1,
    threshold: float = 0.9,
    max_new_tokens: int = 128,
    temperature: float = 0.2,
    top_p: float = 0.9,
) -> str:
    device = next(model.parameters()).device
    enc = tok(prompt, return_tensors="pt")
    input_ids = enc["input_ids"].to(device)
    attn = enc.get("attention_mask", torch.ones_like(input_ids)).to(device)
    out_ids = [input_ids]
    steps = 0


from .common_imports import *

class EarlyExitHead(torch.nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, 1)

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        # hidden_state: [B, H]
        return torch.sigmoid(self.classifier(hidden_state))

def generate_with_early_exit(
    tok,
    model,
    prompt: str,
    head: EarlyExitHead,
    *,
    layer_index: int = -1,
    threshold: float = 0.9,
    max_new_tokens: int = 128,
    temperature: float = 0.2,
    top_p: float = 0.9,
) -> str:
    device = next(model.parameters()).device
    enc = tok(prompt, return_tensors="pt")


from .common_imports import *
try:
    from .tasks_foundry import _REGISTRY as _FOUNDRY_REG, _mk_tests as _FOUNDRY_MK, TaskTemplate as _FT
except Exception:
    _FOUNDRY_REG = []  # type: ignore

_PRIMS = [
    {
        "name": "sum_to_n",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(1, 5000)]},
        "ref": lambda n: n * (n + 1) // 2,
        "desc": "Return the sum of integers from 1 to n (inclusive).",
    },
    {
        "name": "is_prime",
        "sig": "def solve(n:int)->bool:",
        "gen": lambda: {"args": [__import__("random").randint(0, 5000)]},
        "ref": lambda n: n > 1 and all(n % d for d in range(2, int(n ** 0.5) + 1)),
        "desc": "Return True if n is prime, else False.",
    },
    {
        "name": "rev_str",
        "sig": "def solve(s:str)->str:",
        "gen": lambda: {"args": ["".join(__import__("random").choice("abcxyz") for _ in range(__import__("random").randint(1, 20)))]},
        "ref": lambda s: s[::-1],
        "desc": "Return the reverse of the input string s.",
    },
]

# Additional very simple primitives for level 0

_PRIMS_EASY = [
    {
        "name": "identity_int",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: n,
        "desc": "Return n unchanged.",
    },
    {
        "name": "add_one",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: n + 1,
        "desc": "Return n + 1.",
    },
    {
        "name": "double",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: n * 2,
        "desc": "Return 2 * n.",
    },
    {
        "name": "square",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-500, 500)]},
        "ref": lambda n: n * n,
        "desc": "Return n squared.",
    },
    {
        "name": "abs_int",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: abs(n),
        "desc": "Return the absolute value of n.",
    },
    {
        "name": "negate",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: -n,
        "desc": "Return -n.",
    },
    {
        "name": "is_even",
        "sig": "def solve(n:int)->bool:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: (n % 2) == 0,
        "desc": "Return True if n is even, else False.",
    },
]

def sample_task(level: int) -> Dict[str, Any]:
    import random

    # Primitive-only mode: skip foundry/task templates entirely
    try:
        prim_only = str(os.environ.get("SELFPLAY_PRIMITIVE_ONLY", "")).strip() not in ("", "0", "false", "False")
        prim_only = prim_only or (str(os.environ.get("SELFPLAY_BLOCK_IMPORTS", "")).strip() not in ("", "0", "false", "False"))
    except Exception:
        prim_only = False

    # Prefer foundry templates when available and not in primitive-only mode
    try:
        if _FOUNDRY_REG and (not prim_only) and level >= 0:
            # Optional tag filter from env
            tag_env = str(os.environ.get("SELFPLAY_TASK_TAGS", "")).strip()
            want_tags = [s.strip() for s in tag_env.split(",") if s.strip()] if tag_env else []
            # Filter by difficulty: allow within level+2
            def _ok(t: _FT) -> bool:  # type: ignore
                diff_ok = getattr(t, "difficulty", 0) <= min(5, int(level) + 2)
                tag_ok = (not want_tags) or any((tg in getattr(t, "tags", [])) for tg in want_tags)
                return bool(diff_ok and tag_ok)
            cand = [t for t in _FOUNDRY_REG if _ok(t)]
            if not cand:
                cand = list(_FOUNDRY_REG)
            # Seeded RNG for reproducibility when SELFPLAY_SEED is set
            try:
                # Episode-scoped seed has priority for reproducibility across steps
                seed_env = os.environ.get("SELFPLAY_SEED_EP")
                if seed_env is None or str(seed_env).strip() == "":
                    seed_env = os.environ.get("SELFPLAY_SEED")
                seed_val = int(seed_env) if seed_env is not None and str(seed_env).strip() != "" else None
            except Exception:
                seed_val = None
            rng = random.Random(seed_val) if seed_val is not None else random.Random()
            T = random.choice(cand)
            inputs, oracle = T.sampler(rng)
            # Base tests
            tests = _FOUNDRY_MK(inputs, oracle)
            # Metamorphic expansion
            try:
                for mr in getattr(T, "metamorphics", []) or []:
                    try:
                        for (inp, out) in mr(inputs, oracle):
                            tests.append({"inp": list(inp), "out": out})
                    except Exception:
                        continue
            except Exception:
                pass
            # cap tests ~8..32
            if len(tests) < 8:
                base = list(tests)
                while len(tests) < 8 and base:
                    tests.append(base[len(tests) % len(base)])
            tests = tests[:32]
            return {"name": T.name, "signature": T.signature, "tests": tests, "description": getattr(T, "description", T.name)}
    except Exception:
        pass

    # Fallback to primitives
    pool = (_PRIMS_EASY if level <= 0 else (_PRIMS_EASY + _PRIMS))
    t = random.choice(pool)
    cases = [t["gen"]() for _ in range(8)]
    tests = [{"inp": c["args"], "out": t["ref"](*c["args"])} for c in cases]
    return {"name": t["name"], "signature": t["sig"], "tests": tests, "description": t.get("desc", t["name"]).strip()}

def render_prompt(spec: Dict[str, Any]) -> str:
    t = spec["name"]
    sig = spec["signature"]
    try:
        prim_mode = str(os.environ.get("SELFPLAY_PRIMITIVE_ONLY", "")).strip() not in ("", "0", "false", "False")
        prim_mode = prim_mode or (str(os.environ.get("SELFPLAY_BLOCK_IMPORTS", "")).strip() not in ("", "0", "false", "False"))
    except Exception:
        prim_mode = False
    if prim_mode:
        return (
            "You are a senior Python engineer. Output ONLY code.\n"
            f"Task: {t}\n"
            f"Signature:\n```python\n{sig}\n```\n"
            "Rules:\n"
            "- Primitive-only: leetcode style, NO imports. Use arithmetic, control-flow, slicing, comprehensions, and pure python.\n"
            "Checklist: exact signature, no preamble/backticks, no imports.\n"
            "Emit only the function code block.\nAnswer:"
        )
    else:
        return (
            "You are a senior Python engineer. Output ONLY code.\n"
            f"Task: {t}\n"
            f"Signature:\n```python\n{sig}\n```\n"
            "Rules:\n"
            "- Self-contained solution. If you need helpers, define them or import from the Python standard library (e.g., 'from functools import reduce').\n"
            "- Allowed: standard library imports. Disallowed: external packages, file/network I/O, subprocess, side effects.\n"
            "- If you use stdlib helpers (e.g., reduce, lru_cache, Counter), include the import.\n"
            "- Prefer pure function behavior and deterministic outputs.\n"
            "Disallowed (zero credit): any prefix like 'python', markdown fences, wrong function name, comments, or explanations.\n"
            "Checklist: exact signature, no preamble/backticks, include stdlib imports if used.\n"
            "Emit only code: optional imports then the single function matching the signature.\nAnswer:"
        )


        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: abs(n),
        "desc": "Return the absolute value of n.",
    },
    {
        "name": "negate",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: -n,
        "desc": "Return -n.",
    },
    {
        "name": "is_even",
        "sig": "def solve(n:int)->bool:",
        "gen": lambda: {"args": [__import__("random").randint(-1000, 1000)]},
        "ref": lambda n: (n % 2) == 0,
        "desc": "Return True if n is even, else False.",
    },
]

def sample_task(level: int) -> Dict[str, Any]:
    import random

    # Primitive-only mode: skip foundry/task templates entirely
    try:
        prim_only = str(os.environ.get("SELFPLAY_PRIMITIVE_ONLY", "")).strip() not in ("", "0", "false", "False")
        prim_only = prim_only or (str(os.environ.get("SELFPLAY_BLOCK_IMPORTS", "")).strip() not in ("", "0", "false", "False"))
    except Exception:
        prim_only = False

    # Prefer foundry templates when available and not in primitive-only mode
    try:
        if _FOUNDRY_REG and (not prim_only) and level >= 0:
            # Optional tag filter from env
            tag_env = str(os.environ.get("SELFPLAY_TASK_TAGS", "")).strip()
            want_tags = [s.strip() for s in tag_env.split(",") if s.strip()] if tag_env else []
            # Filter by difficulty: allow within level+2
            def _ok(t: _FT) -> bool:  # type: ignore
                diff_ok = getattr(t, "difficulty", 0) <= min(5, int(level) + 2)
                tag_ok = (not want_tags) or any((tg in getattr(t, "tags", [])) for tg in want_tags)
                return bool(diff_ok and tag_ok)
            cand = [t for t in _FOUNDRY_REG if _ok(t)]
            if not cand:
                cand = list(_FOUNDRY_REG)
            # Seeded RNG for reproducibility when SELFPLAY_SEED is set
            try:
                # Episode-scoped seed has priority for reproducibility across steps
                seed_env = os.environ.get("SELFPLAY_SEED_EP")
                if seed_env is None or str(seed_env).strip() == "":
                    seed_env = os.environ.get("SELFPLAY_SEED")
                seed_val = int(seed_env) if seed_env is not None and str(seed_env).strip() != "" else None
            except Exception:
                seed_val = None
            rng = random.Random(seed_val) if seed_val is not None else random.Random()
            T = random.choice(cand)
            inputs, oracle = T.sampler(rng)
            # Base tests
            tests = _FOUNDRY_MK(inputs, oracle)
            # Metamorphic expansion
            try:
                for mr in getattr(T, "metamorphics", []) or []:
                    try:
                        for (inp, out) in mr(inputs, oracle):
                            tests.append({"inp": list(inp), "out": out})
                    except Exception:
                        continue
            except Exception:
                pass
            # cap tests ~8..32
            if len(tests) < 8:
                base = list(tests)
                while len(tests) < 8 and base:
                    tests.append(base[len(tests) % len(base)])
            tests = tests[:32]
            return {"name": T.name, "signature": T.signature, "tests": tests, "description": getattr(T, "description", T.name)}
    except Exception:
        pass

    # Fallback to primitives
    pool = (_PRIMS_EASY if level <= 0 else (_PRIMS_EASY + _PRIMS))
    t = random.choice(pool)
    cases = [t["gen"]() for _ in range(8)]
    tests = [{"inp": c["args"], "out": t["ref"](*c["args"])} for c in cases]
    return {"name": t["name"], "signature": t["sig"], "tests": tests, "description": t.get("desc", t["name"]).strip()}

def render_prompt(spec: Dict[str, Any]) -> str:
    t = spec["name"]
    sig = spec["signature"]
    try:
        prim_mode = str(os.environ.get("SELFPLAY_PRIMITIVE_ONLY", "")).strip() not in ("", "0", "false", "False")
        prim_mode = prim_mode or (str(os.environ.get("SELFPLAY_BLOCK_IMPORTS", "")).strip() not in ("", "0", "false", "False"))
    except Exception:
        prim_mode = False
    if prim_mode:
        return (
            "You are a senior Python engineer. Output ONLY code.\n"
            f"Task: {t}\n"
            f"Signature:\n```python\n{sig}\n```\n"
            "Rules:\n"
            "- Primitive-only: leetcode style, NO imports. Use arithmetic, control-flow, slicing, comprehensions, and pure python.\n"
            "Checklist: exact signature, no preamble/backticks, no imports.\n"
            "Emit only the function code block.\nAnswer:"
        )


            except Exception:
                pass
            # cap tests ~8..32
            if len(tests) < 8:
                base = list(tests)
                while len(tests) < 8 and base:
                    tests.append(base[len(tests) % len(base)])
            tests = tests[:32]
            return {"name": T.name, "signature": T.signature, "tests": tests, "description": getattr(T, "description", T.name)}
    except Exception:
        pass

    # Fallback to primitives
    pool = (_PRIMS_EASY if level <= 0 else (_PRIMS_EASY + _PRIMS))
    t = random.choice(pool)
    cases = [t["gen"]() for _ in range(8)]
    tests = [{"inp": c["args"], "out": t["ref"](*c["args"])} for c in cases]
    return {"name": t["name"], "signature": t["sig"], "tests": tests, "description": t.get("desc", t["name"]).strip()}

def render_prompt(spec: Dict[str, Any]) -> str:
    t = spec["name"]
    sig = spec["signature"]
    try:
        prim_mode = str(os.environ.get("SELFPLAY_PRIMITIVE_ONLY", "")).strip() not in ("", "0", "false", "False")
        prim_mode = prim_mode or (str(os.environ.get("SELFPLAY_BLOCK_IMPORTS", "")).strip() not in ("", "0", "false", "False"))
    except Exception:
        prim_mode = False
    if prim_mode:
        return (
            "You are a senior Python engineer. Output ONLY code.\n"
            f"Task: {t}\n"
            f"Signature:\n```python\n{sig}\n```\n"
            "Rules:\n"
            "- Primitive-only: leetcode style, NO imports. Use arithmetic, control-flow, slicing, comprehensions, and pure python.\n"
            "Checklist: exact signature, no preamble/backticks, no imports.\n"
            "Emit only the function code block.\nAnswer:"
        )
    else:
        return (
            "You are a senior Python engineer. Output ONLY code.\n"
            f"Task: {t}\n"
            f"Signature:\n```python\n{sig}\n```\n"
            "Rules:\n"
            "- Self-contained solution. If you need helpers, define them or import from the Python standard library (e.g., 'from functools import reduce').\n"
            "- Allowed: standard library imports. Disallowed: external packages, file/network I/O, subprocess, side effects.\n"
            "- If you use stdlib helpers (e.g., reduce, lru_cache, Counter), include the import.\n"
            "- Prefer pure function behavior and deterministic outputs.\n"
            "Disallowed (zero credit): any prefix like 'python', markdown fences, wrong function name, comments, or explanations.\n"
            "Checklist: exact signature, no preamble/backticks, include stdlib imports if used.\n"
            "Emit only code: optional imports then the single function matching the signature.\nAnswer:"
        )