    def ref(arr: List[int]) -> int:
        return sum(x for x in (a + 1 for a in arr) if x % 2 == 0)
    return xs, ref

register(
    TaskTemplate(
        name="pipeline_map_filter_reduce",
        signature="def solve(arr:list[int])->int:",
        sampler=_pipeline_sampler,
        tags=["array", "map", "filter", "reduce", "composed"],
        difficulty=2,
    )
)

# --- GCD (Euclid) with metamorphics --- #

def _gcd_sampler(rng: Rand) -> Tuple[List[List[Any]], Callable[..., Any]]:
    xs: List[List[Any]] = [[rng.randint(0, 10_000), rng.randint(0, 10_000)] for _ in range(12)]
    def _g(a: int, b: int) -> int:
        a = int(a); b = int(b)
        while b:
            a, b = b, a % b
        return abs(a)
    return xs, _g

def _gcd_mr(inputs: List[List[Any]], ref: Callable[..., Any]) -> List[Tuple[List[Any], Any]]:
    extra: List[Tuple[List[Any], Any]] = []
    for a, b in inputs[:8]:
        a = int(a); b = int(b)
        extra.append(([b, a], ref(b, a)))  # symmetry
        extra.append(([a, 0], abs(a)))     # boundary
        if a and b:
            k = random.randint(2, 5)
            extra.append(([a * k, b * k], ref(a * k, b * k)))  # scaling
    return extra

register(
    TaskTemplate(
        name="gcd_euclid",
        signature="def solve(a:int,b:int)->int:",
        sampler=_gcd_sampler,


from .common_imports import *

def props_rev_str(f) -> bool:
    try:
        import string as _s
        import random as _r
        for _ in range(64):
            s = "".join(_r.choice(_s.printable[:90]) for __ in range(_r.randint(0, 64)))
            r = f(s)
            if (r[::-1] != s) or (len(r) != len(s)):
                return False
            if f(r) != s:
                return False
        return True
    except Exception:
        return False

def tests_sum_to_n(f) -> bool:
    try:
        for n in [0, 1, 2, 5, 100, 10000]:
            if f(n) != n * (n + 1) // 2:
                return False
        return True
    except Exception:
        return False

def tests_rev_str_involution(f) -> bool:
    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:
        def _isqrt(x: int) -> int:
            return int((x) ** 0.5)
    try:
        small = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29}
        for n in range(0, 40):
            want = n in small or (n > 1 and all(n % d for d in range(2, _isqrt(n) + 1)))
            if f(n) != want:
                return False
        return True
    except Exception:
        return False

def tests_pipeline_map_filter_reduce(f) -> bool:
    try:
        import random as _r
        for _ in range(64):
            arr = [_r.randint(-8, 8) for __ in range(32)]
            want = sum(x * x for x in arr if x % 2 == 0)
            got = f(arr)
            if got != want:
                return False
        return True
    except Exception:
        return False

TASK_TESTS: Dict[str, Any] = {
    "sum_to_n": tests_sum_to_n,
    "rev_str_involution": tests_rev_str_involution,
    "is_prime_adv": tests_is_prime_adv,
    "pipeline_map_filter_reduce": tests_pipeline_map_filter_reduce,
}


from .common_imports import *

def props_rev_str(f) -> bool:
    try:
        import string as _s
        import random as _r
        for _ in range(64):
            s = "".join(_r.choice(_s.printable[:90]) for __ in range(_r.randint(0, 64)))
            r = f(s)
            if (r[::-1] != s) or (len(r) != len(s)):
                return False
            if f(r) != s:
                return False
        return True
    except Exception:
        return False

def tests_sum_to_n(f) -> bool:
    try:
        for n in [0, 1, 2, 5, 100, 10000]:
            if f(n) != n * (n + 1) // 2:
                return False
        return True
    except Exception:
        return False

def tests_rev_str_involution(f) -> bool:
    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:


from .common_imports import *

def props_rev_str(f) -> bool:
    try:
        import string as _s
        import random as _r
        for _ in range(64):
            s = "".join(_r.choice(_s.printable[:90]) for __ in range(_r.randint(0, 64)))
            r = f(s)
            if (r[::-1] != s) or (len(r) != len(s)):
                return False
            if f(r) != s:
                return False
        return True
    except Exception:
        return False

def tests_sum_to_n(f) -> bool:
    try:
        for n in [0, 1, 2, 5, 100, 10000]:
            if f(n) != n * (n + 1) // 2:
                return False
        return True
    except Exception:
        return False

def tests_rev_str_involution(f) -> bool:
    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:
        def _isqrt(x: int) -> int:
            return int((x) ** 0.5)
    try:
        small = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29}
        for n in range(0, 40):
            want = n in small or (n > 1 and all(n % d for d in range(2, _isqrt(n) + 1)))
            if f(n) != want:
                return False
        return True
    except Exception:


            r = f(s)
            if (r[::-1] != s) or (len(r) != len(s)):
                return False
            if f(r) != s:
                return False
        return True
    except Exception:
        return False

def tests_sum_to_n(f) -> bool:
    try:
        for n in [0, 1, 2, 5, 100, 10000]:
            if f(n) != n * (n + 1) // 2:
                return False
        return True
    except Exception:
        return False

def tests_rev_str_involution(f) -> bool:
    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:
        def _isqrt(x: int) -> int:
            return int((x) ** 0.5)
    try:
        small = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29}
        for n in range(0, 40):
            want = n in small or (n > 1 and all(n % d for d in range(2, _isqrt(n) + 1)))
            if f(n) != want:
                return False
        return True
    except Exception:
        return False

def tests_pipeline_map_filter_reduce(f) -> bool:


                return False
        return True
    except Exception:
        return False

def tests_sum_to_n(f) -> bool:
    try:
        for n in [0, 1, 2, 5, 100, 10000]:
            if f(n) != n * (n + 1) // 2:
                return False
        return True
    except Exception:
        return False

def tests_rev_str_involution(f) -> bool:
    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:
        def _isqrt(x: int) -> int:
            return int((x) ** 0.5)
    try:
        small = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29}
        for n in range(0, 40):
            want = n in small or (n > 1 and all(n % d for d in range(2, _isqrt(n) + 1)))
            if f(n) != want:
                return False
        return True
    except Exception:
        return False

def tests_pipeline_map_filter_reduce(f) -> bool:
    try:
        import random as _r
        for _ in range(64):
            arr = [_r.randint(-8, 8) for __ in range(32)]
            want = sum(x * x for x in arr if x % 2 == 0)
            got = f(arr)
            if got != want:
                return False
        return True
    except Exception:
        return False

TASK_TESTS: Dict[str, Any] = {
    "sum_to_n": tests_sum_to_n,
    "rev_str_involution": tests_rev_str_involution,
    "is_prime_adv": tests_is_prime_adv,


    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:
        def _isqrt(x: int) -> int:
            return int((x) ** 0.5)
    try:
        small = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29}
        for n in range(0, 40):
            want = n in small or (n > 1 and all(n % d for d in range(2, _isqrt(n) + 1)))
            if f(n) != want:
                return False
        return True
    except Exception:
        return False

def tests_pipeline_map_filter_reduce(f) -> bool:
    try:
        import random as _r
        for _ in range(64):
            arr = [_r.randint(-8, 8) for __ in range(32)]
            want = sum(x * x for x in arr if x % 2 == 0)
            got = f(arr)
            if got != want:
                return False
        return True
    except Exception:
        return False

TASK_TESTS: Dict[str, Any] = {
    "sum_to_n": tests_sum_to_n,
    "rev_str_involution": tests_rev_str_involution,
    "is_prime_adv": tests_is_prime_adv,
    "pipeline_map_filter_reduce": tests_pipeline_map_filter_reduce,
}


def tests_sum_to_n(f) -> bool:
    try:
        for n in [0, 1, 2, 5, 100, 10000]:
            if f(n) != n * (n + 1) // 2:
                return False
        return True
    except Exception:
        return False

def tests_rev_str_involution(f) -> bool:
    return props_rev_str(f)

def tests_is_prime_adv(f) -> bool:
    try:
        from math import isqrt as _isqrt
    except Exception:
        def _isqrt(x: int) -> int:
            return int((x) ** 0.5)
    try:
        small = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29}
        for n in range(0, 40):
            want = n in small or (n > 1 and all(n % d for d in range(2, _isqrt(n) + 1)))
            if f(n) != want:
                return False
        return True
    except Exception:
        return False

def tests_pipeline_map_filter_reduce(f) -> bool:
    try:
        import random as _r
        for _ in range(64):
            arr = [_r.randint(-8, 8) for __ in range(32)]
            want = sum(x * x for x in arr if x % 2 == 0)
            got = f(arr)
            if got != want:
                return False


from .common_imports import *
from .generation import generate_text

