import logging
from typing import Any, Dict, List, Optional
import os

_MODEL_CACHE: Dict[str, "HFClient"] = {}

class HFClient:
	def __init__(self, model_id: str = "meta-llama/Llama-3.1-8B-Instruct", device_map: Optional[str] = "auto", torch_dtype: Optional[str] = None) -> None:
		try:
			import torch  # type: ignore
			from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
		except Exception as e:
			raise RuntimeError("transformers/torch not installed: pip install -r requirements.txt") from e
		self.model_id = model_id
		self.tokenizer = AutoTokenizer.from_pretrained(model_id)
		# Allow env overrides and torchrun rank binding
		env_device_map = os.environ.get("HF_DEVICE_MAP") or device_map
		try:
			# If running under torchrun and device_map wasn't explicitly set, bind to this rank's GPU
			if (env_device_map is None or str(env_device_map).strip().lower() == "auto") and os.environ.get("LOCAL_RANK") is not None:
				lr = int(os.environ.get("LOCAL_RANK", "0") or 0)
				env_device_map = f"cuda:{lr}" if lr >= 0 else env_device_map
		except Exception:
			pass
		env_dtype = os.environ.get("HF_TORCH_DTYPE") or torch_dtype
		dtype = None
		if env_dtype == "float16":
			dtype = torch.float16
		elif env_dtype == "bfloat16":
			dtype = torch.bfloat16
		elif env_dtype is None and torch.cuda.is_available():
			# Prefer BF16 if supported, else FP16 for faster GPU inference
			try:
				if hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
					dtype = torch.bfloat16
				else:
					dtype = torch.float16
			except Exception:
				dtype = torch.float16
		try:
			self.model = AutoModelForCausalLM.from_pretrained(
				model_id,
				device_map=env_device_map,
				torch_dtype=dtype,
			)
		except Exception as e:
			# Graceful low-memory fallback: load on CPU if CUDA OOM or device map fails
			try:
				self.model = AutoModelForCausalLM.from_pretrained(
					model_id,
					device_map="cpu",
				)
			except Exception as e2:
				raise e
		# Save dtype for autocast
		self._dtype = dtype
		# Prefer faster attention kernels when available; disable sliding window warnings
		try:
			# Explicitly set attention implementation if configurable
			if hasattr(self.model, "config"):
				cfg = self.model.config
				# Disable sliding window settings that sdpa does not support
				if hasattr(cfg, "sliding_window"):
					try:
						setattr(cfg, "sliding_window", None)
					except Exception:
						pass
				if hasattr(cfg, "use_sliding_window"):
					try:
						setattr(cfg, "use_sliding_window", False)
					except Exception:
						pass
				if hasattr(cfg, "sliding_window_size"):
					try:
						setattr(cfg, "sliding_window_size", 0)
					except Exception:
						pass
				if hasattr(cfg, "use_cache"):
					setattr(cfg, "use_cache", True)
				# Prefer flash-attn2 if available and requested; else SDPA
				attn_impl = "sdpa"
				if os.environ.get("HF_FLASH_ATTN", "0") in ("1", "true", "True"):
					try:
						import flash_attn  # type: ignore  # noqa: F401
						attn_impl = "flash_attention_2"
					except Exception:
						attn_impl = "sdpa"
				if hasattr(cfg, "attn_implementation"):
					setattr(cfg, "attn_implementation", attn_impl)
				if hasattr(self.model, "generation_config"):
					if hasattr(self.model.generation_config, "use_cache"):
						self.model.generation_config.use_cache = True  # type: ignore[attr-defined]
					if hasattr(self.model.generation_config, "sliding_window"):
						try:
							self.model.generation_config.sliding_window = None  # type: ignore[attr-defined]
						except Exception:
							pass
		except Exception:
			pass
		# Ensure model is on GPU if available
		try:
			if torch.cuda.is_available():
				# Respect torchrun LOCAL_RANK when present
				try:
					if os.environ.get("LOCAL_RANK") is not None:
						lr = int(os.environ.get("LOCAL_RANK", "0") or 0)
						if lr >= 0 and hasattr(torch.cuda, "set_device"):
							torch.cuda.set_device(lr)
				except Exception:
					pass
				current = str(self.model.device)
				if "cuda" not in current:
					# Attempt to move to CUDA; if it fails due to OOM, continue on CPU
					try:
						self.model.to("cuda")
					except Exception:
						pass
			# Enable TF32 on Ampere+ for faster matmul where acceptable
			try:
				torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
				torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
				if hasattr(torch, "set_float32_matmul_precision"):
					torch.set_float32_matmul_precision("high")  # type: ignore[attr-defined]
			except Exception:
				pass
		except Exception:
			pass

	@classmethod
	def get_cached(cls, model_id: str) -> "HFClient":
		client = _MODEL_CACHE.get(model_id)
		if client is None:
			client = HFClient(model_id=model_id)
			_MODEL_CACHE[model_id] = client
		return client

	def attach_adapter(self, adapter_dir: Optional[str]) -> None:
		"""Optionally load a PEFT adapter onto the base model if available."""
		if not adapter_dir:
			return
		try:
			from peft import PeftModel  # type: ignore
			self.model = PeftModel.from_pretrained(self.model, adapter_dir)
		except Exception:
			# Silently continue if PEFT not installed or adapter invalid
			pass


					pass
				current = str(self.model.device)
				if "cuda" not in current:
					# Attempt to move to CUDA; if it fails due to OOM, continue on CPU
					try:
						self.model.to("cuda")
					except Exception:
						pass
			# Enable TF32 on Ampere+ for faster matmul where acceptable
			try:
				torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
				torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
				if hasattr(torch, "set_float32_matmul_precision"):
					torch.set_float32_matmul_precision("high")  # type: ignore[attr-defined]
			except Exception:
				pass
		except Exception:
			pass

	@classmethod
	def get_cached(cls, model_id: str) -> "HFClient":
		client = _MODEL_CACHE.get(model_id)
		if client is None:
			client = HFClient(model_id=model_id)
			_MODEL_CACHE[model_id] = client
		return client

	def attach_adapter(self, adapter_dir: Optional[str]) -> None:
		"""Optionally load a PEFT adapter onto the base model if available."""
		if not adapter_dir:
			return
		try:
			from peft import PeftModel  # type: ignore
			self.model = PeftModel.from_pretrained(self.model, adapter_dir)
		except Exception:
			# Silently continue if PEFT not installed or adapter invalid
			pass

	def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		import torch  # type: ignore
		inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
		do_sample = temperature > 0.0
		# Try grammar-constrained decoding via Outlines if requested
		if grammar:
			try:
				import outlines  # type: ignore


						pass
			# Enable TF32 on Ampere+ for faster matmul where acceptable
			try:
				torch.backends.cuda.matmul.allow_tf32 = True  # type: ignore[attr-defined]
				torch.backends.cudnn.allow_tf32 = True  # type: ignore[attr-defined]
				if hasattr(torch, "set_float32_matmul_precision"):
					torch.set_float32_matmul_precision("high")  # type: ignore[attr-defined]
			except Exception:
				pass
		except Exception:
			pass

	@classmethod
	def get_cached(cls, model_id: str) -> "HFClient":
		client = _MODEL_CACHE.get(model_id)
		if client is None:
			client = HFClient(model_id=model_id)
			_MODEL_CACHE[model_id] = client
		return client

	def attach_adapter(self, adapter_dir: Optional[str]) -> None:
		"""Optionally load a PEFT adapter onto the base model if available."""
		if not adapter_dir:
			return
		try:
			from peft import PeftModel  # type: ignore
			self.model = PeftModel.from_pretrained(self.model, adapter_dir)
		except Exception:
			# Silently continue if PEFT not installed or adapter invalid
			pass

	def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		import torch  # type: ignore
		inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
		do_sample = temperature > 0.0
		# Try grammar-constrained decoding via Outlines if requested
		if grammar:
			try:
				import outlines  # type: ignore
				from outlines.models.transformers import Transformers  # type: ignore
				# Map known grammar shorthands to a conservative regex that discourages
				# docstrings and line comments in Python function bodies.
				def _python_body_minimal_regex() -> str:
					# Matches up to max_new_tokens lines consisting of optional indentation
					# and any non-newline characters, but forbids lines starting with
					# triple quotes or '#'. Negative lookaheads apply per line.
					# Note: This is a best-effort regex; Outlines enforces it token by token.
					return r"(?:[ \t]*(?!#)(?!\"\"\"|''')[^\n]*\n?){1," + str(max(1, int(max_new_tokens))) + r"}"
				pattern = None
				if grammar == "python_function_body_minimal":