def heal_code(tok, model, prompt: str, failing_code: str, failure_info: Dict[str, Any], max_new_tokens: int = 256) -> str:
     """Construct a healing prompt that explains why the submission failed and
    reminds the model of exact signature/output constraints.

    We include:
    - First exception string (if any)
    - First input/output mismatch example (if any)
    - Aggregate stats (passed/total, timeouts)
    - A restatement of the exact required signature parsed from the task prompt
    """
    # Extract concise failure signal
    try:
        err = failure_info.get("first_failure") or failure_info.get("err") or ""
    except Exception:
        err = ""

    # Aggregate stats if available
    stats = ""
    try:
        passed = failure_info.get("passed")
        total = failure_info.get("total")
        timeouts = failure_info.get("timeouts")
        have_counts = (passed is not None) or (total is not None) or (timeouts is not None)
        if have_counts:
            ps = int(passed) if passed is not None else 0
            tt = int(total) if total is not None else 0
            to = int(timeouts) if timeouts is not None else 0
            stats = f"\n# Test summary: passed={ps}/{tt}, timeouts={to}"
    except Exception:
        stats = ""

    # Include a concrete failing example when available
    ex = ""
    try:
        mm = failure_info.get("first_mismatch")
        if isinstance(mm, dict):
            ex = ("\n# Example failing test case:\n"
                  f"# input: {mm.get('inp')}\n"
                  f"# expected: {mm.get('want')}\n"
                  f"# got: {mm.get('got')}\n")
    except Exception:
        ex = ""

    # Try to restate the exact signature from the task prompt to avoid drift
    def _extract_sig_from_prompt(txt: str) -> str | None:
        try:
            m = re.search(r"Signature:\s*```python\s*(def\s+solve[^\n]+)\s*```", txt, flags=re.IGNORECASE)
            if m:
                return str(m.group(1)).strip()
        except Exception:
            pass
        return None

    sig_hint = _extract_sig_from_prompt(prompt or "")
    sig_line = (f"# Exact signature (must match): {sig_hint}\n" if sig_hint else "")

    # Strong, explicit output constraints
    constraints = (
        "# Requirements:\n"
        "# - Output ONLY code. No backticks, no comments, no explanations.\n"
        + ("# - Use the exact signature above without changes.\n" if sig_hint else "") +
        "# - Self-contained. Only Python standard library imports if needed.\n"
        "# - No file/network I/O, no subprocess, no side effects.\n"
    )

    heal_prompt = (
        (prompt or "").rstrip() +
        "\n\n# The code failed. Error info (if any):\n" + str(err) + stats + ex + "\n" +
        sig_line +
        constraints +
        "\n# Current code:\n" + failing_code +
        "\n\n# Provide a corrected version of the function now.\n"
    )
    return generate_text(tok, model, heal_prompt, max_new_tokens=max_new_tokens)

_BAD = re.compile(
    r"("  # Group for alternation
    r"^\s*(import\s+|from\s+\w+\s+import\s+)(os|sys|subprocess)\b|"  # imports of os/sys/subprocess
    r"^\s*(exec|eval)\s*\(|"  # exec/eval at start of line
    r"\bopen\s*\(|"  # file open anywhere
    r"^\s*subprocess\.|"  # direct subprocess usage at SOL
    r"while\s*(True|1)\s*:\s*|"  # infinite while
    r"for\s*\(\s*;\s*;\s*\)\s*:\s*|"  # C-style infinite for
    r"recursionlimit"  # attempts to change recursion limit
    r")",
    re.MULTILINE,
)

def _strip_triple_quoted(code: str) -> str:
    """Remove triple-quoted strings to avoid hiding disallowed constructs inside them."""
    try:
        # Remove both """...""" and '''...'''
        code = re.sub(r"\"\"\"[\s\S]*?\"\"\"", "", code)
        code = re.sub(r"'''[\s\S]*?'''", "", code)
    except Exception:
        pass
    return code

def is_disallowed(code: str) -> bool:
    if len(code) > 2000:
        return True
    try:
        scan = _strip_triple_quoted(code)
    except Exception:
        scan = code
    return bool(_BAD.search(scan))


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

def sft_step(tok, model, prompt: str, target: str, lr: float, steps: int, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:
            alpha = 16
        try:
            _inject_lora(model, r=r, alpha=alpha)
            _patch_forwards(model)
        except Exception:
            pass
    try:
        # Enable only LoRA params (custom or PEFT) and freeze base
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
        # Fallback: if nothing enabled (e.g., PEFT naming mismatch), inject custom LoRA now
        if enabled == 0:
            try:
                r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
            except Exception:
                r = 8
            try:
                alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
            except Exception:
                alpha = 16
            _inject_lora(model, r=r, alpha=alpha)
            _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)
            _set_lora_enabled(model, True)
    except Exception:
        pass
    # Optional MetaOpt integration: default enabled unless explicitly disabled
    _v_meta = os.environ.get("AGI_METAOPT")
    use_meta = True if _v_meta is None else (str(_v_meta).strip().lower() in ("1","true","yes","y","on"))
    meta_opt = None
    if use_meta:
        try:
            from agi_dw.core.metaopt.metaopt_grouped import MixtureMetaOptGrouped  # type: ignore
            from agi_dw.core.metaopt.gate_grouped import GateNetGrouped  # type: ignore
            trainable = [p for p in model.parameters() if p.requires_grad]
            device = next(model.parameters()).device.type
            # Cache and reuse meta optimizer to keep state
            meta_opt = getattr(model, "_sft_meta", None)
            if meta_opt is None:
                gate = GateNetGrouped(n_groups=1, hidden=128)
                meta_opt = MixtureMetaOptGrouped(model, param_groups=[{"params": trainable}], gate=gate, base_lr=lr, base_wd=0.01, device=device)
                setattr(model, "_sft_meta", meta_opt)
        except Exception:
            use_meta = False
            meta_opt = None
    # Fallback/default: AdamW (reuse optimizer across calls to preserve momentum)
    opt = None
    if not use_meta:
        opt = getattr(model, "_sft_opt", None)
        if opt is None:
            opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)
            setattr(model, "_sft_opt", opt)
    # Gradient accumulation
    try:
        grad_accum = max(1, int(os.environ.get("SELFPLAY_GRAD_ACCUM", "1") or 1))
    except Exception:
        grad_accum = 1
    # Optional KL clamp vs frozen baseline
    ref_model = None
    try:
        if bool(os.environ.get("SELFPLAY_KL_CLAMP", "0") == "1"):
            ref_model = copy.deepcopy(model).eval()
            for p in ref_model.parameters():
                p.requires_grad_(False)
    except Exception:
        ref_model = None
    model.train()
    last_loss: float | None = None
    device = next(model.parameters()).device
    _dtype = next(model.parameters()).dtype
    # Pre-tokenize concatenated text once to avoid double tokenization
    text = prompt + "\n" + target
    enc = tok(text, return_tensors="pt")
    enc_prompt = tok(prompt + "\n", return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    labels = enc["input_ids"].clone()
    prompt_len = int(enc_prompt["input_ids"].shape[1])
    labels[:, :prompt_len] = -100
    # Mask EOS tokens from contributing to loss
    try:
        eos_id = getattr(tok, "eos_token_id", None)
        if eos_id is not None:
            labels[labels == int(eos_id)] = -100
    except Exception:
        pass
    # If no valid supervised tokens after shift, skip step to avoid constant loss
    try:
        valid_positions = int(((labels[:, 1:] != -100).sum()).item())
        if valid_positions == 0:
            return None
    except Exception:
        pass
    # Training loop with accumulation
    grad_norm_val: float | None = None
    kl_val: float | None = None
    for it in range(steps):
        if it % grad_accum == 0:
            if opt is not None:
                opt.zero_grad()
        if device.type == "cuda" and _dtype in (getattr(torch, "float16"), getattr(torch, "bfloat16")):
            with _autocast_cuda(_dtype):
                with torch.set_grad_enabled(True):
                    out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                    loss = out.loss
                    # Fallback: if loss is detached, recompute CE from logits
                    try:
                        if not getattr(loss, "requires_grad", True):
                            import torch.nn.functional as F  # type: ignore
                            out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                            logits = out2.logits  # [B, T, V]
                            shift_logits = logits[:, :-1, :].contiguous()
                            shift_labels = labels[:, 1:].contiguous()
                            mask = shift_labels != -100
                            if mask.any():
                                loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                            else:
                                loss = logits.new_tensor(0.0)
                    except Exception:
                        pass
                    # Optional KL clamp
                    if ref_model is not None:
                        with torch.inference_mode():
                            ref = ref_model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                        p = torch.log_softmax(out.logits.float(), dim=-1)
                        q = torch.log_softmax(ref.logits.float(), dim=-1)
                        kl = torch.sum(torch.exp(p) * (p - q), dim=-1).mean()
                        try:
                            kl_val = float(kl.detach().item())
                        except Exception:
                            kl_val = None
                        try:
                            kl_max = float(os.environ.get("SELFPLAY_KL_MAX", "1.0") or 1.0)
                        except Exception:
                            kl_max = 1.0
                        loss = loss + torch.clamp(kl, max=kl_max)
        else:
            with torch.set_grad_enabled(True):
                out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                loss = out.loss
                try:
                    if not getattr(loss, "requires_grad", True):
                        import torch.nn.functional as F  # type: ignore
                        out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                        logits = out2.logits
                        shift_logits = logits[:, :-1, :].contiguous()
                        shift_labels = labels[:, 1:].contiguous()
                        mask = shift_labels != -100
                        if mask.any():
                            loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                        else:
                            loss = logits.new_tensor(0.0)
                except Exception:
                    pass
        # If loss is still detached (no trainable path), skip this micro-step safely
        try:
            if not getattr(loss, "requires_grad", True):
                continue
        except Exception:
            pass
        loss.backward()
        # Step on accumulation interval
        if (it % grad_accum) == (grad_accum - 1):
            if use_meta and meta_opt is not None:
                try:
                    meta_opt.step(float(loss.detach().item()))
                except Exception:
                    pass
                # zero grads for next accumulation window
                for p in model.parameters():
                    if p.grad is not None:
                        p.grad.detach_(); p.grad.zero_()
            else:
                try:
                    clip = float(os.environ.get("SELFPLAY_CLIP_NORM", "1.0") or 1.0)
                except Exception:
                    clip = 1.0
                try:
                    _gn = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip)  # type: ignore
                    try:
                        grad_norm_val = float(getattr(_gn, "item", lambda: _gn)()) if _gn is not None else None  # type: ignore
                    except Exception:
                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    # Log training metrics once per call
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "sft", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val, "kl_to_ref": kl_val}})
    except Exception:
        pass
    return last_loss

def contrastive_step(tok, model, prompt: str, neg: str, lr: float, steps: int, margin: float = 0.1, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:
            alpha = 16
        try:
            _inject_lora(model, r=r, alpha=alpha)
            _patch_forwards(model)
        except Exception:
            pass
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
        if enabled == 0:
            try:
                r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
            except Exception:
                r = 8
            try:
                alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
            except Exception:
                alpha = 16
            _inject_lora(model, r=r, alpha=alpha)
            _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)
            _set_lora_enabled(model, True)
    except Exception:
        pass
    # Optional MetaOpt integration: default enabled unless explicitly disabled
    _v_meta = os.environ.get("AGI_METAOPT")
    use_meta = True if _v_meta is None else (str(_v_meta).strip().lower() in ("1","true","yes","y","on"))
    meta_opt = None
    if use_meta:
        try:
            from agi_dw.core.metaopt.metaopt_grouped import MixtureMetaOptGrouped  # type: ignore
            from agi_dw.core.metaopt.gate_grouped import GateNetGrouped  # type: ignore
            trainable = [p for p in model.parameters() if p.requires_grad]
            device = next(model.parameters()).device.type
            meta_opt = getattr(model, "_ctr_meta", None)
            if meta_opt is None:
                gate = GateNetGrouped(n_groups=1, hidden=128)
                meta_opt = MixtureMetaOptGrouped(model, param_groups=[{"params": trainable}], gate=gate, base_lr=lr, base_wd=0.01, device=device)
                setattr(model, "_ctr_meta", meta_opt)
        except Exception:
            use_meta = False
            meta_opt = None
    opt = None
    if not use_meta:
        opt = getattr(model, "_ctr_opt", None)
        if opt is None:
            opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)
            setattr(model, "_ctr_opt", opt)
    try:
        grad_accum = max(1, int(os.environ.get("SELFPLAY_GRAD_ACCUM", "1") or 1))
    except Exception:
        grad_accum = 1
    model.train()
    last_loss: float | None = None
    device = next(model.parameters()).device
    _dtype = next(model.parameters()).dtype
    # Pre-tokenize once
    text = prompt + "\n" + neg
    enc = tok(text, return_tensors="pt")
    enc_prompt = tok(prompt + "\n", return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    labels = enc["input_ids"].clone()
    prompt_len = int(enc_prompt["input_ids"].shape[1])
    labels[:, :prompt_len] = -100
    # If no valid supervised tokens (e.g., empty neg), skip to avoid no-grad loss
    try:
        if int((labels != -100).sum().item()) == 0:
            return None
    except Exception:
        pass
    try:
        eos_id = getattr(tok, "eos_token_id", None)
        if eos_id is not None:
            labels[labels == int(eos_id)] = -100
    except Exception:
        pass
    # If no valid supervised positions after shift, skip to avoid constant, no-grad loss
    try:
        valid_positions = int(((labels[:, 1:] != -100).sum()).item())
        if valid_positions == 0:
            return None
    except Exception:
        pass
    grad_norm_val: float | None = None
    for it in range(steps):
        if it % grad_accum == 0:
            if opt is not None:
                opt.zero_grad()
        if device.type == "cuda" and _dtype in (getattr(torch, "float16"), getattr(torch, "bfloat16")):
            with _autocast_cuda(_dtype):
                with torch.set_grad_enabled(True):
                    out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                    # Encourage model to keep negative likelihood below margin by making loss zero once margin reached
                    base_loss = out.loss
                    try:
                        if not getattr(base_loss, "requires_grad", True):
                            import torch.nn.functional as F  # type: ignore
                            out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                            logits = out2.logits
                            shift_logits = logits[:, :-1, :].contiguous()
                            shift_labels = labels[:, 1:].contiguous()
                            mask = shift_labels != -100
                            if mask.any():
                                base_loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                            else:
                                base_loss = logits.new_tensor(0.0)
                    except Exception:
                        pass
                    loss = torch.relu(margin - base_loss)
        else:
            with torch.set_grad_enabled(True):
                out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                base_loss = out.loss
                try:
                    if not getattr(base_loss, "requires_grad", True):
                        import torch.nn.functional as F  # type: ignore
                        out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                        logits = out2.logits
                        shift_logits = logits[:, :-1, :].contiguous()
                        shift_labels = labels[:, 1:].contiguous()
                        mask = shift_labels != -100
                        if mask.any():
                            base_loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                        else:
                            base_loss = logits.new_tensor(0.0)
                except Exception:
                    pass
                loss = torch.relu(margin - base_loss)
        try:
            if not getattr(loss, "requires_grad", True):
                continue
        except Exception:
            pass
        loss.backward()
        if (it % grad_accum) == (grad_accum - 1):
            if use_meta and meta_opt is not None:
                try:
                    meta_opt.step(float(loss.detach().item()))
                except Exception:
                    pass
                for p in model.parameters():
                    if p.grad is not None:
                        p.grad.detach_(); p.grad.zero_()
            else:
                try:
                    clip = float(os.environ.get("SELFPLAY_CLIP_NORM", "1.0") or 1.0)
                except Exception:
                    clip = 1.0
                try:
                    _gn = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip)  # type: ignore
                    try:
                        grad_norm_val = float(getattr(_gn, "item", lambda: _gn)()) if _gn is not None else None  # type: ignore
                    except Exception:
                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "contrastive", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val}})
    except Exception:
        pass
    return last_loss

def _autocast_cuda(dtype: torch.dtype):
    """Return a CUDA autocast context, compatible with both torch.amp.autocast and torch.autocast."""
    try:
        return torch.amp.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]
    except AttributeError:
        return torch.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]

# ========================= New training utilities =========================

def _enable_lora_train_only(model: torch.nn.Module) -> None:
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loram
# ... truncated ...


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