generate

	@classmethod
	def get_cached(cls, model_id: str) -> "HFClient":
		client = _MODEL_CACHE.get(model_id)
		if client is None:
			client = HFClient(model_id=model_id)
			_MODEL_CACHE[model_id] = client
		return client

	def attach_adapter(self, adapter_dir: Optional[str]) -> None:
		"""Optionally load a PEFT adapter onto the base model if available."""
		if not adapter_dir:
			return
		try:
			from peft import PeftModel  # type: ignore
			self.model = PeftModel.from_pretrained(self.model, adapter_dir)
		except Exception:
			# Silently continue if PEFT not installed or adapter invalid
			pass

	def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		import torch  # type: ignore
		inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
		do_sample = temperature > 0.0
		# Try grammar-constrained decoding via Outlines if requested
		if grammar:
			try:
				import outlines  # type: ignore
				from outlines.models.transformers import Transformers  # type: ignore
				# Map known grammar shorthands to a conservative regex that discourages
				# docstrings and line comments in Python function bodies.
				def _python_body_minimal_regex() -> str:
					# Matches up to max_new_tokens lines consisting of optional indentation
					# and any non-newline characters, but forbids lines starting with
					# triple quotes or '#'. Negative lookaheads apply per line.
					# Note: This is a best-effort regex; Outlines enforces it token by token.
					return r"(?:[ \t]*(?!#)(?!\"\"\"|''')[^\n]*\n?){1," + str(max(1, int(max_new_tokens))) + r"}"
				pattern = None
				if grammar == "python_function_body_minimal":
					pattern = _python_body_minimal_regex()
				else:
					# Allow passing a raw regex as grammar string
					pattern = str(grammar)
				omodel = Transformers(self.model, self.tokenizer)
				gen = outlines.generate.regex(omodel, pattern)
				# Outlines handles sampling internally; pass temperature when sampling
				kwargs: Dict[str, Any] = {}
				if do_sample:
					kwargs["temperature"] = float(temperature)
					if top_p is not None:
						kwargs["top_p"] = float(top_p)
					if top_k is not None:
						kwargs["top_k"] = int(top_k)
				# Execute constrained generation. If anything fails, fall back below.
				text = str(gen(prompt, max_tokens=int(max_new_tokens), **kwargs))
				# Heuristic to drop echoed prompt
				if text.startswith(prompt):
					text = text[len(prompt):]
				text = text.lstrip("\n").rstrip()
				# Optional post-hoc stop handling for safety
				try:
					if stop:
						cut = len(text)
						for s in stop:
							if not s:
								continue
							idx = text.find(s)
							if idx != -1:
								cut = min(cut, idx)
						text = text[:cut]
				except Exception:
					pass
				return text
			except Exception:
				# Fall back to vanilla generation below
				pass
		with torch.inference_mode():
			gen_kwargs: Dict[str, Any] = {
				"max_new_tokens": max_new_tokens,
				"do_sample": do_sample,
				"pad_token_id": self.tokenizer.eos_token_id,
			}
			if do_sample:
				gen_kwargs["temperature"] = temperature
				if top_p is not None:
					gen_kwargs["top_p"] = float(top_p)
				if top_k is not None:
					gen_kwargs["top_k"] = int(top_k)
			# Remove sampler-only params to avoid warnings
			if not do_sample:
				if hasattr(self.model.generation_config, "top_p"):
					try:
						self.model.generation_config.top_p = None  # type: ignore[attr-defined]
					except Exception:
						pass
				if hasattr(self.model.generation_config, "top_k"):
					try:
						self.model.generation_config.top_k = None  # type: ignore[attr-defined]
					except Exception:
						pass
				if hasattr(self.model.generation_config, "temperature"):
					try:
						self.model.generation_config.temperature = None  # type: ignore[attr-defined]
					except Exception:
						pass
			# Autocast to configured dtype
			try:
				if torch.cuda.is_available() and self._dtype in (torch.float16, torch.bfloat16):
					with torch.amp.autocast("cuda", dtype=self._dtype):  # type: ignore[attr-defined]
						output = self.model.generate(**inputs, **gen_kwargs)
				else:
					output = self.model.generate(**inputs, **gen_kwargs)
			except Exception:
				output = self.model.generate(**inputs, **gen_kwargs)
		text = self.tokenizer.decode(output[0], skip_special_tokens=True)
		# Heuristic to return only the completion while preserving leading indentation
		if text.startswith(prompt):
			text = text[len(prompt) :]
		else:
			text = text
		# Optional post-hoc stop truncation by strings
		try:
			if stop:
				cut = len(text)
				for s in stop:
					if not s:
						continue
					idx = text.find(s)
					if idx != -1:
						cut = min(cut, idx)
				text = text[:cut]
		except Exception:
			pass
		return text.lstrip("\n").rstrip()

	def generate_batch(self, prompts: List[str], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[str]:
		import torch  # type: ignore
		do_sample = temperature > 0.0
		enc = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True)
		enc = {k: v.to(self.model.device) for k, v in enc.items()}
		gen_kwargs: Dict[str, Any] = {
			"max_new_tokens": max_new_tokens,
			"do_sample": do_sample,
			"pad_token_id": self.tokenizer.eos_token_id,
		}
		if do_sample:
			gen_kwargs["temperature"] = temperature
			if top_p is not None:
				gen_kwargs["top_p"] = float(top_p)
			if top_k is not None:
				gen_kwargs["top_k"] = int(top_k)
		with torch.inference_mode():
			if torch.cuda.is_available() and self._dtype in (torch.float16, torch.bfloat16):
				with torch.amp.autocast("cuda", dtype=self._dtype):  # type: ignore[attr-defined]


		# Heuristic to return only the completion while preserving leading indentation
		if text.startswith(prompt):
			text = text[len(prompt) :]
		else:
			text = text
		# Optional post-hoc stop truncation by strings
		try:
			if stop:
				cut = len(text)
				for s in stop:
					if not s:
						continue
					idx = text.find(s)
					if idx != -1:
						cut = min(cut, idx)
				text = text[:cut]
		except Exception:
			pass
		return text.lstrip("\n").rstrip()

	def generate_batch(self, prompts: List[str], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None) -> List[str]:
		import torch  # type: ignore
		do_sample = temperature > 0.0
		enc = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True)
		enc = {k: v.to(self.model.device) for k, v in enc.items()}
		gen_kwargs: Dict[str, Any] = {
			"max_new_tokens": max_new_tokens,
			"do_sample": do_sample,
			"pad_token_id": self.tokenizer.eos_token_id,
		}
		if do_sample:
			gen_kwargs["temperature"] = temperature
			if top_p is not None:
				gen_kwargs["top_p"] = float(top_p)
			if top_k is not None:
				gen_kwargs["top_k"] = int(top_k)
		with torch.inference_mode():
			if torch.cuda.is_available() and self._dtype in (torch.float16, torch.bfloat16):
				with torch.amp.autocast("cuda", dtype=self._dtype):  # type: ignore[attr-defined]
					out = self.model.generate(**enc, **gen_kwargs)
			else:
				out = self.model.generate(**enc, **gen_kwargs)
		texts = self.tokenizer.batch_decode(out, skip_special_tokens=True)
		# Trim echoes where possible, preserving leading indentation
		cleaned: List[str] = []
		for p, t in zip(prompts, texts):
			if t.startswith(p):
				cleaned.append(t[len(p):].lstrip("\n").rstrip())
			else:
				cleaned.append(t.lstrip("\n").rstrip())
		return cleaned

	def chat(self, messages: List[Dict[str, str]], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		# Try chat template if available
		prompt = None
		if hasattr(self.tokenizer, "apply_chat_template"):
			try:
				prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
			except Exception:
				prompt = None
		if prompt is None:
			# Fallback to simple concat
			parts: List[str] = []
			for m in messages:
				role = m.get("role", "user")
				content = m.get("content", "")
				parts.append(f"{role}: {content}")
			prompt = "\n".join(parts) + "\nassistant:"
		# If a grammar is requested, route through generate() with grammar to enable
		# constrained decoding for chat-style prompts as well.
		return self.generate(prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, stop=stop, grammar=grammar)


chat

			if top_p is not None:
				gen_kwargs["top_p"] = float(top_p)
			if top_k is not None:
				gen_kwargs["top_k"] = int(top_k)
		with torch.inference_mode():
			if torch.cuda.is_available() and self._dtype in (torch.float16, torch.bfloat16):
				with torch.amp.autocast("cuda", dtype=self._dtype):  # type: ignore[attr-defined]
					out = self.model.generate(**enc, **gen_kwargs)
			else:
				out = self.model.generate(**enc, **gen_kwargs)
		texts = self.tokenizer.batch_decode(out, skip_special_tokens=True)
		# Trim echoes where possible, preserving leading indentation
		cleaned: List[str] = []
		for p, t in zip(prompts, texts):
			if t.startswith(p):
				cleaned.append(t[len(p):].lstrip("\n").rstrip())
			else:
				cleaned.append(t.lstrip("\n").rstrip())
		return cleaned

	def chat(self, messages: List[Dict[str, str]], max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		# Try chat template if available
		prompt = None
		if hasattr(self.tokenizer, "apply_chat_template"):
			try:
				prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
			except Exception:
				prompt = None
		if prompt is None:
			# Fallback to simple concat
			parts: List[str] = []
			for m in messages:
				role = m.get("role", "user")
				content = m.get("content", "")
				parts.append(f"{role}: {content}")
			prompt = "\n".join(parts) + "\nassistant:"
		# If a grammar is requested, route through generate() with grammar to enable
		# constrained decoding for chat-style prompts as well.
		return self.generate(prompt, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, stop=stop, grammar=grammar)


		if not adapter_dir:
			return
		try:
			from peft import PeftModel  # type: ignore
			self.model = PeftModel.from_pretrained(self.model, adapter_dir)
		except Exception:
			# Silently continue if PEFT not installed or adapter invalid
			pass

	def generate(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.0, top_p: Optional[float] = None, top_k: Optional[int] = None, stop: Optional[List[str]] = None, grammar: Optional[str] = None) -> str:
		import torch  # type: ignore
		inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
		do_sample = temperature > 0.0
		# Try grammar-constrained decoding via Outlines if requested
		if grammar:
			try:
				import outlines  # type: ignore
				from outlines.models.transformers import Transformers  # type: ignore
				# Map known grammar shorthands to a conservative regex that discourages
				# docstrings and line comments in Python function bodies.
				def _python_body_minimal_regex() -> str:
					# Matches up to max_new_tokens lines consisting of optional indentation
					# and any non-newline characters, but forbids lines starting with
					# triple quotes or '#'. Negative lookaheads apply per line.
					# Note: This is a best-effort regex; Outlines enforces it token by token.
					return r"(?:[ \t]*(?!#)(?!\"\"\"|''')[^\n]*\n?){1," + str(max(1, int(max_new_tokens))) + r"}"
				pattern = None
				if grammar == "python_function_body_minimal":
					pattern = _python_body_minimal_regex()
				else:
					# Allow passing a raw regex as grammar string
					pattern = str(grammar)
				omodel = Transformers(self.model, self.tokenizer)
				gen = outlines.generate.regex(omodel, pattern)
				# Outlines handles sampling internally; pass temperature when sampling
				kwargs: Dict[str, Any] = {}
				if do_sample:
					kwargs["temperature"] = float(temperature)
					if top_p is not None:
						kwargs["top_p"] = float(top_p)
					if top_k is not None:
						kwargs["top_k"] = int(top_k)
				# Execute constrained generation. If anything fails, fall back below.
				text = str(gen(prompt, max_tokens=int(max_new_tokens), **kwargs))
				# Heuristic to drop echoed prompt
				if text.startswith(prompt):


from __future__ import annotations
import logging
import os

import threading
from pathlib import Path
from typing import Dict, Optional, Tuple, Any

class AdapterCache:
	"""
	Thread-safe cache for loading/retrieving LoRA/PEFT adapters on demand.
	Maintains a mapping (base_model_id, adapter_dir) -> (tokenizer, model).
	Caller owns inference usage; this class does not schedule devices.
	"""

	_lock = threading.Lock()
	_cache: Dict[Tuple[str, str], Tuple[Any, Any]] = {}

	@classmethod
	def get(cls, base_model: str, adapter_dir: str) -> Tuple[Any, Any]:
		key = (str(base_model), str(adapter_dir))
		with cls._lock:
			if key in cls._cache:
				return cls._cache[key]
			# Lazy load tokenizer+model with adapter
			from transformers import AutoTokenizer, AutoModelForCausalLM  # type: ignore
			from peft import PeftModel  # type: ignore
			import torch  # type: ignore
			tok = AutoTokenizer.from_pretrained(base_model)
			if getattr(tok, "pad_token_id", None) is None:
				try:
					tok.pad_token_id = tok.eos_token_id
				except Exception:
					pass
			# Respect torchrun LOCAL_RANK when present to bind each process to its GPU
			device_map: Optional[str] = os.environ.get("HF_DEVICE_MAP") or "auto"
			try:
				if (device_map is None or str(device_map).strip().lower() == "auto") and os.environ.get("LOCAL_RANK") is not None:
					lr = int(os.environ.get("LOCAL_RANK", "0") or 0)
					if lr >= 0:
						device_map = f"cuda:{lr}"
			except Exception:
				pass
			base = AutoModelForCausalLM.from_pretrained(
				base_model,
				device_map=device_map,
				torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
			)
			peft_model = PeftModel.from_pretrained(base, adapter_dir)
			cls._cache[key] = (tok, peft_model)
			return cls._cache[key]


import logging
from typing import Any, Dict, Optional, List, Tuple
from pathlib import Path
from datetime import datetime

from agi_dw.core.llm.hf_client import HFClient
from agi_dw.core.utils.prompt_logger import PromptLogger, get_prompt_logger  # type: ignore
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore
from agi_dw.tools.artifact_cache import is_enabled as cache_enabled, get_cached_text as cache_get, set_cached_text as cache_set  # type: ignore

try:
	import yaml  # type: ignore
except Exception:
	yaml = None

# Reuse actuator parsing/coercion to salvage slightly malformed outputs
try:
	from agi_dw.core.actuator.parse import coerce_flat_yaml  # type: ignore
except Exception:
	def coerce_flat_yaml(_: str) -> Dict[str, Any]:  # fallback no-op
		return {}

import re
import json
import os

def verify_sequence(
	sequence: Dict[str, Any],
	state_requirements: Optional[Dict[int, List[str]]] = None,
	validation_points: Optional[List[Dict[str, Any]]] = None
) -> Dict[str, Any]:
	"""Verify sequence execution against requirements and validation points."""
	results = {
		"state_requirements": {
			"satisfied": [],
			"missing": []
		},
		"validation_points": []
	}

	# Check state requirements
	if state_requirements:
		for step_idx, required in state_requirements.items():
			step = sequence.get("steps", [])[step_idx] if step_idx < len(sequence.get("steps", [])) else None
			if step:
				state = step.get("state", {})
				for req in required:
					if _check_state_requirement(state, req):
						results["state_requirements"]["satisfied"].append(req)
					else:
						results["state_requirements"]["missing"].append(req)

	# Check validation points
	if validation_points:
		for point in validation_points:
			step_idx = point.get("step")
			check = point.get("check")
			step = sequence.get("steps", [])[step_idx] if step_idx < len(sequence.get("steps", [])) else None

			if step and check:
				passed, details = _check_validation_point(step, check)
				results["validation_points"].append({
					"step": step_idx,
					"check": check,
					"passed": passed,
					"details": details
				})

	return results

def _check_state_requirement(state: Dict[str, Any], requirement: str) -> bool:
	"""Check if a state requirement is satisfied."""
	# Handle form state
	if requirement.startswith("form."):
		field = requirement.split(".", 1)[1]
		return bool(state.get("form_state", {}).get(field))

	# Handle navigation state
	if requirement.startswith("nav."):
		page = requirement.split(".", 1)[1]
		return page in state.get("navigation_history", [])

	# Handle element state
	if requirement.startswith("element."):
		selector = requirement.split(".", 1)[1]
		return bool(state.get("dom_state", {}).get(selector))

	return False

def _check_validation_point(step: Dict[str, Any], check: str) -> Tuple[bool, str]:
	"""Check if a validation point passes."""
	state = step.get("state", {})
	observation = step.get("observation", {})

	# Handle form validation
	if check.startswith("form."):
		field = check.split(".", 1)[1]
		if field in state.get("form_state", {}):
			return True, f"Form field {field} is set"
		return False, f"Form field {field} not found"

	# Handle element validation
	if check.startswith("element."):
		selector = check.split(".", 1)[1]
		if selector in state.get("dom_state", {}):
			return True, f"Element {selector} found"
		return False, f"Element {selector} not found"

	# Handle content validation
	if check.startswith("content."):
		text = check.split(".", 1)[1]
		if text.lower() in str(observation.get("content", "")).lower():
			return True, f"Content '{text}' found"
		return False, f"Content '{text}' not found"

	return False, "Unknown validation check"

def verify_trace_snippet(
	trace: Dict[str, Any],
	model: str = "qwen3:8b",
	timeout_sec: int = 8,
	use_llm: bool = True,
	require_llm: bool = False,
	backend: str = "hf",
	log_prompts: bool = False,
	adapter_dir: str | None = None,
	structured_mode: str = "none",
) -> Dict[str, Any]:
	"""
	LLM verifier with optional strict mode. HF backend only.
	Prompts require YAML output; we parse YAML (or JSON fallback) to a dict.
	"""
	if use_llm:
		with trace_span("verify_trace", {"backend": backend, "structured": structured_mode}):
			logger = get_prompt_logger("verifier", bool(log_prompts), echo=bool(log_prompts))
			def _log(kind: str, text: str) -> None:
				logger.log_text(kind, text)
			def _red(s: str) -> str:
				try:
					from agi_dw.core.utils.redact import redact_text  # type: ignore
					rs, _ = redact_text(s)
					return rs
				except Exception:
					return s
			if backend == "hf":
				try:
					# Prefer local PEFT adapter if provided
					if adapter_dir:
						try:
							from agi_dw.core.llm.adapter_cache import AdapterCache  # type: ignore
							tok, peft_model = AdapterCache.get(model, adapter_dir)
							prompt = (
								"You are a strict verifier. Given a trace snippet, output ONLY a YAML mapping with keys"
								" success_prob (float 0..1), risk (float 0..1), critique (short string). No prose.\n"
								"Example:\nsuccess_prob: 0.72\nrisk: 0.18\ncritique: brief reason\n\n"
								"Trace (YAML or JSON content below; do NOT echo it):\n" + str(trace)
							)
							if log_prompts:
								_log("prompt", prompt)
							enc = tok(prompt, return_tensors="pt").to(peft_model.device)
							with torch.inference_mode():
								out_ids = peft_model.generate(
									**enc,
									max_new_tokens=200,
									do_sample=False,
									num_beams=1,
								)
							text = tok.decode(out_ids[0], skip_special_tokens=True)
							try:
								meter_cost("verify_llm", 1.0)
							except Exception:
								pass
							if log_prompts:
								_log("response", text)
							parsed: Dict[str, Any] = _robust_struct_parse(text)
						except Exception as e:
							if log_prompts:
								print(f"[HF VERIFY WARN] Adapter load failed: {e}")
							parsed = {}
							text = ""
					else:
						parsed = {}
						text = ""
					# Fallback to shared HF client (base model)
					if not parsed:
						# Optional structured decoding via Outlines (JSON schema)
						if structured_mode == "json":
							try:
								# Lazy import outlines; optional dependency
								from outlines import models as _out_models  # type: ignore
								from outlines import generate as _out_generate  # type: ignore
								schema = {
									"type": "object",
									"properties": {
										"success_prob": {"type": "number", "minimum": 0.0, "maximum": 1.0},
										"risk": {"type": "number", "minimum": 0.0, "maximum": 1.0},
										"critique": {"type": "string"},
										"validation_results": {
											"type": "object",
											"properties": {
												"state_requirements": {
													"type": "object",
													"properties": {
														"satisfied": {"type": "array", "items": {"type": "string"}},
														"missing": {"type": "array", "items": {"type": "string"}}
													}
												},
												"validation_points": {
													"type": "array",
													"items": {
														"type": "object",
														"properties": {
															"step": {"type": "integer"},
															"check": {"type": "string"},
															"passed": {"type": "boolean"},
															"details": {"type": "string"}
														},
														"required": ["step", "check", "passed"]
													}
												}
											}
										}
									},
									"required": ["success_prob", "risk", "critique"],
									"additionalProperties": False,
								}
								base_prompt = (
									"You are a strict verifier. Given a trace snippet, return ONLY a JSON object with:\n"
									"1. success_prob (0..1): Overall success probability\n"
									"2. risk (0..1): Overall risk assessment\n"
									"3. critique: Brief explanation\n"
									"4. validation_results:\n"
									"   - state_requirements: Track satisfied and missing state elements\n"
									"   - validation_points: Array of validation checks with step, check, passed, and details\n"
									"Example:\n"
									'{"success_prob":0.8,"risk":0.2,"critique":"Login successful","validation_results":{'
									'"state_requirements":{"satisfied":["login_page"],"missing":[]},'
									'"validation_points":[{"step":0,"check":"Form visible","passed":true,"details":"Found #login-form"}]}}\n\n'
									"Trace (YAML/JSON below; do NOT echo it):\n" + str(trace)
								)
								if log_prompts:
									_log("prompt", base_prompt)
								# Optional cache for structured verify
								use_cache = cache_enabled("AGI_VERIFY_CACHE")
								ttl = int(os.environ.get("AGI_VERIFY_CACHE_TTL", "1800") or 1800) if use_cache else 0
								text = ""
								if use_cache:
									cached = cache_get("verifier", [model, "json", base_prompt], ttl)
									if isinstance(cached, str) and cached.strip():
										text = cached
								if not text:
									mdl = _out_models.transformers(model)
									generator = _out_generate.json(mdl, schema)
									text = generator(base_prompt)
									try:
										meter_cost("verify_structured", 1.0)
									except Exception:
										pass
									if use_cache and text:
										cache_set("verifier", [model, "json", base_prompt], str(text))
									if log_prompts:
										_log("response", text)
									parsed = _robust_json_parse(text)
							except Exception:
								parsed = {}
								text = ""
						# Standard YAML-first path
						if not parsed:
							client = HFClient.get_cached(model)
							base_prompt = (
								"You are a strict verifier. Given a trace snippet, output ONLY a YAML mapping with keys"
								" success_prob (float 0..1), risk (float 0..1), critique (short string). No prose.\n"
								"Example:\n"
								"success_prob: 0.72\n"
								"risk: 0.18\n"
								"critique: brief reason\n\n"
								"Trace (YAML or JSON content below; do NOT echo it):\n" + str(trace)
							)
							if log_prompts:
								_log("prompt", base_prompt)
							# Optional cache for YAML verify
							use_cache = cache_enabled("AGI_VERIFY_CACHE")
							ttl = int(os.environ.get("AGI_VERIFY_CACHE_TTL", "1800") or 1800) if use_cache else 0
							text = ""
							if use_cache:
								cached = cache_get("verifier", [model, "yaml", base_prompt], ttl)
								if isinstance(cached, str) and cached.strip():
									text = cached
							if not text:
								text = client.generate(base_prompt, max_new_tokens=200, temperature=0.0)
								try:
									meter_cost("verify_yaml", 1.0)
								except Exception:
									pass
								if use_cache and text:
									cache_set("verifier", [model, "yaml", base_prompt], str(text))
							if log_prompts:
								_log("response", text)
							parsed = _robust_struct_parse(text)
				if "success_prob" not in parsed or "risk" not in parsed:
					retry_prompt = (
						"Return ONLY YAML with exactly these keys: success_prob, risk, critique. No backticks.\n"
						"Example:\nsuccess_prob: 0.5\nrisk: 0.5\ncritique: short\n\nNow respond:"
					)
					if log_prompts:
						_log("prompt", retry_prompt)
					# Optional cache for retry prompt
					use_cache = cache_enabled("AGI_VERIFY_CACHE")
					ttl = int(os.environ.get("AGI_VERIFY_CACHE_TTL", "1800") or 1800) if use_cache else 0
					text_retry = ""
					if use_cache:
						cached = cache_get("verifier", [model, "retry", retry_prompt], ttl)
						if isinstance(cached, str) and cached.strip():
							text_retry = cached
					if not text_retry:
						text_retry = client.generate(retry_prompt, max_new_tokens=120, temperature=0.0)
						try:
							meter_cost("verify_retry", 1.0)
						except Exception:
							pass
					if log_prompts:
						_log("response", text_retry)
					parsed = _robust_struct_parse(text_retry)
					if ("success_prob" not in parsed or "risk" not in parsed) and text_retry:
						coerced = coerce_flat_yaml(text_retry)
						if "success_prob" in coerced and "risk" in coerced:
							parsed = coerced
					# Final regex fallback before failing strict
					if ("success_prob" not in parsed or "risk" not in parsed) and (text_retry or text):
						fallback_src = text_retry or text
						m1 = re.search(r"success_prob\s*[:=]\s*([0-9]*\.?[0-9]+)", fallback_src)
						m2 = re.search(r"risk\s*[:=]\s*([0-9]*\.?[0-9]+)", fallback_src)
						if m1 and m2:
							parsed = {
								"success_prob": float(m1.group(1)),
								"risk": float(m2.group(1)),
								"critique": "",
							}
				# If LLM responded but still unparsable, return conservative defaults instead of failing hard
				if ("success_prob" not in parsed or "risk" not in parsed) and (text_retry or text):
					status = (trace.get("result") or {}).get("status")
					fallback_risk = 0.2 if status == "ok" else 0.6
					fallback_sp = 0.8 if status == "ok" else 0.3
					parsed = {"success_prob": fallback_sp, "risk": fallback_risk, "critique": "unparsable-llm"}
				if "success_prob" in parsed and "risk" in parsed:
					return {
						"success_prob": float(parsed.get("success_prob", 0.5)),
						"risk": float(parsed.get("risk", 0.5)),
						"critique": str(parsed.get("critique", "")),
					}
				if require_llm:
					raise RuntimeError("HF LLM returned no valid YAML verdict")
			except Exception as e:
				if require_llm:
					raise RuntimeError(f"HF LLM call failed: {e}")
		else:
			# HF-only
			pass
	# Heuristic fallback
	return {"success_prob": 0.5, "risk": 0.5, "critique": "no-llm"}

def _robust_json_parse(text: str) -> Dict[str, Any]:
	try:
		return json.loads(text)
	except Exception:
		return {}

def _robust_struct_parse(text: str) -> Dict[str, Any]:
	# YAML first
	if yaml is not None:
		try:
			y = yaml.safe_load(text)
			if isinstance(y, dict):
				return y
		except Exception:
			pass
	# JSON fallback
	t = text.strip()
	if t.startswith("{") and t.endswith("}"):
		try:
			return json.loads(t)
		except Exception:
			return {}
	start = t.find("{")
	end = t.rfind("}")
	if start != -1 and end != -1 and end > start:
		try:
			return json.loads(t[start : end + 1])
		except Exception:
			return {}
	return {}


from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional

@dataclass
class VerifierServiceConfig:
	model: str
	backend: str = "hf"
	adapter_dir: Optional[str] = None
	adapter_bank: Optional[str] = None
	structured_mode: str = "json"
	timeout_sec: int = 30
	strict: bool = False
	calibrate: bool = False
	calib_model: Optional[str] = None
	log_prompts: bool = False

def _resolve_adapter(root_dir: Path, bank_name: Optional[str], adapter_dir: Optional[str]) -> Optional[str]:
	if adapter_dir:
		return adapter_dir
	if not bank_name:
		return adapter_dir
	try:
		from agi_dw.core.llm.adapter_router import pick_from_bank  # type: ignore
		bank = pick_from_bank(root_dir, bank_name)
		return bank.get("verifier", adapter_dir)
	except Exception:
		return adapter_dir

def verify(trace: Dict[str, Any], cfg: VerifierServiceConfig) -> Dict[str, Any]:
	"""Run verifier with optional calibration. Returns {success_prob, risk, critique}."""
	root_dir = Path(__file__).resolve().parents[2]
	adapter_dir = _resolve_adapter(root_dir, cfg.adapter_bank, cfg.adapter_dir)
	res: Dict[str, Any]
	try:
		from agi_dw.core.verifier.llm_verifier import verify_trace_snippet  # type: ignore
		vres = verify_trace_snippet(
			trace,
			model=cfg.model,
			timeout_sec=int(cfg.timeout_sec),
			use_llm=True,
			require_llm=bool(cfg.strict),
			backend=cfg.backend,
			log_prompts=bool(cfg.log_prompts),
			adapter_dir=adapter_dir,
			structured_mode=str(cfg.structured_mode),
		)
		res = {
			"success_prob": float(vres.get("success_prob", 0.5)),
			"risk": float(vres.get("risk", 0.5)),
			"critique": str(vres.get("critique", "")),
		}
	except Exception:
		res = {"success_prob": 0.5, "risk": 0.5, "critique": "verify-error"}

	# Optional calibration
	if cfg.calibrate:
		try:
			cal_path = Path(cfg.calib_model or "")
			if cal_path.exists():
				from joblib import load as joblib_load  # type: ignore
				pack = joblib_load(cal_path)
				iso = pack.get("iso")
				if iso is not None and hasattr(iso, "predict"):
					orig = float(res.get("risk", 0.5))
					cal = float(iso.predict([orig])[0])
					res["risk"] = max(0.0, min(1.0, cal))
		except Exception:
			pass
	return res

def quick_risk(obs: Dict[str, Any], plan: Dict[str, Any], cfg: VerifierServiceConfig) -> float:
	"""Lightweight risk on (obs, plan) snippet."""
	trace = {"obs": obs, "plan": plan, "action": {}, "result": {"status": "pending"}}
	try:
		res = verify(trace, cfg)
		return float(res.get("risk", 0.5))
	except Exception:
		return 0.5


VerifierServiceConfig

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional

@dataclass
class VerifierServiceConfig:
	model: str
	backend: str = "hf"
	adapter_dir: Optional[str] = None
	adapter_bank: Optional[str] = None
	structured_mode: str = "json"
	timeout_sec: int = 30
	strict: bool = False
	calibrate: bool = False
	calib_model: Optional[str] = None
	log_prompts: bool = False

def _resolve_adapter(root_dir: Path, bank_name: Optional[str], adapter_dir: Optional[str]) -> Optional[str]:
	if adapter_dir:
		return adapter_dir
	if not bank_name:
		return adapter_dir
	try:
		from agi_dw.core.llm.adapter_router import pick_from_bank  # type: ignore
		bank = pick_from_bank(root_dir, bank_name)
		return bank.get("verifier", adapter_dir)
	except Exception:
		return adapter_dir

def verify(trace: Dict[str, Any], cfg: VerifierServiceConfig) -> Dict[str, Any]:
	"""Run verifier with optional calibration. Returns {success_prob, risk, critique}."""
	root_dir = Path(__file__).resolve().parents[2]
	adapter_dir = _resolve_adapter(root_dir, cfg.adapter_bank, cfg.adapter_dir)
	res: Dict[str, Any]


from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional

@dataclass
class VerifierServiceConfig:
	model: str
	backend: str = "hf"
	adapter_dir: Optional[str] = None
	adapter_bank: Optional[str] = None
	structured_mode: str = "json"
	timeout_sec: int = 30
	strict: bool = False
	calibrate: bool = False
	calib_model: Optional[str] = None
	log_prompts: bool = False

def _resolve_adapter(root_dir: Path, bank_name: Optional[str], adapter_dir: Optional[str]) -> Optional[str]:
	if adapter_dir:
		return adapter_dir
	if not bank_name:
		return adapter_dir
	try:
		from agi_dw.core.llm.adapter_router import pick_from_bank  # type: ignore
		bank = pick_from_bank(root_dir, bank_name)
		return bank.get("verifier", adapter_dir)
	except Exception:
		return adapter_dir

def verify(trace: Dict[str, Any], cfg: VerifierServiceConfig) -> Dict[str, Any]:
	"""Run verifier with optional calibration. Returns {success_prob, risk, critique}."""
	root_dir = Path(__file__).resolve().parents[2]
	adapter_dir = _resolve_adapter(root_dir, cfg.adapter_bank, cfg.adapter_dir)
	res: Dict[str, Any]
	try:
		from agi_dw.core.verifier.llm_verifier import verify_trace_snippet  # type: ignore
		vres = verify_trace_snippet(
			trace,
			model=cfg.model,
			timeout_sec=int(cfg.timeout_sec),
			use_llm=True,
			require_llm=bool(cfg.strict),
			backend=cfg.backend,
			log_prompts=bool(cfg.log_prompts),
			adapter_dir=adapter_dir,
			structured_mode=str(cfg.structured_mode),
		)


verify

	timeout_sec: int = 30
	strict: bool = False
	calibrate: bool = False
	calib_model: Optional[str] = None
	log_prompts: bool = False

def _resolve_adapter(root_dir: Path, bank_name: Optional[str], adapter_dir: Optional[str]) -> Optional[str]:
	if adapter_dir:
		return adapter_dir
	if not bank_name:
		return adapter_dir
	try:
		from agi_dw.core.llm.adapter_router import pick_from_bank  # type: ignore
		bank = pick_from_bank(root_dir, bank_name)
		return bank.get("verifier", adapter_dir)
	except Exception:
		return adapter_dir

def verify(trace: Dict[str, Any], cfg: VerifierServiceConfig) -> Dict[str, Any]:
	"""Run verifier with optional calibration. Returns {success_prob, risk, critique}."""
	root_dir = Path(__file__).resolve().parents[2]
	adapter_dir = _resolve_adapter(root_dir, cfg.adapter_bank, cfg.adapter_dir)
	res: Dict[str, Any]
	try:
		from agi_dw.core.verifier.llm_verifier import verify_trace_snippet  # type: ignore
		vres = verify_trace_snippet(
			trace,
			model=cfg.model,
			timeout_sec=int(cfg.timeout_sec),
			use_llm=True,
			require_llm=bool(cfg.strict),
			backend=cfg.backend,
			log_prompts=bool(cfg.log_prompts),
			adapter_dir=adapter_dir,
			structured_mode=str(cfg.structured_mode),
		)
		res = {
			"success_prob": float(vres.get("success_prob", 0.5)),
			"risk": float(vres.get("risk", 0.5)),
			"critique": str(vres.get("critique", "")),
		}
	except Exception:
		res = {"success_prob": 0.5, "risk": 0.5, "critique": "verify-error"}

	# Optional calibration
	if cfg.calibrate:
		try:
			cal_path = Path(cfg.calib_model or "")
			if cal_path.exists():
				from joblib import load as joblib_load  # type: ignore
				pack = joblib_load(cal_path)
				iso = pack.get("iso")
				if iso is not None and hasattr(iso, "predict"):
					orig = float(res.get("risk", 0.5))
					cal = float(iso.predict([orig])[0])
					res["risk"] = max(0.0, min(1.0, cal))
		except Exception:
			pass
	return res

def quick_risk(obs: Dict[str, Any], plan: Dict[str, Any], cfg: VerifierServiceConfig) -> float:
	"""Lightweight risk on (obs, plan) snippet."""
	trace = {"obs": obs, "plan": plan, "action": {}, "result": {"status": "pending"}}
	try:
		res = verify(trace, cfg)
		return float(res.get("risk", 0.5))
	except Exception:
		return 0.5


	except Exception:
		res = {"success_prob": 0.5, "risk": 0.5, "critique": "verify-error"}

	# Optional calibration
	if cfg.calibrate:
		try:
			cal_path = Path(cfg.calib_model or "")
			if cal_path.exists():
				from joblib import load as joblib_load  # type: ignore
				pack = joblib_load(cal_path)
				iso = pack.get("iso")
				if iso is not None and hasattr(iso, "predict"):
					orig = float(res.get("risk", 0.5))
					cal = float(iso.predict([orig])[0])
					res["risk"] = max(0.0, min(1.0, cal))
		except Exception:
			pass
	return res

def quick_risk(obs: Dict[str, Any], plan: Dict[str, Any], cfg: VerifierServiceConfig) -> float:
	"""Lightweight risk on (obs, plan) snippet."""
	trace = {"obs": obs, "plan": plan, "action": {}, "result": {"status": "pending"}}
	try:
		res = verify(trace, cfg)
		return float(res.get("risk", 0.5))
	except Exception:
		return 0.5


from __future__ import annotations
import logging

from typing import Dict, Any

class AgentRole:
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:  # noqa: D401
		"""Perform role-specific action given a context and return a message."""
		raise NotImplementedError

class WebBrowserRole(AgentRole):
	def __init__(self, t5_model_path: str, il_path: str | None = None) -> None:
		self.t5_model_path = t5_model_path
		self.il_path = il_path

	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RepairConfig, select_action  # type: ignore
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		act_cfg = ActuatorConfig(mode="t5", t5_model=str(self.t5_model_path), il_path=str(self.il_path or ""), dom_structured=True)
		repair = RepairConfig(domain="dom", prefer_obs_args=True, default_url=str((obs.get("meta") or {}).get("url", "")), default_selector=str((obs.get("meta") or {}).get("selector", "")))
		extra = RouterExtras(domain="dom")
		action, _ = select_action(obs, plan, act_cfg, extra, verifier_cfg=None, wm_prior_cfg=None, wm_screen_cfg=None, repair_cfg=repair)
		return {"role": "WebBrowser", "action": action}

class CoderRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		# Placeholder: would integrate dev loop; for now emit a stub suggestion
		return {"role": "Coder", "suggestion": "Consider writing or running tests, then implement minimal fix."}

class DocWriterRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		return {"role": "DocWriter", "summary": {
			"task": str(obs.get("content", "")),
			"steps": plan.get("subgoals") or plan.get("steps") or [],
		}}


AgentRole

from __future__ import annotations
import logging

from typing import Dict, Any

class AgentRole:
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:  # noqa: D401
		"""Perform role-specific action given a context and return a message."""
		raise NotImplementedError

class WebBrowserRole(AgentRole):
	def __init__(self, t5_model_path: str, il_path: str | None = None) -> None:
		self.t5_model_path = t5_model_path
		self.il_path = il_path

	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RepairConfig, select_action  # type: ignore
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		act_cfg = ActuatorConfig(mode="t5", t5_model=str(self.t5_model_path), il_path=str(self.il_path or ""), dom_structured=True)
		repair = RepairConfig(domain="dom", prefer_obs_args=True, default_url=str((obs.get("meta") or {}).get("url", "")), default_selector=str((obs.get("meta") or {}).get("selector", "")))
		extra = RouterExtras(domain="dom")
		action, _ = select_action(obs, plan, act_cfg, extra, verifier_cfg=None, wm_prior_cfg=None, wm_screen_cfg=None, repair_cfg=repair)
		return {"role": "WebBrowser", "action": action}

class CoderRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:


CoderRole

		"""Perform role-specific action given a context and return a message."""
		raise NotImplementedError

class WebBrowserRole(AgentRole):
	def __init__(self, t5_model_path: str, il_path: str | None = None) -> None:
		self.t5_model_path = t5_model_path
		self.il_path = il_path

	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RepairConfig, select_action  # type: ignore
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		act_cfg = ActuatorConfig(mode="t5", t5_model=str(self.t5_model_path), il_path=str(self.il_path or ""), dom_structured=True)
		repair = RepairConfig(domain="dom", prefer_obs_args=True, default_url=str((obs.get("meta") or {}).get("url", "")), default_selector=str((obs.get("meta") or {}).get("selector", "")))
		extra = RouterExtras(domain="dom")
		action, _ = select_action(obs, plan, act_cfg, extra, verifier_cfg=None, wm_prior_cfg=None, wm_screen_cfg=None, repair_cfg=repair)
		return {"role": "WebBrowser", "action": action}

class CoderRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		# Placeholder: would integrate dev loop; for now emit a stub suggestion
		return {"role": "Coder", "suggestion": "Consider writing or running tests, then implement minimal fix."}

class DocWriterRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		return {"role": "DocWriter", "summary": {
			"task": str(obs.get("content", "")),
			"steps": plan.get("subgoals") or plan.get("steps") or [],
		}}


DocWriterRole

		self.t5_model_path = t5_model_path
		self.il_path = il_path

	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RepairConfig, select_action  # type: ignore
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		act_cfg = ActuatorConfig(mode="t5", t5_model=str(self.t5_model_path), il_path=str(self.il_path or ""), dom_structured=True)
		repair = RepairConfig(domain="dom", prefer_obs_args=True, default_url=str((obs.get("meta") or {}).get("url", "")), default_selector=str((obs.get("meta") or {}).get("selector", "")))
		extra = RouterExtras(domain="dom")
		action, _ = select_action(obs, plan, act_cfg, extra, verifier_cfg=None, wm_prior_cfg=None, wm_screen_cfg=None, repair_cfg=repair)
		return {"role": "WebBrowser", "action": action}

class CoderRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		# Placeholder: would integrate dev loop; for now emit a stub suggestion
		return {"role": "Coder", "suggestion": "Consider writing or running tests, then implement minimal fix."}

class DocWriterRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		return {"role": "DocWriter", "summary": {
			"task": str(obs.get("content", "")),
			"steps": plan.get("subgoals") or plan.get("steps") or [],
		}}


act

		self.il_path = il_path

	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RepairConfig, select_action  # type: ignore
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		act_cfg = ActuatorConfig(mode="t5", t5_model=str(self.t5_model_path), il_path=str(self.il_path or ""), dom_structured=True)
		repair = RepairConfig(domain="dom", prefer_obs_args=True, default_url=str((obs.get("meta") or {}).get("url", "")), default_selector=str((obs.get("meta") or {}).get("selector", "")))
		extra = RouterExtras(domain="dom")
		action, _ = select_action(obs, plan, act_cfg, extra, verifier_cfg=None, wm_prior_cfg=None, wm_screen_cfg=None, repair_cfg=repair)
		return {"role": "WebBrowser", "action": action}

class CoderRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		# Placeholder: would integrate dev loop; for now emit a stub suggestion
		return {"role": "Coder", "suggestion": "Consider writing or running tests, then implement minimal fix."}

class DocWriterRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		return {"role": "DocWriter", "summary": {
			"task": str(obs.get("content", "")),
			"steps": plan.get("subgoals") or plan.get("steps") or [],
		}}


from __future__ import annotations
import logging

from typing import Dict, Any

class AgentRole:
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:  # noqa: D401
		"""Perform role-specific action given a context and return a message."""
		raise NotImplementedError

class WebBrowserRole(AgentRole):
	def __init__(self, t5_model_path: str, il_path: str | None = None) -> None:
		self.t5_model_path = t5_model_path
		self.il_path = il_path

	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RepairConfig, select_action  # type: ignore
		obs = context.get("obs", {})
		plan = context.get("plan", {})
		act_cfg = ActuatorConfig(mode="t5", t5_model=str(self.t5_model_path), il_path=str(self.il_path or ""), dom_structured=True)
		repair = RepairConfig(domain="dom", prefer_obs_args=True, default_url=str((obs.get("meta") or {}).get("url", "")), default_selector=str((obs.get("meta") or {}).get("selector", "")))
		extra = RouterExtras(domain="dom")
		action, _ = select_action(obs, plan, act_cfg, extra, verifier_cfg=None, wm_prior_cfg=None, wm_screen_cfg=None, repair_cfg=repair)
		return {"role": "WebBrowser", "action": action}

class CoderRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:
		# Placeholder: would integrate dev loop; for now emit a stub suggestion
		return {"role": "Coder", "suggestion": "Consider writing or running tests, then implement minimal fix."}

class DocWriterRole(AgentRole):
	def act(self, context: Dict[str, Any]) -> Dict[str, Any]:


from typing import Dict, List
import math
import torch
import torch.nn as nn

from .curvature import update_curvature_diag, sophia_update_with_curv
from .bandit import UCB1

class MixtureMetaOptGrouped:
    def __init__(self, model: nn.Module, param_groups: List[Dict], gate: nn.Module, base_lr: float = 3e-4, base_wd: float = 0.01, device: str = "cuda"):
        self.model = model
        self.groups = param_groups
        self.gate = gate.to(device)
        self.base_lr = base_lr
        self.base_wd = base_wd
        self.device = device

        # per-param states
        self.st_adam: Dict[int, Dict[str, torch.Tensor]] = {}
        self.st_lion: Dict[int, Dict[str, torch.Tensor]] = {}
        self.st_soph: Dict[int, Dict[str, torch.Tensor]] = {}
        self.curv: Dict[int, torch.Tensor] = {}

        # group stats
        self.grad_ema: List[float] = [0.0] * len(self.groups)
        self.last_loss: float | None = None
        self.step_idx: int = 0
        self.T_hint: int = 1_000_000

        # bandits per group and last chosen arm storage
        self.bandits: List[UCB1] = [UCB1(n_arms=3) for _ in self.groups]
        self._last_arm: Dict[int, int] = {}

    def _group_grad_norm(self, params: List[torch.Tensor]) -> float:
        s = 0.0
        for p in params:
            if p.grad is None:
                continue
            s += p.grad.detach().float().pow(2).sum().item()
        return float(math.sqrt(s + 1e-12))

    def _features(self, loss: float, gnorm: float, gi: int) -> torch.Tensor:
        self.grad_ema[gi] = 0.95 * self.grad_ema[gi] + 0.05 * gnorm
        dloss = 0.0 if self.last_loss is None else (loss - self.last_loss)
        cos_proxy = (gnorm / (self.grad_ema[gi] + 1e-8))
        t_feat = float(self.step_idx / max(1, self.T_hint))
        self.last_loss = loss
        return torch.tensor([
            [
                float(math.log(gnorm + 1e-8)),
                float(math.log(self.grad_ema[gi] + 1e-8)),
                float(loss),
                float(dloss),
                float(cos_proxy),
                t_feat,
            ]
        ], device=self.device)

    @torch.no_grad()
    def _apply_blended_updates(self, params: List[torch.Tensor], mix: torch.Tensor, lr: float, wd: float) -> None:
        w_adam, w_lion, w_soph = mix.tolist()
        for p in params:
            if p.grad is None:
                continue
            pid = id(p)
            g = p.grad
            p0 = p.detach().clone()

            # AdamW proposal
            stA = self.st_adam.setdefault(pid, {})
            m = stA.setdefault("m", torch.zeros_like(p))
            v = stA.setdefault("v", torch.zeros_like(p))
            t = stA.setdefault("t", 0) + 1
            stA["t"] = t
            m.mul_(0.9).add_(g, alpha=0.1)
            v.mul_(0.999).addcmul_(g, g, value=0.001)
            m_hat = m / (1 - 0.9 ** t)
            v_hat = v / (1 - 0.999 ** t)
            p_adam = p0.clone()
            if wd != 0.0:
                p_adam.add_(p_adam, alpha=-lr * wd)
            p_adam.addcdiv_(m_hat, v_hat.sqrt().add_(1e-8), value=-lr)

            # Lion proposal
            stL = self.st_lion.setdefault(pid, {})
            mL = stL.setdefault("m", torch.zeros_like(p))
            u = 0.9 * mL + 0.1 * g
            p_lion = p0.clone()
            if wd != 0.0:
                p_lion.add_(p_lion, alpha=-lr * wd)
            p_lion.add_(u.sign(), alpha=-lr)
            mL.mul_(0.99).add_(g, alpha=0.01)

            # Sophia proposal using diagonal curvature
            stS = self.st_soph.setdefault(pid, {})
            h = self.curv.get(pid, torch.ones_like(p))
            p_soph = p0.clone()
            sophia_update_with_curv(p_soph, g, stS, lr=lr, wd=wd, curv_diag=h)

            # Blend
            p.copy_(w_adam * p_adam + w_lion * p_lion + w_soph * p_soph)

    def step(self, loss: float) -> None:
        self.step_idx += 1
        # Update curvature diag (Fisher by default)
        update_curvature_diag(self.model, self.curv, beta=0.99, mode="fisher")

        for gi, grp in enumerate(self.groups):
            params = grp["params"]
            gnorm = self._group_grad_norm(params)
            feats = self._features(loss, gnorm, gi)
            group_idx = torch.tensor([gi], device=self.device)

            mix_logits, lrs, wds = self.gate(feats, group_idx)
            arm = self.bandits[gi].select()
            self._last_arm[gi] = arm
            boost = torch.zeros_like(mix_logits)
            boost[0, arm] += 0.75  # gentle prior toward bandit choice
            mix = (mix_logits + boost).softmax(-1).squeeze(0)

            lr = float(self.base_lr * lrs.item())
            wd = float(self.base_wd * wds.item())
            self._apply_blended_updates(params, mix, lr, wd)

    def record_bandit_reward(self, gi: int, reward: float) -> None:
        arm = self._last_arm.get(gi)
        if arm is not None:
            self.bandits[gi].update(arm, reward)


from typing import Dict, List
import math
import torch
import torch.nn as nn

from .curvature import update_curvature_diag, sophia_update_with_curv
from .bandit import UCB1

class MixtureMetaOptGrouped:
    def __init__(self, model: nn.Module, param_groups: List[Dict], gate: nn.Module, base_lr: float = 3e-4, base_wd: float = 0.01, device: str = "cuda"):
        self.model = model
        self.groups = param_groups
        self.gate = gate.to(device)
        self.base_lr = base_lr
        self.base_wd = base_wd
        self.device = device

        # per-param states
        self.st_adam: Dict[int, Dict[str, torch.Tensor]] = {}
        self.st_lion: Dict[int, Dict[str, torch.Tensor]] = {}
        self.st_soph: Dict[int, Dict[str, torch.Tensor]] = {}
        self.curv: Dict[int, torch.Tensor] = {}

        # group stats
        self.grad_ema: List[float] = [0.0] * len(self.groups)
        self.last_loss: float | None = None
        self.step_idx: int = 0
        self.T_hint: int = 1_000_000

        # bandits per group and last chosen arm storage
        self.bandits: List[UCB1] = [UCB1(n_arms=3) for _ in self.groups]
        self._last_arm: Dict[int, int] = {}

    def _group_grad_norm(self, params: List[torch.Tensor]) -> float:
        s = 0.0
        for p in params:
            if p.grad is None:
                continue
            s += p.grad.detach().float().pow(2).sum().item()
        return float(math.sqrt(s + 1e-12))

    def _features(self, loss: float, gnorm: float, gi: int) -> torch.Tensor:
        self.grad_ema[gi] = 0.95 * self.grad_ema[gi] + 0.05 * gnorm
        dloss = 0.0 if self.last_loss is None else (loss - self.last_loss)
        cos_proxy = (gnorm / (self.grad_ema[gi] + 1e-8))
        t_feat = float(self.step_idx / max(1, self.T_hint))
        self.last_loss = loss
        return torch.tensor([
            [
                float(math.log(gnorm + 1e-8)),
                float(math.log(self.grad_ema[gi] + 1e-8)),
                float(loss),


        self.base_lr = base_lr
        self.base_wd = base_wd
        self.device = device

        # per-param states
        self.st_adam: Dict[int, Dict[str, torch.Tensor]] = {}
        self.st_lion: Dict[int, Dict[str, torch.Tensor]] = {}
        self.st_soph: Dict[int, Dict[str, torch.Tensor]] = {}
        self.curv: Dict[int, torch.Tensor] = {}

        # group stats
        self.grad_ema: List[float] = [0.0] * len(self.groups)
        self.last_loss: float | None = None
        self.step_idx: int = 0
        self.T_hint: int = 1_000_000

        # bandits per group and last chosen arm storage
        self.bandits: List[UCB1] = [UCB1(n_arms=3) for _ in self.groups]
        self._last_arm: Dict[int, int] = {}

    def _group_grad_norm(self, params: List[torch.Tensor]) -> float:
        s = 0.0
        for p in params:
            if p.grad is None:
                continue
            s += p.grad.detach().float().pow(2).sum().item()
        return float(math.sqrt(s + 1e-12))

    def _features(self, loss: float, gnorm: float, gi: int) -> torch.Tensor:
        self.grad_ema[gi] = 0.95 * self.grad_ema[gi] + 0.05 * gnorm
        dloss = 0.0 if self.last_loss is None else (loss - self.last_loss)
        cos_proxy = (gnorm / (self.grad_ema[gi] + 1e-8))
        t_feat = float(self.step_idx / max(1, self.T_hint))
        self.last_loss = loss
        return torch.tensor([
            [
                float(math.log(gnorm + 1e-8)),
                float(math.log(self.grad_ema[gi] + 1e-8)),
                float(loss),
                float(dloss),
                float(cos_proxy),
                t_feat,
            ]
        ], device=self.device)

    @torch.no_grad()
    def _apply_blended_updates(self, params: List[torch.Tensor], mix: torch.Tensor, lr: float, wd: float) -> None:


        self.curv: Dict[int, torch.Tensor] = {}

        # group stats
        self.grad_ema: List[float] = [0.0] * len(self.groups)
        self.last_loss: float | None = None
        self.step_idx: int = 0
        self.T_hint: int = 1_000_000

        # bandits per group and last chosen arm storage
        self.bandits: List[UCB1] = [UCB1(n_arms=3) for _ in self.groups]
        self._last_arm: Dict[int, int] = {}

    def _group_grad_norm(self, params: List[torch.Tensor]) -> float:
        s = 0.0
        for p in params:
            if p.grad is None:
                continue
            s += p.grad.detach().float().pow(2).sum().item()
        return float(math.sqrt(s + 1e-12))

    def _features(self, loss: float, gnorm: float, gi: int) -> torch.Tensor:
        self.grad_ema[gi] = 0.95 * self.grad_ema[gi] + 0.05 * gnorm
        dloss = 0.0 if self.last_loss is None else (loss - self.last_loss)
        cos_proxy = (gnorm / (self.grad_ema[gi] + 1e-8))
        t_feat = float(self.step_idx / max(1, self.T_hint))
        self.last_loss = loss
        return torch.tensor([
            [
                float(math.log(gnorm + 1e-8)),
                float(math.log(self.grad_ema[gi] + 1e-8)),
                float(loss),
                float(dloss),
                float(cos_proxy),
                t_feat,
            ]
        ], device=self.device)

    @torch.no_grad()
    def _apply_blended_updates(self, params: List[torch.Tensor], mix: torch.Tensor, lr: float, wd: float) -> None:
        w_adam, w_lion, w_soph = mix.tolist()
        for p in params:
            if p.grad is None:
                continue
            pid = id(p)
            g = p.grad
            p0 = p.detach().clone()

            # AdamW proposal
            stA = self.st_adam.setdefault(pid, {})
            m = stA.setdefault("m", torch.zeros_like(p))
            v = stA.setdefault("v", torch.zeros_like(p))
            t = stA.setdefault("t", 0) + 1
            stA["t"] = t
            m.mul_(0.9).add_(g, alpha=0.1)
            v.mul_(0.999).addcmul_(g, g, value=0.001)
            m_hat = m / (1 - 0.9 ** t)


        return float(math.sqrt(s + 1e-12))

    def _features(self, loss: float, gnorm: float, gi: int) -> torch.Tensor:
        self.grad_ema[gi] = 0.95 * self.grad_ema[gi] + 0.05 * gnorm
        dloss = 0.0 if self.last_loss is None else (loss - self.last_loss)
        cos_proxy = (gnorm / (self.grad_ema[gi] + 1e-8))
        t_feat = float(self.step_idx / max(1, self.T_hint))
        self.last_loss = loss
        return torch.tensor([
            [
                float(math.log(gnorm + 1e-8)),
                float(math.log(self.grad_ema[gi] + 1e-8)),
                float(loss),
                float(dloss),
                float(cos_proxy),
                t_feat,
            ]
        ], device=self.device)

    @torch.no_grad()
    def _apply_blended_updates(self, params: List[torch.Tensor], mix: torch.Tensor, lr: float, wd: float) -> None:
        w_adam, w_lion, w_soph = mix.tolist()
        for p in params:
            if p.grad is None:
                continue
            pid = id(p)
            g = p.grad
            p0 = p.detach().clone()

            # AdamW proposal
            stA = self.st_adam.setdefault(pid, {})
            m = stA.setdefault("m", torch.zeros_like(p))
            v = stA.setdefault("v", torch.zeros_like(p))
            t = stA.setdefault("t", 0) + 1
            stA["t"] = t
            m.mul_(0.9).add_(g, alpha=0.1)
            v.mul_(0.999).addcmul_(g, g, value=0.001)
            m_hat = m / (1 - 0.9 ** t)
            v_hat = v / (1 - 0.999 ** t)
            p_adam = p0.clone()
            if wd != 0.0:
                p_adam.add_(p_adam, alpha=-lr * wd)
            p_adam.addcdiv_(m_hat, v_hat.sqrt().add_(1e-8), value=-lr)

            # Lion proposal
            stL = self.st_lion.setdefault(pid, {})
            mL = stL.setdefault("m", torch.zeros_like(p))
            u = 0.9 * mL + 0.1 * g
            p_lion = p0.clone()
            if wd != 0.0:
                p_lion.add_(p_lion, alpha=-lr * wd)
            p_lion.add_(u.sign(), alpha=-lr)
            mL.mul_(0.99).add_(g, alpha=0.01)

            # Sophia proposal using diagonal curvature
            stS = self.st_soph.setdefault(pid, {})
            h = self.curv.get(pid, torch.ones_like(p))
            p_soph = p0.clone()
            sophia_update_with_curv(p_soph, g, stS, lr=lr, wd=wd, curv_diag=h)

            # Blend
            p.copy_(w_adam * p_adam + w_lion * p_lion + w_soph * p_soph)

    def step(self, loss: float) -> None:
        self.step_idx += 1
        # Update curvature diag (Fisher by default)
        update_curvature_diag(self.model, self.curv, beta=0.99, mode="fisher")

        for gi, grp in enumerate(self.groups):
            params = grp["params"]
            gnorm = self._group_grad_norm(params)
            feats = self._features(loss, gnorm, gi)
            group_idx = torch.tensor([gi], device=self.device)

            mix_logits, lrs, wds = self.gate(feats, group_idx)
            arm = self.bandits[gi].select()
            self._last_arm[gi] = arm
            boost = torch.zeros_like(mix_logits)
            boost[0, arm] += 0.75  # gentle prior toward bandit choice
            mix = (mix_logits + boost).softmax(-1).squeeze(0)

            lr = float(self.base_lr * lrs.item())


step

            # Lion proposal
            stL = self.st_lion.setdefault(pid, {})
            mL = stL.setdefault("m", torch.zeros_like(p))
            u = 0.9 * mL + 0.1 * g
            p_lion = p0.clone()
            if wd != 0.0:
                p_lion.add_(p_lion, alpha=-lr * wd)
            p_lion.add_(u.sign(), alpha=-lr)
            mL.mul_(0.99).add_(g, alpha=0.01)

            # Sophia proposal using diagonal curvature
            stS = self.st_soph.setdefault(pid, {})
            h = self.curv.get(pid, torch.ones_like(p))
            p_soph = p0.clone()
            sophia_update_with_curv(p_soph, g, stS, lr=lr, wd=wd, curv_diag=h)

            # Blend
            p.copy_(w_adam * p_adam + w_lion * p_lion + w_soph * p_soph)

    def step(self, loss: float) -> None:
        self.step_idx += 1
        # Update curvature diag (Fisher by default)
        update_curvature_diag(self.model, self.curv, beta=0.99, mode="fisher")

        for gi, grp in enumerate(self.groups):
            params = grp["params"]
            gnorm = self._group_grad_norm(params)
            feats = self._features(loss, gnorm, gi)
            group_idx = torch.tensor([gi], device=self.device)

            mix_logits, lrs, wds = self.gate(feats, group_idx)
            arm = self.bandits[gi].select()
            self._last_arm[gi] = arm
            boost = torch.zeros_like(mix_logits)
            boost[0, arm] += 0.75  # gentle prior toward bandit choice
            mix = (mix_logits + boost).softmax(-1).squeeze(0)

            lr = float(self.base_lr * lrs.item())
            wd = float(self.base_wd * wds.item())
            self._apply_blended_updates(params, mix, lr, wd)

    def record_bandit_reward(self, gi: int, reward: float) -> None:
        arm = self._last_arm.get(gi)
        if arm is not None:
            self.bandits[gi].update(arm, reward)


        # Update curvature diag (Fisher by default)
        update_curvature_diag(self.model, self.curv, beta=0.99, mode="fisher")

        for gi, grp in enumerate(self.groups):
            params = grp["params"]
            gnorm = self._group_grad_norm(params)
            feats = self._features(loss, gnorm, gi)
            group_idx = torch.tensor([gi], device=self.device)

            mix_logits, lrs, wds = self.gate(feats, group_idx)
            arm = self.bandits[gi].select()
            self._last_arm[gi] = arm
            boost = torch.zeros_like(mix_logits)
            boost[0, arm] += 0.75  # gentle prior toward bandit choice
            mix = (mix_logits + boost).softmax(-1).squeeze(0)

            lr = float(self.base_lr * lrs.item())
            wd = float(self.base_wd * wds.item())
            self._apply_blended_updates(params, mix, lr, wd)

    def record_bandit_reward(self, gi: int, reward: float) -> None:
        arm = self._last_arm.get(gi)
        if arm is not None:
            self.bandits[gi].update(arm, reward)


from typing import List, Dict, Any, Set
import torch.nn as nn

def make_param_groups(transformer_model: nn.Module) -> List[Dict[str, Any]]:
    """
    Build parameter groups for layer-wise gating.
    - Embeddings
    - Per transformer block (GPT-style `.h` or common `.layers`)
    - Output head
    - Misc (fallback)
    """
    groups: List[Dict[str, Any]] = []

    # Embeddings
    emb_params = []
    for n, p in transformer_model.named_parameters():
        if ("embed" in n) or ("wte" in n) or ("wpe" in n) or ("embeddings" in n):
            emb_params.append(p)
    if emb_params:
        groups.append({"name": "emb", "params": emb_params})

    # Per-transformer block (try GPT `.h`, then generic `.layers`)
    blocks = list(getattr(transformer_model, "h", []))
    if not blocks and hasattr(transformer_model, "layers"):
        try:
            blocks = list(transformer_model.layers)
        except Exception:
            blocks = []
    for li, block in enumerate(blocks):
        gparams = [p for _, p in block.named_parameters()]
        if gparams:
            groups.append({"name": f"block_{li}", "params": gparams})

    # Output head
    head_params = []
    for n, p in transformer_model.named_parameters():
        if ("lm_head" in n) or n.endswith("out_proj.weight") or n.endswith("out_proj.bias"):
            head_params.append(p)
    if head_params:
        groups.append({"name": "head", "params": head_params})

    # Fallback for anything ungrouped
    grouped: Set[int] = set([id(p) for g in groups for p in g["params"]])
    rest = [p for p in transformer_model.parameters() if id(p) not in grouped]
    if rest:
        groups.append({"name": "misc", "params": rest})

    return groups


import torch
from typing import Dict

@torch.no_grad()
def update_curvature_diag(model: torch.nn.Module, curv_state: Dict[int, torch.Tensor], beta: float = 0.99, mode: str = "fisher", probe_prob: float = 0.05) -> None:
    """
    Maintain a diagonal curvature proxy per-parameter via EMA.

    - mode="fisher": update with on-policy gradient squares (Fisher diagonal proxy)
    - mode="hutch": occasionally nudge scale using a Hutchinson-like probe (stabilizer)
    """
    for i, p in enumerate(model.parameters()):
        if p.grad is None:
            continue
        g = p.grad.detach()
        h = curv_state.setdefault(id(p), torch.zeros_like(p))
        # Fisher diagonal EMA
        h.mul_(beta).addcmul_(g, g, value=(1.0 - beta))

        # Optional: cheap Hutchinson probe stabilizer (no closure to recompute grads here)
        if mode == "hutch" and torch.rand(()) < probe_prob:
            v = torch.empty_like(p).bernoulli_(0.5).mul_(2.0).add_(-1.0)  # Rademacher
            # Without a closure, approximate by gentle scale nudging
            h.mul_(0.99).add_(0.01 * (v * v))

@torch.no_grad()
def sophia_update_with_curv(p: torch.Tensor, g: torch.Tensor, state: dict, lr: float, wd: float, curv_diag: torch.Tensor, eps: float = 1e-8, temp: float = 1.0) -> None:
    """
    Sophia-style diagonal-curvature scaled update with decoupled weight decay.
    """
    if wd != 0.0:
        p.add_(p, alpha=-lr * wd)
    denom = curv_diag.abs().mul_(temp).add_(eps)
    p.addcdiv_(g, denom, value=-lr)


    - mode="fisher": update with on-policy gradient squares (Fisher diagonal proxy)
    - mode="hutch": occasionally nudge scale using a Hutchinson-like probe (stabilizer)
    """
    for i, p in enumerate(model.parameters()):
        if p.grad is None:
            continue
        g = p.grad.detach()
        h = curv_state.setdefault(id(p), torch.zeros_like(p))
        # Fisher diagonal EMA
        h.mul_(beta).addcmul_(g, g, value=(1.0 - beta))

        # Optional: cheap Hutchinson probe stabilizer (no closure to recompute grads here)
        if mode == "hutch" and torch.rand(()) < probe_prob:
            v = torch.empty_like(p).bernoulli_(0.5).mul_(2.0).add_(-1.0)  # Rademacher
            # Without a closure, approximate by gentle scale nudging
            h.mul_(0.99).add_(0.01 * (v * v))

@torch.no_grad()
def sophia_update_with_curv(p: torch.Tensor, g: torch.Tensor, state: dict, lr: float, wd: float, curv_diag: torch.Tensor, eps: float = 1e-8, temp: float = 1.0) -> None:
    """
    Sophia-style diagonal-curvature scaled update with decoupled weight decay.
    """
    if wd != 0.0:
        p.add_(p, alpha=-lr * wd)
    denom = curv_diag.abs().mul_(temp).add_(eps)
    p.addcdiv_(g, denom, value=-lr)


import math
from typing import List

class UCB1:
    def __init__(self, n_arms: int = 3):
        self.n: List[int] = [0] * n_arms
        self.value: List[float] = [0.0] * n_arms
        self.total: int = 0

    def select(self, c: float = 2.0) -> int:
        self.total += 1
        for a in range(len(self.n)):
            if self.n[a] == 0:
                return a
        ucb = [self.value[a] + c * math.sqrt(max(1.0, math.log(self.total)) / self.n[a]) for a in range(len(self.n))]
        return max(range(len(self.n)), key=lambda a: ucb[a])

    def update(self, arm: int, reward: float) -> None:
        self.n[arm] += 1
        n = self.n[arm]
        self.value[arm] += (reward - self.value[arm]) / n


import torch
import torch.nn as nn

class GateNetGrouped(nn.Module):
    def __init__(self, n_groups: int, hidden: int = 128):
        super().__init__()
        self.group_embed = nn.Embedding(n_groups, 16)
        self.mlp = nn.Sequential(
            nn.Linear(6 + 16, hidden), nn.GELU(),
            nn.Linear(hidden, hidden), nn.GELU(),
        )
        self.mix_head = nn.Linear(hidden, 3)  # AdamW, Lion, Sophia
        self.lrs_head = nn.Linear(hidden, 1)
        self.wds_head = nn.Linear(hidden, 1)

    def forward(self, feats: torch.Tensor, group_idx: torch.Tensor):
        # feats: [B, 6], group_idx: [B]
        ge = self.group_embed(group_idx)  # [B, 16]
        z = torch.cat([feats, ge], dim=-1)
        h = self.mlp(z)
        mix_logits = self.mix_head(h)
        lr_scale = torch.sigmoid(self.lrs_head(h)) * 2.0
        wd_scale = torch.sigmoid(self.wds_head(h)) * 2.0
        return mix_logits, lr_scale.squeeze(-1), wd_scale.squeeze(-1)


from __future__ import annotations
import torch
from torch.nn.utils import clip_grad_norm_
from typing import Any, Dict

try:
    from transformers import Trainer  # type: ignore
except Exception:  # pragma: no cover
    Trainer = object  # type: ignore

from .group_utils import make_param_groups
from .gate_grouped import GateNetGrouped
from .metaopt_grouped import MixtureMetaOptGrouped

class MetaOptTrainer(Trainer):  # type: ignore
    """
    HuggingFace Trainer subclass that applies MixtureMetaOptGrouped instead of optimizer.step.
    - Computes loss and backward via standard Trainer infra (Accelerate-aware)
    - Clips grads, then calls meta-optimizer step
    - Skips the regular optimizer step when enabled
    """

    def __init__(self, *args: Any, metaopt: bool = False, meta_base_lr: float = 3e-4, meta_base_wd: float = 0.01, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._metaopt_enabled: bool = bool(metaopt)
        self._meta: MixtureMetaOptGrouped | None = None
        if self._metaopt_enabled:
            device = "cuda" if torch.cuda.is_available() else "cpu"
            groups = make_param_groups(self.model)
            gate = GateNetGrouped(n_groups=len(groups), hidden=128)
            self._meta = MixtureMetaOptGrouped(self.model, groups, gate, base_lr=float(meta_base_lr), base_wd=float(meta_base_wd), device=device)

    def training_step(self, model: torch.nn.Module, inputs: Dict[str, Any], num_items_in_batch: int | None = None, *args: Any, **kwargs: Any) -> torch.Tensor:  # type: ignore[override]
        if not self._metaopt_enabled or self._meta is None:
            # Delegate to HF Trainer (supports extra args like num_items_in_batch)
            return super().training_step(model, inputs, *(() if num_items_in_batch is None else (num_items_in_batch,)), *args, **kwargs)  # type: ignore[misc]

        model.train()
        inputs = self._prepare_inputs(inputs)
        outputs = model(**inputs)
        loss = outputs.get("loss") if isinstance(outputs, dict) else outputs[0]

        # Backward using accelerator for mixed precision/DDP compatibility
        self.accelerator.backward(loss)
        # Clip then meta-step
        try:
            clip_grad_norm_(model.parameters(), 1.0)
        except Exception:
            pass
        self._meta.step(float(loss.detach().item()))
        return loss.detach()

    def optimizer_step(self, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
        if not self._metaopt_enabled:
            return super().optimizer_step(*args, **kwargs)  # type: ignore[misc]
        # Skip regular optimizer stepping when meta-optimizer is enabled
        return None


from typing import Any, Dict

try:
    from transformers import Trainer  # type: ignore
except Exception:  # pragma: no cover
    Trainer = object  # type: ignore

from .group_utils import make_param_groups
from .gate_grouped import GateNetGrouped
from .metaopt_grouped import MixtureMetaOptGrouped

class MetaOptTrainer(Trainer):  # type: ignore
    """
    HuggingFace Trainer subclass that applies MixtureMetaOptGrouped instead of optimizer.step.
    - Computes loss and backward via standard Trainer infra (Accelerate-aware)
    - Clips grads, then calls meta-optimizer step
    - Skips the regular optimizer step when enabled
    """

    def __init__(self, *args: Any, metaopt: bool = False, meta_base_lr: float = 3e-4, meta_base_wd: float = 0.01, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._metaopt_enabled: bool = bool(metaopt)
        self._meta: MixtureMetaOptGrouped | None = None
        if self._metaopt_enabled:
            device = "cuda" if torch.cuda.is_available() else "cpu"
            groups = make_param_groups(self.model)
            gate = GateNetGrouped(n_groups=len(groups), hidden=128)
            self._meta = MixtureMetaOptGrouped(self.model, groups, gate, base_lr=float(meta_base_lr), base_wd=float(meta_base_wd), device=device)

    def training_step(self, model: torch.nn.Module, inputs: Dict[str, Any], num_items_in_batch: int | None = None, *args: Any, **kwargs: Any) -> torch.Tensor:  # type: ignore[override]
        if not self._metaopt_enabled or self._meta is None:
            # Delegate to HF Trainer (supports extra args like num_items_in_batch)
            return super().training_step(model, inputs, *(() if num_items_in_batch is None else (num_items_in_batch,)), *args, **kwargs)  # type: ignore[misc]

        model.train()
        inputs = self._prepare_inputs(inputs)
        outputs = model(**inputs)
        loss = outputs.get("loss") if isinstance(outputs, dict) else outputs[0]

        # Backward using accelerator for mixed precision/DDP compatibility
        self.accelerator.backward(loss)
        # Clip then meta-step
        try:
            clip_grad_norm_(model.parameters(), 1.0)
        except Exception:
            pass
        self._meta.step(float(loss.detach().item()))
        return loss.detach()


class MetaOptTrainer(Trainer):  # type: ignore
    """
    HuggingFace Trainer subclass that applies MixtureMetaOptGrouped instead of optimizer.step.
    - Computes loss and backward via standard Trainer infra (Accelerate-aware)
    - Clips grads, then calls meta-optimizer step
    - Skips the regular optimizer step when enabled
    """

    def __init__(self, *args: Any, metaopt: bool = False, meta_base_lr: float = 3e-4, meta_base_wd: float = 0.01, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._metaopt_enabled: bool = bool(metaopt)
        self._meta: MixtureMetaOptGrouped | None = None
        if self._metaopt_enabled:
            device = "cuda" if torch.cuda.is_available() else "cpu"
            groups = make_param_groups(self.model)
            gate = GateNetGrouped(n_groups=len(groups), hidden=128)
            self._meta = MixtureMetaOptGrouped(self.model, groups, gate, base_lr=float(meta_base_lr), base_wd=float(meta_base_wd), device=device)

    def training_step(self, model: torch.nn.Module, inputs: Dict[str, Any], num_items_in_batch: int | None = None, *args: Any, **kwargs: Any) -> torch.Tensor:  # type: ignore[override]
        if not self._metaopt_enabled or self._meta is None:
            # Delegate to HF Trainer (supports extra args like num_items_in_batch)
            return super().training_step(model, inputs, *(() if num_items_in_batch is None else (num_items_in_batch,)), *args, **kwargs)  # type: ignore[misc]

        model.train()
        inputs = self._prepare_inputs(inputs)
        outputs = model(**inputs)
        loss = outputs.get("loss") if isinstance(outputs, dict) else outputs[0]

        # Backward using accelerator for mixed precision/DDP compatibility
        self.accelerator.backward(loss)
        # Clip then meta-step
        try:
            clip_grad_norm_(model.parameters(), 1.0)
        except Exception:
            pass
        self._meta.step(float(loss.detach().item()))
        return loss.detach()

    def optimizer_step(self, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
        if not self._metaopt_enabled:
            return super().optimizer_step(*args, **kwargs)  # type: ignore[misc]
        # Skip regular optimizer stepping when meta-optimizer is enabled
        return None


    def training_step(self, model: torch.nn.Module, inputs: Dict[str, Any], num_items_in_batch: int | None = None, *args: Any, **kwargs: Any) -> torch.Tensor:  # type: ignore[override]
        if not self._metaopt_enabled or self._meta is None:
            # Delegate to HF Trainer (supports extra args like num_items_in_batch)
            return super().training_step(model, inputs, *(() if num_items_in_batch is None else (num_items_in_batch,)), *args, **kwargs)  # type: ignore[misc]

        model.train()
        inputs = self._prepare_inputs(inputs)
        outputs = model(**inputs)
        loss = outputs.get("loss") if isinstance(outputs, dict) else outputs[0]

        # Backward using accelerator for mixed precision/DDP compatibility
        self.accelerator.backward(loss)
        # Clip then meta-step
        try:
            clip_grad_norm_(model.parameters(), 1.0)
        except Exception:
            pass
        self._meta.step(float(loss.detach().item()))
        return loss.detach()

    def optimizer_step(self, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
        if not self._metaopt_enabled:
            return super().optimizer_step(*args, **kwargs)  # type: ignore[misc]
        # Skip regular optimizer stepping when meta-optimizer is enabled
        return None


import copy
import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from typing import Callable, Dict, Any

def clone_model_for_unroll(model: torch.nn.Module) -> torch.nn.Module:
    shadow = copy.deepcopy(model)
    for p in shadow.parameters():
        p.requires_grad_(True)
    return shadow

@torch.no_grad()
def _adamw_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    v = st.setdefault("v", torch.zeros_like(p))
    t = st.setdefault("t", 0) + 1
    st["t"] = t
    m.mul_(0.9).add_(g, alpha=0.1)
    v.mul_(0.999).addcmul_(g, g, value=0.001)
    mhat = m / (1 - 0.9 ** t)
    vhat = v / (1 - 0.999 ** t)
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(mhat, vhat.sqrt().add_(1e-8), value=-lr)

@torch.no_grad()
def _lion_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    u = 0.9 * m + 0.1 * g
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.add_(u.sign(), alpha=-lr)
    m.mul_(0.99).add_(g, alpha=0.01)

@torch.no_grad()
def _sophia_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float, h: torch.Tensor, eps: float = 1e-8) -> None:
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(g, h.abs().add_(eps), value=-lr)

def short_unroll_reinforce(
    base_model: torch.nn.Module,
    make_groups_fn: Callable[[torch.nn.Module], Any],
    gate: torch.nn.Module,
    batcher: Callable[[str, int], tuple],
    curv_beta: float = 0.99,
    base_lr: float = 3e-4,
    base_wd: float = 0.01,
    steps: int = 200,
    bsz: int = 128,
    device: str = "cuda",
    gate_opt_lr: float = 1e-3,
) -> Dict[str, float]:
    """
    Run a short shadow unroll and update the gate via REINFORCE using -val_loss as reward.
    """
    gate.train()
    gate_opt = torch.optim.AdamW(gate.parameters(), lr=gate_opt_lr)

    # 1) Clone model and groups
    model = clone_model_for_unroll(base_model).to(device)
    groups = make_groups_fn(model)

    # 2) Per-param state + curvature store (diagonal Fisher)
    stA: Dict[int, Dict[str, torch.Tensor]] = {}
    stL: Dict[int, Dict[str, torch.Tensor]] = {}
    stS: Dict[int, Dict[str, torch.Tensor]] = {}
    Hdiag: Dict[int, torch.Tensor] = {}

    # 3) Sample a discrete arm per group (fixed during this short unroll)
    logps = []
    arms = []
    with torch.no_grad():
        for gi, _ in enumerate(groups):
            feats = torch.zeros(1, 6, device=device)  # cold start features
            logits, lr_s, wd_s = gate(feats, torch.tensor([gi], device=device))
            pi = logits.softmax(-1)
            dist = torch.distributions.Categorical(pi.squeeze(0))
            a = dist.sample()
            logp = dist.log_prob(a)
            logps.append(logp)
            arms.append((int(a.item()), float(lr_s.item()), float(wd_s.item())))

    # 4) Unroll steps with chosen arms
    for _ in range(steps):
        xb, yb = batcher("train", bsz)
        logits, loss = model(xb, yb)
        for p in model.parameters():
            if p.grad is not None:
                p.grad.zero_()
        loss.backward()
        clip_grad_norm_(model.parameters(), 1.0)

        # curvature (Fisher diag)
        with torch.no_grad():
            for p in model.parameters():
                if p.grad is None:
                    continue
                pid = id(p)
                h = Hdiag.setdefault(pid, torch.zeros_like(p))
                g = p.grad
                h.mul_(curv_beta).addcmul_(g, g, value=(1 - curv_beta))

        # apply chosen arm per group
        with torch.no_grad():
            for gi, grp in enumerate(groups):
                arm, lr_s, wd_s = arms[gi]
                lr = base_lr * max(0.25, min(2.0, lr_s))
                wd = base_wd * max(0.0, min(2.0, wd_s))
                for p in grp["params"]:
                    if p.grad is None:
                        continue
                    g = p.grad
                    pid = id(p)
                    if arm == 0:
                        _adamw_single(p, g, stA.setdefault(pid, {}), lr, wd)
                    elif arm == 1:
                        _lion_single(p, g, stL.setdefault(pid, {}), lr, wd)
                    else:
                        _sophia_single(p, g, stS.setdefault(pid, {}), lr, wd, Hdiag.get(pid, torch.ones_like(p)))

    # 5) Meta-reward from validation loss
    with torch.no_grad():
        xb, yb = batcher("val", 512)
        _, vloss = model(xb, yb)
        reward = -vloss  # maximize negative val loss

    # 6) REINFORCE update
    adv = reward.detach() - vloss.detach()  # centered within run
    loss_meta = -(adv) * torch.stack(logps).mean()
    gate_opt.zero_grad()
    loss_meta.backward()
    gate_opt.step()

    return {"reward": float(reward), "val_loss": float(vloss)}


import copy
import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from typing import Callable, Dict, Any

def clone_model_for_unroll(model: torch.nn.Module) -> torch.nn.Module:
    shadow = copy.deepcopy(model)
    for p in shadow.parameters():
        p.requires_grad_(True)
    return shadow

@torch.no_grad()
def _adamw_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    v = st.setdefault("v", torch.zeros_like(p))
    t = st.setdefault("t", 0) + 1
    st["t"] = t
    m.mul_(0.9).add_(g, alpha=0.1)
    v.mul_(0.999).addcmul_(g, g, value=0.001)
    mhat = m / (1 - 0.9 ** t)
    vhat = v / (1 - 0.999 ** t)
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(mhat, vhat.sqrt().add_(1e-8), value=-lr)

@torch.no_grad()
def _lion_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))


import copy
import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from typing import Callable, Dict, Any

def clone_model_for_unroll(model: torch.nn.Module) -> torch.nn.Module:
    shadow = copy.deepcopy(model)
    for p in shadow.parameters():
        p.requires_grad_(True)
    return shadow

@torch.no_grad()
def _adamw_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    v = st.setdefault("v", torch.zeros_like(p))
    t = st.setdefault("t", 0) + 1
    st["t"] = t
    m.mul_(0.9).add_(g, alpha=0.1)
    v.mul_(0.999).addcmul_(g, g, value=0.001)
    mhat = m / (1 - 0.9 ** t)
    vhat = v / (1 - 0.999 ** t)
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(mhat, vhat.sqrt().add_(1e-8), value=-lr)

@torch.no_grad()
def _lion_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    u = 0.9 * m + 0.1 * g
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.add_(u.sign(), alpha=-lr)
    m.mul_(0.99).add_(g, alpha=0.01)

@torch.no_grad()
def _sophia_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float, h: torch.Tensor, eps: float = 1e-8) -> None:
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(g, h.abs().add_(eps), value=-lr)

def short_unroll_reinforce(


        p.requires_grad_(True)
    return shadow

@torch.no_grad()
def _adamw_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    v = st.setdefault("v", torch.zeros_like(p))
    t = st.setdefault("t", 0) + 1
    st["t"] = t
    m.mul_(0.9).add_(g, alpha=0.1)
    v.mul_(0.999).addcmul_(g, g, value=0.001)
    mhat = m / (1 - 0.9 ** t)
    vhat = v / (1 - 0.999 ** t)
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(mhat, vhat.sqrt().add_(1e-8), value=-lr)

@torch.no_grad()
def _lion_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    u = 0.9 * m + 0.1 * g
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.add_(u.sign(), alpha=-lr)
    m.mul_(0.99).add_(g, alpha=0.01)

@torch.no_grad()
def _sophia_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float, h: torch.Tensor, eps: float = 1e-8) -> None:
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(g, h.abs().add_(eps), value=-lr)

def short_unroll_reinforce(
    base_model: torch.nn.Module,
    make_groups_fn: Callable[[torch.nn.Module], Any],
    gate: torch.nn.Module,
    batcher: Callable[[str, int], tuple],
    curv_beta: float = 0.99,
    base_lr: float = 3e-4,
    base_wd: float = 0.01,
    steps: int = 200,
    bsz: int = 128,
    device: str = "cuda",


    m.mul_(0.9).add_(g, alpha=0.1)
    v.mul_(0.999).addcmul_(g, g, value=0.001)
    mhat = m / (1 - 0.9 ** t)
    vhat = v / (1 - 0.999 ** t)
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(mhat, vhat.sqrt().add_(1e-8), value=-lr)

@torch.no_grad()
def _lion_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float) -> None:
    m = st.setdefault("m", torch.zeros_like(p))
    u = 0.9 * m + 0.1 * g
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.add_(u.sign(), alpha=-lr)
    m.mul_(0.99).add_(g, alpha=0.01)

@torch.no_grad()
def _sophia_single(p: torch.Tensor, g: torch.Tensor, st: Dict[str, torch.Tensor], lr: float, wd: float, h: torch.Tensor, eps: float = 1e-8) -> None:
    if wd:
        p.add_(p, alpha=-lr * wd)
    p.addcdiv_(g, h.abs().add_(eps), value=-lr)

def short_unroll_reinforce(
    base_model: torch.nn.Module,
    make_groups_fn: Callable[[torch.nn.Module], Any],
    gate: torch.nn.Module,
    batcher: Callable[[str, int], tuple],
    curv_beta: float = 0.99,
    base_lr: float = 3e-4,
    base_wd: float = 0.01,
    steps: int = 200,
    bsz: int = 128,
    device: str = "cuda",
    gate_opt_lr: float = 1e-3,
) -> Dict[str, float]:
    """
    Run a short shadow unroll and update the gate via REINFORCE using -val_loss as reward.
    """
    gate.train()
    gate_opt = torch.optim.AdamW(gate.parameters(), lr=gate_opt_lr)