def sft_step(tok, model, prompt: str, target: str, lr: float, steps: int, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:
            alpha = 16
        try:
            _inject_lora(model, r=r, alpha=alpha)
            _patch_forwards(model)
        except Exception:
            pass
    try:
        # Enable only LoRA params (custom or PEFT) and freeze base
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
        # Fallback: if nothing enabled (e.g., PEFT naming mismatch), inject custom LoRA now
        if enabled == 0:
            try:
                r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
            except Exception:
                r = 8
            try:
                alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
            except Exception:
                alpha = 16
            _inject_lora(model, r=r, alpha=alpha)
            _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)
            _set_lora_enabled(model, True)
    except Exception:
        pass
    # Optional MetaOpt integration: default enabled unless explicitly disabled
    _v_meta = os.environ.get("AGI_METAOPT")
    use_meta = True if _v_meta is None else (str(_v_meta).strip().lower() in ("1","true","yes","y","on"))
    meta_opt = None
    if use_meta:
        try:
            from agi_dw.core.metaopt.metaopt_grouped import MixtureMetaOptGrouped  # type: ignore
            from agi_dw.core.metaopt.gate_grouped import GateNetGrouped  # type: ignore
            trainable = [p for p in model.parameters() if p.requires_grad]
            device = next(model.parameters()).device.type
            # Cache and reuse meta optimizer to keep state
            meta_opt = getattr(model, "_sft_meta", None)
            if meta_opt is None:
                gate = GateNetGrouped(n_groups=1, hidden=128)
                meta_opt = MixtureMetaOptGrouped(model, param_groups=[{"params": trainable}], gate=gate, base_lr=lr, base_wd=0.01, device=device)
                setattr(model, "_sft_meta", meta_opt)
        except Exception:
            use_meta = False
            meta_opt = None
    # Fallback/default: AdamW (reuse optimizer across calls to preserve momentum)
    opt = None
    if not use_meta:
        opt = getattr(model, "_sft_opt", None)
        if opt is None:
            opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)
            setattr(model, "_sft_opt", opt)
    # Gradient accumulation
    try:
        grad_accum = max(1, int(os.environ.get("SELFPLAY_GRAD_ACCUM", "1") or 1))
    except Exception:
        grad_accum = 1
    # Optional KL clamp vs frozen baseline
    ref_model = None
    try:
        if bool(os.environ.get("SELFPLAY_KL_CLAMP", "0") == "1"):
            ref_model = copy.deepcopy(model).eval()
            for p in ref_model.parameters():
                p.requires_grad_(False)
    except Exception:
        ref_model = None
    model.train()
    last_loss: float | None = None
    device = next(model.parameters()).device
    _dtype = next(model.parameters()).dtype
    # Pre-tokenize concatenated text once to avoid double tokenization
    text = prompt + "\n" + target
    enc = tok(text, return_tensors="pt")
    enc_prompt = tok(prompt + "\n", return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    labels = enc["input_ids"].clone()
    prompt_len = int(enc_prompt["input_ids"].shape[1])
    labels[:, :prompt_len] = -100
    # Mask EOS tokens from contributing to loss
    try:
        eos_id = getattr(tok, "eos_token_id", None)
        if eos_id is not None:
            labels[labels == int(eos_id)] = -100
    except Exception:
        pass
    # If no valid supervised tokens after shift, skip step to avoid constant loss
    try:
        valid_positions = int(((labels[:, 1:] != -100).sum()).item())
        if valid_positions == 0:
            return None
    except Exception:
        pass
    # Training loop with accumulation
    grad_norm_val: float | None = None
    kl_val: float | None = None
    for it in range(steps):
        if it % grad_accum == 0:
            if opt is not None:
                opt.zero_grad()
        if device.type == "cuda" and _dtype in (getattr(torch, "float16"), getattr(torch, "bfloat16")):
            with _autocast_cuda(_dtype):
                with torch.set_grad_enabled(True):
                    out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                    loss = out.loss
                    # Fallback: if loss is detached, recompute CE from logits
                    try:
                        if not getattr(loss, "requires_grad", True):
                            import torch.nn.functional as F  # type: ignore
                            out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                            logits = out2.logits  # [B, T, V]
                            shift_logits = logits[:, :-1, :].contiguous()
                            shift_labels = labels[:, 1:].contiguous()
                            mask = shift_labels != -100
                            if mask.any():
                                loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                            else:
                                loss = logits.new_tensor(0.0)
                    except Exception:
                        pass
                    # Optional KL clamp
                    if ref_model is not None:
                        with torch.inference_mode():
                            ref = ref_model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                        p = torch.log_softmax(out.logits.float(), dim=-1)
                        q = torch.log_softmax(ref.logits.float(), dim=-1)
                        kl = torch.sum(torch.exp(p) * (p - q), dim=-1).mean()
                        try:
                            kl_val = float(kl.detach().item())
                        except Exception:
                            kl_val = None
                        try:
                            kl_max = float(os.environ.get("SELFPLAY_KL_MAX", "1.0") or 1.0)
                        except Exception:
                            kl_max = 1.0
                        loss = loss + torch.clamp(kl, max=kl_max)
        else:
            with torch.set_grad_enabled(True):
                out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                loss = out.loss
                try:
                    if not getattr(loss, "requires_grad", True):
                        import torch.nn.functional as F  # type: ignore
                        out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                        logits = out2.logits
                        shift_logits = logits[:, :-1, :].contiguous()
                        shift_labels = labels[:, 1:].contiguous()
                        mask = shift_labels != -100
                        if mask.any():
                            loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                        else:
                            loss = logits.new_tensor(0.0)
                except Exception:
                    pass
        # If loss is still detached (no trainable path), skip this micro-step safely
        try:
            if not getattr(loss, "requires_grad", True):
                continue
        except Exception:
            pass
        loss.backward()
        # Step on accumulation interval
        if (it % grad_accum) == (grad_accum - 1):
            if use_meta and meta_opt is not None:
                try:
                    meta_opt.step(float(loss.detach().item()))
                except Exception:
                    pass
                # zero grads for next accumulation window
                for p in model.parameters():
                    if p.grad is not None:
                        p.grad.detach_(); p.grad.zero_()
            else:
                try:
                    clip = float(os.environ.get("SELFPLAY_CLIP_NORM", "1.0") or 1.0)
                except Exception:
                    clip = 1.0
                try:
                    _gn = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip)  # type: ignore
                    try:
                        grad_norm_val = float(getattr(_gn, "item", lambda: _gn)()) if _gn is not None else None  # type: ignore
                    except Exception:
                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    # Log training metrics once per call
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "sft", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val, "kl_to_ref": kl_val}})
    except Exception:
        pass
    return last_loss

def contrastive_step(tok, model, prompt: str, neg: str, lr: float, steps: int, margin: float = 0.1, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:
            alpha = 16


                        grad_norm_val = float(getattr(_gn, "item", lambda: _gn)()) if _gn is not None else None  # type: ignore
                    except Exception:
                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    # Log training metrics once per call
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "sft", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val, "kl_to_ref": kl_val}})
    except Exception:
        pass
    return last_loss

def contrastive_step(tok, model, prompt: str, neg: str, lr: float, steps: int, margin: float = 0.1, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:
            alpha = 16
        try:
            _inject_lora(model, r=r, alpha=alpha)
            _patch_forwards(model)
        except Exception:
            pass
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
        if enabled == 0:
            try:
                r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
            except Exception:
                r = 8
            try:
                alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
            except Exception:
                alpha = 16
            _inject_lora(model, r=r, alpha=alpha)
            _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)
            _set_lora_enabled(model, True)
    except Exception:
        pass
    # Optional MetaOpt integration: default enabled unless explicitly disabled
    _v_meta = os.environ.get("AGI_METAOPT")
    use_meta = True if _v_meta is None else (str(_v_meta).strip().lower() in ("1","true","yes","y","on"))
    meta_opt = None
    if use_meta:
        try:
            from agi_dw.core.metaopt.metaopt_grouped import MixtureMetaOptGrouped  # type: ignore
            from agi_dw.core.metaopt.gate_grouped import GateNetGrouped  # type: ignore
            trainable = [p for p in model.parameters() if p.requires_grad]
            device = next(model.parameters()).device.type
            meta_opt = getattr(model, "_ctr_meta", None)
            if meta_opt is None:
                gate = GateNetGrouped(n_groups=1, hidden=128)
                meta_opt = MixtureMetaOptGrouped(model, param_groups=[{"params": trainable}], gate=gate, base_lr=lr, base_wd=0.01, device=device)
                setattr(model, "_ctr_meta", meta_opt)
        except Exception:
            use_meta = False
            meta_opt = None
    opt = None
    if not use_meta:
        opt = getattr(model, "_ctr_opt", None)
        if opt is None:
            opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr, weight_decay=0.01)
            setattr(model, "_ctr_opt", opt)
    try:
        grad_accum = max(1, int(os.environ.get("SELFPLAY_GRAD_ACCUM", "1") or 1))
    except Exception:
        grad_accum = 1
    model.train()
    last_loss: float | None = None
    device = next(model.parameters()).device
    _dtype = next(model.parameters()).dtype
    # Pre-tokenize once
    text = prompt + "\n" + neg
    enc = tok(text, return_tensors="pt")
    enc_prompt = tok(prompt + "\n", return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    labels = enc["input_ids"].clone()
    prompt_len = int(enc_prompt["input_ids"].shape[1])
    labels[:, :prompt_len] = -100
    # If no valid supervised tokens (e.g., empty neg), skip to avoid no-grad loss
    try:
        if int((labels != -100).sum().item()) == 0:
            return None
    except Exception:
        pass
    try:
        eos_id = getattr(tok, "eos_token_id", None)
        if eos_id is not None:
            labels[labels == int(eos_id)] = -100
    except Exception:
        pass
    # If no valid supervised positions after shift, skip to avoid constant, no-grad loss
    try:
        valid_positions = int(((labels[:, 1:] != -100).sum()).item())
        if valid_positions == 0:
            return None
    except Exception:
        pass
    grad_norm_val: float | None = None
    for it in range(steps):
        if it % grad_accum == 0:
            if opt is not None:
                opt.zero_grad()
        if device.type == "cuda" and _dtype in (getattr(torch, "float16"), getattr(torch, "bfloat16")):
            with _autocast_cuda(_dtype):
                with torch.set_grad_enabled(True):
                    out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                    # Encourage model to keep negative likelihood below margin by making loss zero once margin reached
                    base_loss = out.loss
                    try:
                        if not getattr(base_loss, "requires_grad", True):
                            import torch.nn.functional as F  # type: ignore
                            out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                            logits = out2.logits
                            shift_logits = logits[:, :-1, :].contiguous()
                            shift_labels = labels[:, 1:].contiguous()
                            mask = shift_labels != -100
                            if mask.any():
                                base_loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                            else:
                                base_loss = logits.new_tensor(0.0)
                    except Exception:
                        pass
                    loss = torch.relu(margin - base_loss)
        else:
            with torch.set_grad_enabled(True):
                out = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"), labels=labels)
                base_loss = out.loss
                try:
                    if not getattr(base_loss, "requires_grad", True):
                        import torch.nn.functional as F  # type: ignore
                        out2 = model(input_ids=enc["input_ids"], attention_mask=enc.get("attention_mask"))
                        logits = out2.logits
                        shift_logits = logits[:, :-1, :].contiguous()
                        shift_labels = labels[:, 1:].contiguous()
                        mask = shift_labels != -100
                        if mask.any():
                            base_loss = F.cross_entropy(shift_logits[mask], shift_labels[mask], reduction="mean")
                        else:
                            base_loss = logits.new_tensor(0.0)
                except Exception:
                    pass
                loss = torch.relu(margin - base_loss)
        try:
            if not getattr(loss, "requires_grad", True):
                continue
        except Exception:
            pass
        loss.backward()
        if (it % grad_accum) == (grad_accum - 1):
            if use_meta and meta_opt is not None:
                try:
                    meta_opt.step(float(loss.detach().item()))
                except Exception:
                    pass
                for p in model.parameters():
                    if p.grad is not None:
                        p.grad.detach_(); p.grad.zero_()
            else:
                try:
                    clip = float(os.environ.get("SELFPLAY_CLIP_NORM", "1.0") or 1.0)
                except Exception:
                    clip = 1.0
                try:
                    _gn = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip)  # type: ignore
                    try:
                        grad_norm_val = float(getattr(_gn, "item", lambda: _gn)()) if _gn is not None else None  # type: ignore
                    except Exception:
                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "contrastive", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val}})
    except Exception:
        pass
    return last_loss

def _autocast_cuda(dtype: torch.dtype):
    """Return a CUDA autocast context, compatible with both torch.amp.autocast and torch.autocast."""
    try:
        return torch.amp.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]
    except AttributeError:
        return torch.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]

# ========================= New training utilities =========================

def _enable_lora_train_only(model: torch.nn.Module) -> None:
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1


                    try:
                        grad_norm_val = float(getattr(_gn, "item", lambda: _gn)()) if _gn is not None else None  # type: ignore
                    except Exception:
                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "contrastive", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val}})
    except Exception:
        pass
    return last_loss

def _autocast_cuda(dtype: torch.dtype):
    """Return a CUDA autocast context, compatible with both torch.amp.autocast and torch.autocast."""
    try:
        return torch.amp.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]
    except AttributeError:
        return torch.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]

# ========================= New training utilities =========================

def _enable_lora_train_only(model: torch.nn.Module) -> None:
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
    except Exception:
        pass

def _token_logprobs(tok, model, prompt: str, completion: str) -> Tuple[torch.Tensor, torch.Tensor]:


        except Exception:
            pass
    model.eval()
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "contrastive", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val}})
    except Exception:
        pass
    return last_loss

def _autocast_cuda(dtype: torch.dtype):
    """Return a CUDA autocast context, compatible with both torch.amp.autocast and torch.autocast."""
    try:
        return torch.amp.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]
    except AttributeError:
        return torch.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]

# ========================= New training utilities =========================

def _enable_lora_train_only(model: torch.nn.Module) -> None:
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
    except Exception:
        pass

def _token_logprobs(tok, model, prompt: str, completion: str) -> Tuple[torch.Tensor, torch.Tensor]:
    device = next(model.parameters()).device
    full = (prompt or "") + (completion or "")
    enc = tok(full, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        out = model(**enc)
        logits = out.logits  # (1, T, V)
    ids = enc["input_ids"][0]
    prompt_len = len(tok(prompt, return_tensors="pt")["input_ids"][0])
    target_ids = ids[prompt_len:]
    logits_shifted = logits[:, prompt_len - 1 : -1, :]
    logprobs = logits_shifted.log_softmax(-1)[0]
    tgt = target_ids
    lp = logprobs.gather(-1, tgt.view(-1, 1)).squeeze(-1)
    return lp, tgt

def dpo_step(tok, model, prompt: str, good: str, bad: str, lr: float, beta: float = 0.1) -> float | None:
    """One DPO-style step on a single (prompt, good, bad) pair."""


        return torch.autocast("cuda", dtype=dtype)  # type: ignore[attr-defined]

# ========================= New training utilities =========================

def _enable_lora_train_only(model: torch.nn.Module) -> None:
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:
                p.requires_grad_(True)
                enabled += 1
            else:
                p.requires_grad_(False)
        _set_lora_enabled(model, True)
    except Exception:
        pass

def _token_logprobs(tok, model, prompt: str, completion: str) -> Tuple[torch.Tensor, torch.Tensor]:
    device = next(model.parameters()).device
    full = (prompt or "") + (completion or "")
    enc = tok(full, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        out = model(**enc)
        logits = out.logits  # (1, T, V)
    ids = enc["input_ids"][0]
    prompt_len = len(tok(prompt, return_tensors="pt")["input_ids"][0])
    target_ids = ids[prompt_len:]
    logits_shifted = logits[:, prompt_len - 1 : -1, :]
    logprobs = logits_shifted.log_softmax(-1)[0]
    tgt = target_ids
    lp = logprobs.gather(-1, tgt.view(-1, 1)).squeeze(-1)
    return lp, tgt

def dpo_step(tok, model, prompt: str, good: str, bad: str, lr: float, beta: float = 0.1) -> float | None:
    """One DPO-style step on a single (prompt, good, bad) pair."""
    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp_g, _ = _token_logprobs(tok, model, prompt, good)
    lp_b, _ = _token_logprobs(tok, model, prompt, bad)
    logp_g = lp_g.sum()
    logp_b = lp_b.sum()
    pref_margin = (logp_g - logp_b)
    loss = torch.nn.functional.softplus(-float(beta) * pref_margin)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return float(loss.detach().item())


    except Exception:
        pass

def _token_logprobs(tok, model, prompt: str, completion: str) -> Tuple[torch.Tensor, torch.Tensor]:
    device = next(model.parameters()).device
    full = (prompt or "") + (completion or "")
    enc = tok(full, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        out = model(**enc)
        logits = out.logits  # (1, T, V)
    ids = enc["input_ids"][0]
    prompt_len = len(tok(prompt, return_tensors="pt")["input_ids"][0])
    target_ids = ids[prompt_len:]
    logits_shifted = logits[:, prompt_len - 1 : -1, :]
    logprobs = logits_shifted.log_softmax(-1)[0]
    tgt = target_ids
    lp = logprobs.gather(-1, tgt.view(-1, 1)).squeeze(-1)
    return lp, tgt

def dpo_step(tok, model, prompt: str, good: str, bad: str, lr: float, beta: float = 0.1) -> float | None:
    """One DPO-style step on a single (prompt, good, bad) pair."""
    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp_g, _ = _token_logprobs(tok, model, prompt, good)
    lp_b, _ = _token_logprobs(tok, model, prompt, bad)
    logp_g = lp_g.sum()
    logp_b = lp_b.sum()
    pref_margin = (logp_g - logp_b)
    loss = torch.nn.functional.softplus(-float(beta) * pref_margin)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return float(loss.detach().item())

def kto_step(tok, model, prompt: str, good: str, bad: str, lr: float, tau: float = 1.0) -> float | None:
    """KTO-style pairwise margin loss: max(0, tau - (logp_good - logp_bad))."""
    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp_g, _ = _token_logprobs(tok, model, prompt, good)
    lp_b, _ = _token_logprobs(tok, model, prompt, bad)
    logp_g = lp_g.sum()
    logp_b = lp_b.sum()
    margin = float(tau)
    loss = torch.clamp(margin - (logp_g - logp_b), min=0.0)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return float(loss.detach().item())


def dpo_step(tok, model, prompt: str, good: str, bad: str, lr: float, beta: float = 0.1) -> float | None:
    """One DPO-style step on a single (prompt, good, bad) pair."""
    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp_g, _ = _token_logprobs(tok, model, prompt, good)
    lp_b, _ = _token_logprobs(tok, model, prompt, bad)
    logp_g = lp_g.sum()
    logp_b = lp_b.sum()
    pref_margin = (logp_g - logp_b)
    loss = torch.nn.functional.softplus(-float(beta) * pref_margin)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return float(loss.detach().item())

def kto_step(tok, model, prompt: str, good: str, bad: str, lr: float, tau: float = 1.0) -> float | None:
    """KTO-style pairwise margin loss: max(0, tau - (logp_good - logp_bad))."""
    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp_g, _ = _token_logprobs(tok, model, prompt, good)
    lp_b, _ = _token_logprobs(tok, model, prompt, bad)
    logp_g = lp_g.sum()
    logp_b = lp_b.sum()
    margin = float(tau)
    loss = torch.clamp(margin - (logp_g - logp_b), min=0.0)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return float(loss.detach().item())

# ---------------- UL2 span corruption ----------------

def ul2_mask_spans(ids: torch.LongTensor, mask_frac: float = 0.15, mean_span: int = 3, sentinel_id: int = 32000, pad_id: Optional[int] = None) -> Tuple[torch.LongTensor, torch.LongTensor]:
    """Create UL2-style span masks and replace masked spans with sentinel tokens.

    Returns (corrupted_ids, target_ids) where target_ids contains original tokens for masked positions and -100 elsewhere.
    """
    T = int(ids.shape[0])
    num_to_mask = max(1, int(T * float(mask_frac)))
    rng = torch.randint(0, T, (num_to_mask,))
    mask = torch.zeros(T, dtype=torch.bool, device=ids.device)
    for i in rng.tolist():
        span = max(1, int(torch.poisson(torch.tensor(float(mean_span))).item()))
        j = min(T, i + span)
        mask[i:j] = True
    target = torch.full_like(ids, -100)
    target[mask] = ids[mask]
    corrupted = ids.clone()
    # Collapse contiguous masked runs to a single sentinel token; fill rest with pad if provided


    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp_g, _ = _token_logprobs(tok, model, prompt, good)
    lp_b, _ = _token_logprobs(tok, model, prompt, bad)
    logp_g = lp_g.sum()
    logp_b = lp_b.sum()
    margin = float(tau)
    loss = torch.clamp(margin - (logp_g - logp_b), min=0.0)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    return float(loss.detach().item())

# ---------------- UL2 span corruption ----------------

def ul2_mask_spans(ids: torch.LongTensor, mask_frac: float = 0.15, mean_span: int = 3, sentinel_id: int = 32000, pad_id: Optional[int] = None) -> Tuple[torch.LongTensor, torch.LongTensor]:
    """Create UL2-style span masks and replace masked spans with sentinel tokens.

    Returns (corrupted_ids, target_ids) where target_ids contains original tokens for masked positions and -100 elsewhere.
    """
    T = int(ids.shape[0])
    num_to_mask = max(1, int(T * float(mask_frac)))
    rng = torch.randint(0, T, (num_to_mask,))
    mask = torch.zeros(T, dtype=torch.bool, device=ids.device)
    for i in rng.tolist():
        span = max(1, int(torch.poisson(torch.tensor(float(mean_span))).item()))
        j = min(T, i + span)
        mask[i:j] = True
    target = torch.full_like(ids, -100)
    target[mask] = ids[mask]
    corrupted = ids.clone()
    # Collapse contiguous masked runs to a single sentinel token; fill rest with pad if provided
    runs: List[Tuple[int, int]] = []
    in_run = False
    a = 0
    for t in range(T):
        if mask[t] and not in_run:
            in_run = True
            a = t
        elif (not mask[t]) and in_run:
            runs.append((a, t))
            in_run = False
    if in_run:
        runs.append((a, T))
    for a, b in runs:
        corrupted[a] = int(sentinel_id)
        if b - a > 1 and pad_id is not None:
            corrupted[a + 1 : b] = int(pad_id)
    return corrupted, target

def ul2_span_corruption_loss(tok, model, prompt: str, *, mask_frac: float = 0.15, mean_span: int = 3, sentinel_id: int = 32000, lr: Optional[float] = None) -> float:
    device = next(model.parameters()).device
    enc = tok(prompt, return_tensors="pt")
    input_ids = enc["input_ids"][0].to(device)
    pad_id = getattr(tok, "pad_token_id", None)
    corrupted, target = ul2_mask_spans(input_ids, mask_frac=mask_frac, mean_span=mean_span, sentinel_id=sentinel_id, pad_id=pad_id)
    attn = torch.ones_like(corrupted).unsqueeze(0)
    out = model(input_ids=corrupted.unsqueeze(0), attention_mask=attn)
    logits = out.logits[0]
    loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1), ignore_index=-100)
    if lr is not None:
        _enable_lora_train_only(model)
        opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    return float(loss.detach().item())


    corrupted = ids.clone()
    # Collapse contiguous masked runs to a single sentinel token; fill rest with pad if provided
    runs: List[Tuple[int, int]] = []
    in_run = False
    a = 0
    for t in range(T):
        if mask[t] and not in_run:
            in_run = True
            a = t
        elif (not mask[t]) and in_run:
            runs.append((a, t))
            in_run = False
    if in_run:
        runs.append((a, T))
    for a, b in runs:
        corrupted[a] = int(sentinel_id)
        if b - a > 1 and pad_id is not None:
            corrupted[a + 1 : b] = int(pad_id)
    return corrupted, target

def ul2_span_corruption_loss(tok, model, prompt: str, *, mask_frac: float = 0.15, mean_span: int = 3, sentinel_id: int = 32000, lr: Optional[float] = None) -> float:
    device = next(model.parameters()).device
    enc = tok(prompt, return_tensors="pt")
    input_ids = enc["input_ids"][0].to(device)
    pad_id = getattr(tok, "pad_token_id", None)
    corrupted, target = ul2_mask_spans(input_ids, mask_frac=mask_frac, mean_span=mean_span, sentinel_id=sentinel_id, pad_id=pad_id)
    attn = torch.ones_like(corrupted).unsqueeze(0)
    out = model(input_ids=corrupted.unsqueeze(0), attention_mask=attn)
    logits = out.logits[0]
    loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1), ignore_index=-100)
    if lr is not None:
        _enable_lora_train_only(model)
        opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    return float(loss.detach().item())

# ---------------- Retriever head + triplet training ----------------

class SimpleRetrieverHead(torch.nn.Module):
    def __init__(self, hidden_size: int, proj_dim: int = 768):
        super().__init__()
        self.proj = torch.nn.Linear(hidden_size, proj_dim, bias=False)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        summed = (hidden_states * mask).sum(dim=1)
        denom = mask.sum(dim=1).clamp(min=1.0)
        pooled = summed / denom
        return torch.nn.functional.normalize(self.proj(pooled), p=2, dim=-1)

def retriever_triplet_step(tok, model, retr_head: SimpleRetrieverHead, q: str, pos: str, neg: str, *, margin: float = 0.2, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    retr_head.train().to(device)
    opt = torch.optim.AdamW(retr_head.parameters(), lr=float(lr))


SimpleRetrieverHead

    device = next(model.parameters()).device
    enc = tok(prompt, return_tensors="pt")
    input_ids = enc["input_ids"][0].to(device)
    pad_id = getattr(tok, "pad_token_id", None)
    corrupted, target = ul2_mask_spans(input_ids, mask_frac=mask_frac, mean_span=mean_span, sentinel_id=sentinel_id, pad_id=pad_id)
    attn = torch.ones_like(corrupted).unsqueeze(0)
    out = model(input_ids=corrupted.unsqueeze(0), attention_mask=attn)
    logits = out.logits[0]
    loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1), ignore_index=-100)
    if lr is not None:
        _enable_lora_train_only(model)
        opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    return float(loss.detach().item())

# ---------------- Retriever head + triplet training ----------------

class SimpleRetrieverHead(torch.nn.Module):
    def __init__(self, hidden_size: int, proj_dim: int = 768):
        super().__init__()
        self.proj = torch.nn.Linear(hidden_size, proj_dim, bias=False)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        summed = (hidden_states * mask).sum(dim=1)
        denom = mask.sum(dim=1).clamp(min=1.0)
        pooled = summed / denom
        return torch.nn.functional.normalize(self.proj(pooled), p=2, dim=-1)

def retriever_triplet_step(tok, model, retr_head: SimpleRetrieverHead, q: str, pos: str, neg: str, *, margin: float = 0.2, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    retr_head.train().to(device)
    opt = torch.optim.AdamW(retr_head.parameters(), lr=float(lr))
    def _encode(txt: str) -> Tuple[torch.Tensor, torch.Tensor]:
        enc = tok(txt, return_tensors="pt", truncation=True, max_length=512)
        return enc["input_ids"].to(device), enc.get("attention_mask", torch.ones_like(enc["input_ids"])) .to(device)
    qi, qm = _encode(q)
    pi, pm = _encode(pos)
    ni, nm = _encode(neg)
    with torch.no_grad():
        qhs = model(input_ids=qi, attention_mask=qm, output_hidden_states=True).hidden_states[-1]
        phs = model(input_ids=pi, attention_mask=pm, output_hidden_states=True).hidden_states[-1]
        nhs = model(input_ids=ni, attention_mask=nm, output_hidden_states=True).hidden_states[-1]
    qz = retr_head(qhs, qm)
    pz = retr_head(phs, pm)
    nz = retr_head(nhs, nm)
    d_pos = torch.nn.functional.pairwise_distance(qz, pz)


        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    return float(loss.detach().item())

# ---------------- Retriever head + triplet training ----------------

class SimpleRetrieverHead(torch.nn.Module):
    def __init__(self, hidden_size: int, proj_dim: int = 768):
        super().__init__()
        self.proj = torch.nn.Linear(hidden_size, proj_dim, bias=False)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        summed = (hidden_states * mask).sum(dim=1)
        denom = mask.sum(dim=1).clamp(min=1.0)
        pooled = summed / denom
        return torch.nn.functional.normalize(self.proj(pooled), p=2, dim=-1)

def retriever_triplet_step(tok, model, retr_head: SimpleRetrieverHead, q: str, pos: str, neg: str, *, margin: float = 0.2, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    retr_head.train().to(device)
    opt = torch.optim.AdamW(retr_head.parameters(), lr=float(lr))
    def _encode(txt: str) -> Tuple[torch.Tensor, torch.Tensor]:
        enc = tok(txt, return_tensors="pt", truncation=True, max_length=512)
        return enc["input_ids"].to(device), enc.get("attention_mask", torch.ones_like(enc["input_ids"])) .to(device)
    qi, qm = _encode(q)
    pi, pm = _encode(pos)
    ni, nm = _encode(neg)
    with torch.no_grad():
        qhs = model(input_ids=qi, attention_mask=qm, output_hidden_states=True).hidden_states[-1]
        phs = model(input_ids=pi, attention_mask=pm, output_hidden_states=True).hidden_states[-1]
        nhs = model(input_ids=ni, attention_mask=nm, output_hidden_states=True).hidden_states[-1]
    qz = retr_head(qhs, qm)
    pz = retr_head(phs, pm)
    nz = retr_head(nhs, nm)
    d_pos = torch.nn.functional.pairwise_distance(qz, pz)
    d_neg = torch.nn.functional.pairwise_distance(qz, nz)
    loss = torch.clamp(d_pos - d_neg + float(margin), min=0.0).mean()
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- Tool policy imitation ----------------

class ToolPolicyHead(torch.nn.Module):
    def __init__(self, hidden_size: int, num_tools: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, num_tools, bias=True)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return self.classifier(pooled)

def tool_imitation_step(tok, model, tool_head: ToolPolicyHead, prompt: str, api_index: int, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    tool_head.train().to(device)
    opt = torch.optim.AdamW(tool_head.parameters(), lr=float(lr))
    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512)
    ids = enc["input_ids"].to(device)


ToolPolicyHead

    qi, qm = _encode(q)
    pi, pm = _encode(pos)
    ni, nm = _encode(neg)
    with torch.no_grad():
        qhs = model(input_ids=qi, attention_mask=qm, output_hidden_states=True).hidden_states[-1]
        phs = model(input_ids=pi, attention_mask=pm, output_hidden_states=True).hidden_states[-1]
        nhs = model(input_ids=ni, attention_mask=nm, output_hidden_states=True).hidden_states[-1]
    qz = retr_head(qhs, qm)
    pz = retr_head(phs, pm)
    nz = retr_head(nhs, nm)
    d_pos = torch.nn.functional.pairwise_distance(qz, pz)
    d_neg = torch.nn.functional.pairwise_distance(qz, nz)
    loss = torch.clamp(d_pos - d_neg + float(margin), min=0.0).mean()
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- Tool policy imitation ----------------

class ToolPolicyHead(torch.nn.Module):
    def __init__(self, hidden_size: int, num_tools: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, num_tools, bias=True)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return self.classifier(pooled)

def tool_imitation_step(tok, model, tool_head: ToolPolicyHead, prompt: str, api_index: int, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    tool_head.train().to(device)
    opt = torch.optim.AdamW(tool_head.parameters(), lr=float(lr))
    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512)
    ids = enc["input_ids"].to(device)
    mask = enc.get("attention_mask", torch.ones_like(ids)).to(device)
    with torch.no_grad():
        hs = model(input_ids=ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
    logits = tool_head(hs, mask)
    loss = torch.nn.functional.cross_entropy(logits, torch.tensor([int(api_index)], device=device))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- PPO-lite (execution-guided RL) ----------------


    d_pos = torch.nn.functional.pairwise_distance(qz, pz)
    d_neg = torch.nn.functional.pairwise_distance(qz, nz)
    loss = torch.clamp(d_pos - d_neg + float(margin), min=0.0).mean()
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- Tool policy imitation ----------------

class ToolPolicyHead(torch.nn.Module):
    def __init__(self, hidden_size: int, num_tools: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, num_tools, bias=True)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return self.classifier(pooled)

def tool_imitation_step(tok, model, tool_head: ToolPolicyHead, prompt: str, api_index: int, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    tool_head.train().to(device)
    opt = torch.optim.AdamW(tool_head.parameters(), lr=float(lr))
    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512)
    ids = enc["input_ids"].to(device)
    mask = enc.get("attention_mask", torch.ones_like(ids)).to(device)
    with torch.no_grad():
        hs = model(input_ids=ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
    logits = tool_head(hs, mask)
    loss = torch.nn.functional.cross_entropy(logits, torch.tensor([int(api_index)], device=device))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- PPO-lite (execution-guided RL) ----------------

def _sequence_kl(tok, model, ref_model, prompt: str, completion: str) -> torch.Tensor:
    """Compute token-level KL(new || ref) over the completion segment.

    Returns a scalar tensor (mean over positions) for stability.
    """
    if ref_model is None:
        return torch.tensor(0.0, device=next(model.parameters()).device)
    device = next(model.parameters()).device
    full = (prompt or "") + (completion or "")
    enc = tok(full, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        out_ref = ref_model(**enc)
        q = torch.log_softmax(out_ref.logits.float(), dim=-1)
    out_new = model(**enc)
    p = torch.log_softmax(out_new.logits.float(), dim=-1)
    # Only include completion positions in the KL


def tool_imitation_step(tok, model, tool_head: ToolPolicyHead, prompt: str, api_index: int, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    tool_head.train().to(device)
    opt = torch.optim.AdamW(tool_head.parameters(), lr=float(lr))
    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512)
    ids = enc["input_ids"].to(device)
    mask = enc.get("attention_mask", torch.ones_like(ids)).to(device)
    with torch.no_grad():
        hs = model(input_ids=ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
    logits = tool_head(hs, mask)
    loss = torch.nn.functional.cross_entropy(logits, torch.tensor([int(api_index)], device=device))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- PPO-lite (execution-guided RL) ----------------

def _sequence_kl(tok, model, ref_model, prompt: str, completion: str) -> torch.Tensor:
    """Compute token-level KL(new || ref) over the completion segment.

    Returns a scalar tensor (mean over positions) for stability.
    """
    if ref_model is None:
        return torch.tensor(0.0, device=next(model.parameters()).device)
    device = next(model.parameters()).device
    full = (prompt or "") + (completion or "")
    enc = tok(full, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        out_ref = ref_model(**enc)
        q = torch.log_softmax(out_ref.logits.float(), dim=-1)
    out_new = model(**enc)
    p = torch.log_softmax(out_new.logits.float(), dim=-1)
    # Only include completion positions in the KL
    prompt_len = len(tok(prompt, return_tensors="pt")["input_ids"][0])
    p_seg = p[:, prompt_len - 1 : -1, :]
    q_seg = q[:, prompt_len - 1 : -1, :]
    kl = torch.sum(torch.exp(p_seg) * (p_seg - q_seg), dim=-1).mean()
    return kl

def ppo_lite_step(
    tok,
    model,
    prompt: str,
    completion: str,
    *,
    advantage: float,
    lr: float = 1e-4,
    kl_coef: float = 0.01,
    ref_model: Optional[torch.nn.Module] = None,
) -> float:
    """A minimal PPO-like single-example update.

    Optimizes: L = - advantage * sum_t log p_theta(y_t | x, y_<t) + kl_coef * KL(π_theta || π_ref)
    Gradient only flows through LoRA/adapter parameters.
    """
    try:
        _enable_lora_train_only(model)
    except Exception:


    Returns a scalar tensor (mean over positions) for stability.
    """
    if ref_model is None:
        return torch.tensor(0.0, device=next(model.parameters()).device)
    device = next(model.parameters()).device
    full = (prompt or "") + (completion or "")
    enc = tok(full, return_tensors="pt")
    enc = {k: v.to(device) for k, v in enc.items()}
    with torch.no_grad():
        out_ref = ref_model(**enc)
        q = torch.log_softmax(out_ref.logits.float(), dim=-1)
    out_new = model(**enc)
    p = torch.log_softmax(out_new.logits.float(), dim=-1)
    # Only include completion positions in the KL
    prompt_len = len(tok(prompt, return_tensors="pt")["input_ids"][0])
    p_seg = p[:, prompt_len - 1 : -1, :]
    q_seg = q[:, prompt_len - 1 : -1, :]
    kl = torch.sum(torch.exp(p_seg) * (p_seg - q_seg), dim=-1).mean()
    return kl

def ppo_lite_step(
    tok,
    model,
    prompt: str,
    completion: str,
    *,
    advantage: float,
    lr: float = 1e-4,
    kl_coef: float = 0.01,
    ref_model: Optional[torch.nn.Module] = None,
) -> float:
    """A minimal PPO-like single-example update.

    Optimizes: L = - advantage * sum_t log p_theta(y_t | x, y_<t) + kl_coef * KL(π_theta || π_ref)
    Gradient only flows through LoRA/adapter parameters.
    """
    try:
        _enable_lora_train_only(model)
    except Exception:
        pass
    device = next(model.parameters()).device
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=float(lr))
    model.train()
    opt.zero_grad(set_to_none=True)
    lp, _ = _token_logprobs(tok, model, prompt, completion)
    adv = float(advantage)
    pg_loss = - lp.sum() * adv
    # Optional KL anchor to a reference model
    kl_loss = torch.tensor(0.0, device=device)
    if ref_model is not None and float(kl_coef) > 0.0:
        try:
            kl_val = _sequence_kl(tok, model, ref_model, prompt, completion)
            kl_loss = float(kl_coef) * kl_val
        except Exception:
            kl_loss = torch.tensor(0.0, device=device)
    loss = pg_loss + kl_loss
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    try:
        return float(loss.detach().item())
    except Exception:
        return 0.0


                        grad_norm_val = None
                except Exception:
                    grad_norm_val = None
                if opt is not None:
                    opt.step()
        try:
            last_loss = float(loss.detach().item())
        except Exception:
            pass
    model.eval()
    # Log training metrics once per call
    try:
        if mem is not None:
            mem.log({"train_metrics": {"kind": "sft", "loss": (float(last_loss) if last_loss is not None else None), "grad_norm": grad_norm_val, "kl_to_ref": kl_val}})
    except Exception:
        pass
    return last_loss

def contrastive_step(tok, model, prompt: str, neg: str, lr: float, steps: int, margin: float = 0.1, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:
            alpha = 16
        try:
            _inject_lora(model, r=r, alpha=alpha)
            _patch_forwards(model)
        except Exception:
            pass
    try:
        enabled = 0
        for n, p in model.named_parameters():
            is_custom = (".A" in n) or (".B" in n)
            is_peft = ("lora_" in n) or ("loramodule" in n.lower())
            if is_custom or is_peft:


    pi, pm = _encode(pos)
    ni, nm = _encode(neg)
    with torch.no_grad():
        qhs = model(input_ids=qi, attention_mask=qm, output_hidden_states=True).hidden_states[-1]
        phs = model(input_ids=pi, attention_mask=pm, output_hidden_states=True).hidden_states[-1]
        nhs = model(input_ids=ni, attention_mask=nm, output_hidden_states=True).hidden_states[-1]
    qz = retr_head(qhs, qm)
    pz = retr_head(phs, pm)
    nz = retr_head(nhs, nm)
    d_pos = torch.nn.functional.pairwise_distance(qz, pz)
    d_neg = torch.nn.functional.pairwise_distance(qz, nz)
    loss = torch.clamp(d_pos - d_neg + float(margin), min=0.0).mean()
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- Tool policy imitation ----------------

class ToolPolicyHead(torch.nn.Module):
    def __init__(self, hidden_size: int, num_tools: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, num_tools, bias=True)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return self.classifier(pooled)

def tool_imitation_step(tok, model, tool_head: ToolPolicyHead, prompt: str, api_index: int, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    tool_head.train().to(device)
    opt = torch.optim.AdamW(tool_head.parameters(), lr=float(lr))
    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512)
    ids = enc["input_ids"].to(device)
    mask = enc.get("attention_mask", torch.ones_like(ids)).to(device)
    with torch.no_grad():
        hs = model(input_ids=ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
    logits = tool_head(hs, mask)
    loss = torch.nn.functional.cross_entropy(logits, torch.tensor([int(api_index)], device=device))
    opt.zero_grad(set_to_none=True)
    loss.backward()


forward

        phs = model(input_ids=pi, attention_mask=pm, output_hidden_states=True).hidden_states[-1]
        nhs = model(input_ids=ni, attention_mask=nm, output_hidden_states=True).hidden_states[-1]
    qz = retr_head(qhs, qm)
    pz = retr_head(phs, pm)
    nz = retr_head(nhs, nm)
    d_pos = torch.nn.functional.pairwise_distance(qz, pz)
    d_neg = torch.nn.functional.pairwise_distance(qz, nz)
    loss = torch.clamp(d_pos - d_neg + float(margin), min=0.0).mean()
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- Tool policy imitation ----------------

class ToolPolicyHead(torch.nn.Module):
    def __init__(self, hidden_size: int, num_tools: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, num_tools, bias=True)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        return self.classifier(pooled)

def tool_imitation_step(tok, model, tool_head: ToolPolicyHead, prompt: str, api_index: int, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    tool_head.train().to(device)
    opt = torch.optim.AdamW(tool_head.parameters(), lr=float(lr))
    enc = tok(prompt, return_tensors="pt", truncation=True, max_length=512)
    ids = enc["input_ids"].to(device)
    mask = enc.get("attention_mask", torch.ones_like(ids)).to(device)
    with torch.no_grad():
        hs = model(input_ids=ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
    logits = tool_head(hs, mask)
    loss = torch.nn.functional.cross_entropy(logits, torch.tensor([int(api_index)], device=device))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- PPO-lite (execution-guided RL) ----------------


# ---------------- Retriever head + triplet training ----------------

class SimpleRetrieverHead(torch.nn.Module):
    def __init__(self, hidden_size: int, proj_dim: int = 768):
        super().__init__()
        self.proj = torch.nn.Linear(hidden_size, proj_dim, bias=False)

    def forward(self, hidden_states: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
        mask = attn_mask.unsqueeze(-1).float()
        summed = (hidden_states * mask).sum(dim=1)
        denom = mask.sum(dim=1).clamp(min=1.0)
        pooled = summed / denom
        return torch.nn.functional.normalize(self.proj(pooled), p=2, dim=-1)

def retriever_triplet_step(tok, model, retr_head: SimpleRetrieverHead, q: str, pos: str, neg: str, *, margin: float = 0.2, lr: float = 1e-4) -> float:
    device = next(model.parameters()).device
    model.eval()
    retr_head.train().to(device)
    opt = torch.optim.AdamW(retr_head.parameters(), lr=float(lr))
    def _encode(txt: str) -> Tuple[torch.Tensor, torch.Tensor]:
        enc = tok(txt, return_tensors="pt", truncation=True, max_length=512)
        return enc["input_ids"].to(device), enc.get("attention_mask", torch.ones_like(enc["input_ids"])) .to(device)
    qi, qm = _encode(q)
    pi, pm = _encode(pos)
    ni, nm = _encode(neg)
    with torch.no_grad():
        qhs = model(input_ids=qi, attention_mask=qm, output_hidden_states=True).hidden_states[-1]
        phs = model(input_ids=pi, attention_mask=pm, output_hidden_states=True).hidden_states[-1]
        nhs = model(input_ids=ni, attention_mask=nm, output_hidden_states=True).hidden_states[-1]
    qz = retr_head(qhs, qm)
    pz = retr_head(phs, pm)
    nz = retr_head(nhs, nm)
    d_pos = torch.nn.functional.pairwise_distance(qz, pz)
    d_neg = torch.nn.functional.pairwise_distance(qz, nz)
    loss = torch.clamp(d_pos - d_neg + float(margin), min=0.0).mean()
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    return float(loss.detach().item())

# ---------------- Tool policy imitation ----------------


from .common_imports import *

def _resolve_best_wm_model_path() -> str:
    try:
        root = Path(__file__).resolve().parents[1]
        cands = [
            root / "models" / "wm_mlp" / "best.joblib",
            root / "models" / "wm_mlp" / "latest.joblib",
            root / "models" / "wm_mlp" / "wm_mlp.joblib",
        ]
        for p in cands:
            if p.exists():
                return str(p)
    except Exception:
        pass
    return ""


from typing import List, Dict, Any, Tuple, Optional
import sys, io, re, json, math, statistics, itertools, functools, copy, random
import torch
from collections import Counter
from agi_dw.core.llm.hf_client import HFClient  # type: ignore
from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
from transformers.generation.logits_process import (  # type: ignore
    LogitsProcessor,
    LogitsProcessorList,
)
from agi_dw.core.llm.adapter_cache import AdapterCache  # type: ignore
import os
import subprocess
import shlex
from pathlib import Path
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, TimeoutError