				obj = json.loads(txt)
				if args.mode == "cli":
					argv_list = (obj.get("args") or {}).get("argv")
					if isinstance(argv_list, list):
						labels_out.append(" ".join(str(t) for t in argv_list))
					else:
						labels_out.append("")
				else:
					url = str((obj.get("args") or {}).get("url", ""))
					selector = str((obj.get("args") or {}).get("selector", ""))
					# Normalize to reduce label noise
					url = url.strip().rstrip('/')
					selector = re.sub(r"\s+", " ", selector.strip())
					labels_out.append(f"{url} {selector}".strip())
			except Exception:
				labels_out.append("")
		labels = tokenizer(text_target=labels_out, max_length=args.max_out, truncation=True)
		model_inputs["labels"] = labels["input_ids"]
		return model_inputs

	def label_non_empty(example: Dict[str, str]) -> bool:
		txt = example.get("output", "")
		try:
			obj = json.loads(txt)
		except Exception:
			return False
		if args.mode == "cli":
			argv_list = (obj.get("args") or {}).get("argv")
			return bool(isinstance(argv_list, list) and len(argv_list) > 0)
		else:
			url = str((obj.get("args") or {}).get("url", "")).strip()
			selector = str((obj.get("args") or {}).get("selector", "")).strip()
			return bool(url and selector)

	if has_external_val:
		ds_train = ds["train"].filter(label_non_empty)
		ds_validation = ds_val["validation"].filter(label_non_empty)
		tokenized_train = ds_train.map(preprocess, batched=True, remove_columns=["input", "output"])
		tokenized_val = ds_validation.map(preprocess, batched=True, remove_columns=["input", "output"])
	else:
		split = ds["train"].train_test_split(test_size=max(1e-6, min(0.5, args.val_ratio)), seed=args.seed)
		ds_train = split["train"].filter(label_non_empty)
		ds_validation = split["test"].filter(label_non_empty)
		tokenized_train = ds_train.map(preprocess, batched=True, remove_columns=["input", "output"])
		tokenized_val = ds_validation.map(preprocess, batched=True, remove_columns=["input", "output"])

	model = AutoModelForSeq2SeqLM.from_pretrained(args.model)
	# Optional gradient checkpointing
	if args.grad_checkpoint and hasattr(model, "gradient_checkpointing_enable"):
		try:
			model.gradient_checkpointing_enable()
			# use_cache must be False with gradient checkpointing to avoid warnings
			if hasattr(model, "config") and hasattr(model.config, "use_cache"):


		trainer = EWCSeq2SeqTrainer(
			model=model,
			args=training_args,
			train_dataset=tokenized_train,
			eval_dataset=tokenized_val,
			data_collator=collator,
			callbacks=callbacks,
		)
	else:
		trainer = Seq2SeqTrainer(
			model=model,
			args=training_args,
			train_dataset=tokenized_train,
			eval_dataset=tokenized_val,
			data_collator=collator,
			callbacks=callbacks,
		)

	trainer.train()
	# In DDP: save only from primary rank to avoid contention
	def _is_primary() -> bool:
		try:
			ws = int(os.environ.get("WORLD_SIZE", "1") or 1)
			rk = int(os.environ.get("RANK", "0") or 0)
			return (ws <= 1) or (rk == 0)
		except Exception:
			return True
	if _is_primary():
		model.save_pretrained(str(out_dir))
		tokenizer.save_pretrained(str(out_dir))
	print(f"Saved model -> {out_dir}")
	# Write lightweight training metadata
	try:
		from datetime import datetime as _dt  # type: ignore
		meta = {
			"trained_at": _dt.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
			"base_model": str(args.model),
			"mode": str(args.mode),
			"epochs": int(args.epochs),
			"batch_size": int(args.bsz),
			"max_inp": int(args.max_inp),
			"max_out": int(args.max_out),
		}
		(out_dir / "metadata.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
	except Exception:
		pass
	# In DDP runs, cleanly destroy process group to avoid NCCL warnings


					if k in cur:
						try:
							ref_tensors[k] = v.detach().clone().cpu() if hasattr(v, "detach") else v
						except Exception:
							pass
			# Optional Fisher importance
			if str(args.ewc_fisher).strip():
				p = Path(str(args.ewc_fisher).strip())
				if p.exists():
					try:
						obj = json.loads(p.read_text(encoding="utf-8"))
						if isinstance(obj, dict):
							fisher_importance = {str(k): float(v) for k, v in obj.items()}
					except Exception:
						fisher_importance = None
	except Exception:
		ref_tensors = None
		fisher_importance = None

	if float(getattr(args, "ewc_lambda", 0.0) or 0.0) > 0.0 and isinstance(ref_tensors, dict) and len(ref_tensors) > 0:
		class EWCSeq2SeqTrainer(Seq2SeqTrainer):  # type: ignore
			def compute_loss(self, model, inputs, return_outputs=False):  # type: ignore
				out = model(**inputs)
				loss = out.get("loss") if isinstance(out, dict) else out[0]
				try:
					lam = float(getattr(args, "ewc_lambda", 0.0) or 0.0)
					pen = None
					for name, param in model.named_parameters():
						if not param.requires_grad:
							continue
						ref = ref_tensors.get(name) if isinstance(ref_tensors, dict) else None
						if ref is None:
							continue
						w = float(fisher_importance.get(name, 1.0)) if isinstance(fisher_importance, dict) else 1.0
						diff = (param - ref.to(param.device))
						term = (w * (diff * diff)).sum()
						pen = term if pen is None else pen + term
					if pen is not None and lam > 0.0:
						loss = loss + (lam * pen)
				except Exception:
					pass
				return (loss, out) if return_outputs else loss
		trainer = EWCSeq2SeqTrainer(
			model=model,
			args=training_args,
			train_dataset=tokenized_train,
			eval_dataset=tokenized_val,
			data_collator=collator,
			callbacks=callbacks,
		)
	else:
		trainer = Seq2SeqTrainer(
			model=model,
			args=training_args,
			train_dataset=tokenized_train,
			eval_dataset=tokenized_val,
			data_collator=collator,
			callbacks=callbacks,
		)

	trainer.train()
	# In DDP: save only from primary rank to avoid contention


						try:
							ref_tensors[k] = v.detach().clone().cpu() if hasattr(v, "detach") else v
						except Exception:
							pass
			# Optional Fisher importance
			if str(args.ewc_fisher).strip():
				p = Path(str(args.ewc_fisher).strip())
				if p.exists():
					try:
						obj = json.loads(p.read_text(encoding="utf-8"))
						if isinstance(obj, dict):
							fisher_importance = {str(k): float(v) for k, v in obj.items()}
					except Exception:
						fisher_importance = None
	except Exception:
		ref_tensors = None
		fisher_importance = None

	if float(getattr(args, "ewc_lambda", 0.0) or 0.0) > 0.0 and isinstance(ref_tensors, dict) and len(ref_tensors) > 0:
		class EWCSeq2SeqTrainer(Seq2SeqTrainer):  # type: ignore
			def compute_loss(self, model, inputs, return_outputs=False):  # type: ignore
				out = model(**inputs)
				loss = out.get("loss") if isinstance(out, dict) else out[0]
				try:
					lam = float(getattr(args, "ewc_lambda", 0.0) or 0.0)
					pen = None
					for name, param in model.named_parameters():
						if not param.requires_grad:
							continue
						ref = ref_tensors.get(name) if isinstance(ref_tensors, dict) else None
						if ref is None:
							continue
						w = float(fisher_importance.get(name, 1.0)) if isinstance(fisher_importance, dict) else 1.0
						diff = (param - ref.to(param.device))
						term = (w * (diff * diff)).sum()
						pen = term if pen is None else pen + term
					if pen is not None and lam > 0.0:
						loss = loss + (lam * pen)
				except Exception:
					pass
				return (loss, out) if return_outputs else loss
		trainer = EWCSeq2SeqTrainer(
			model=model,
			args=training_args,
			train_dataset=tokenized_train,
			eval_dataset=tokenized_val,
			data_collator=collator,
			callbacks=callbacks,
		)
	else:
		trainer = Seq2SeqTrainer(
			model=model,
			args=training_args,
			train_dataset=tokenized_train,
			eval_dataset=tokenized_val,
			data_collator=collator,
			callbacks=callbacks,
		)

	trainer.train()
	# In DDP: save only from primary rank to avoid contention


import logging
import argparse
import json
from pathlib import Path
from typing import List, Dict

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score

DEFAULT_FEATURE_KEYS = [
	"plan_len",
	"plan_ws_tokens",
	"plan_unique_tokens",
	"plan_token_entropy",
	"obs_cli",
	"obs_dom",
	"num_subgoals",
	"kw_wc",
	"kw_grep",
	"kw_head",
	"kw_tail",
	"kw_sort",
	"kw_uniq",
	"kw_cut",
	# DOM-aware features (will be zero for CLI)
	"dom_url_present",
	"dom_url_len",
	"dom_url_http",
	"dom_url_www",
	"dom_selector_present",
	"dom_selector_len",
	"dom_selector_has_id",
	"dom_selector_has_class",
	"dom_selector_has_attr",
	"dom_selector_spaces",
	"dom_selector_commas",
	"dom_selector_child_ops",
	# DOM confidence features
	"dom_selector_entropy",
	"dom_selector_malformed",
	"dom_verifier_risk",
	# Optional auxiliaries from verifier/WM
	"verifier_risk",
	"wm_success_prob",
	"wm_risk",
]

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", default=str(root / "data" / "skills" / "router_ds.jsonl"))
	ap.add_argument("--negatives", default=str(root / "data" / "skills" / "router_negatives.jsonl"), help="Path to negative examples")
	ap.add_argument("--out", default=str(root / "models" / "router"))
	ap.add_argument("--tune-threshold", action="store_true", help="Tune decision threshold on a validation split and pack it")
	ap.add_argument("--feature-keys", nargs="*", default=DEFAULT_FEATURE_KEYS, help="Override feature keys used for training")
	ap.add_argument("--include-negatives", action="store_true", help="Include negative examples in training")
	args = ap.parse_args()

	rows = load_jsonl(Path(args.data))
	if not rows:
		print("No router data found.")
		return 0

	# Optionally include negative examples
	if args.include_negatives:
		neg_path = Path(args.negatives)
		if neg_path.exists():
			neg_rows = load_jsonl(neg_path)
			rows.extend(neg_rows)
			print(f"Included {len(neg_rows)} negative examples")
		else:
			print(f"Warning: negative examples file not found: {neg_path}")

	X_list: List[List[float]] = []
	y_list: List[int] = []
	feature_keys: List[str] = list(args.feature_keys)
	for r in rows:
		f = r.get("features", {})
		X_list.append([float(f.get(k, 0.0)) for k in feature_keys])
		y_list.append(int(r.get("label", 0)))
		# Optional task label for per-task threshold tuning
		# If unavailable, use empty string
		r.setdefault("task", r.get("task", ""))

	X = np.asarray(X_list, dtype=float)
	y = np.asarray(y_list, dtype=int)

	# Handle degenerate case: single-class labels → cannot train logistic regression
	unique = np.unique(y)
	if unique.shape[0] < 2:
		# Save a stub pack with prior success probability; downstream code will fallback
		prior_p = float(np.mean(y)) if y.size > 0 else 0.5
		out_dir = Path(args.out)
		out_dir.mkdir(parents=True, exist_ok=True)
		try:
			import joblib  # type: ignore
			joblib.dump({"clf": None, "keys": feature_keys, "prior_p": prior_p}, out_dir / "router.joblib")
		except Exception:
			pass
		print(json.dumps({"n_examples": int(len(rows)), "note": "single_class", "prior_p": prior_p, "keys": feature_keys}))
		return 0

	clf = LogisticRegression(max_iter=1000, solver="liblinear")
	clf.fit(X, y)

	# Optional simple threshold tuning on a holdout split
	packed_threshold = None
	packed_per_task: Dict[str, float] = {}
	if args.tune_threshold and X.shape[0] > 10:
		# 80/20 split
		n = X.shape[0]
		split = int(0.8 * n)
		X_tr, X_va = X[:split], X[split:]
		y_tr, y_va = y[:split], y[split:]
		clf.fit(X_tr, y_tr)
		# If predict_proba not available, skip tuning
		if hasattr(clf, "predict_proba") and X_va.shape[0] > 0:
			probs = clf.predict_proba(X_va)[:, 1]
			# Sweep thresholds to maximize balanced accuracy-like utility
			ths = np.linspace(0.1, 0.9, 17)
			best_th, best_score = 0.5, -1.0
			for th in ths:
				preds = (probs >= th).astype(int)
				tp = float(((preds == 1) & (y_va == 1)).sum())
				tn = float(((preds == 0) & (y_va == 0)).sum())
				fp = float(((preds == 1) & (y_va == 0)).sum())
				fn = float(((preds == 0) & (y_va == 1)).sum())
				# Balanced accuracy
				tpr = tp / max(1.0, (tp + fn))
				tnr = tn / max(1.0, (tn + fp))
				score = 0.5 * (tpr + tnr)
				if score > best_score:
					best_score, best_th = score, float(th)
			packed_threshold = float(best_th)

			# Optional: per-task thresholds. Build task list from validation slice if available.
			try:
				tasks_va: List[str] = []
				# Recover tasks from original rows aligned to validation indices
				for i in range(split, n):
					t = rows[i].get("task", "")
					tasks_va.append(str(t) if isinstance(t, str) else "")
				# Compute per-task best threshold if enough examples
				from collections import defaultdict
				idx_by_task: Dict[str, List[int]] = defaultdict(list)
				for j, t in enumerate(tasks_va):
					idx_by_task[t].append(j)
				for t, idxs in idx_by_task.items():
					if len(idxs) < 5:
						continue
					p_t = probs[idxs]
					y_t = y_va[idxs]
					best_t, best_s = 0.5, -1.0
					for th in ths:
						preds_t = (p_t >= th).astype(int)
						tp = float(((preds_t == 1) & (y_t == 1)).sum())
						tn = float(((preds_t == 0) & (y_t == 0)).sum())
						fp = float(((preds_t == 1) & (y_t == 0)).sum())
						fn = float(((preds_t == 0) & (y_t == 1)).sum())
						tpr = tp / max(1.0, (tp + fn))
						tnr = tn / max(1.0, (tn + fp))
						sc = 0.5 * (tpr + tnr)
						if sc > best_s:
							best_s, best_t = sc, float(th)
					packed_per_task[str(t)] = float(best_t)
			except Exception:
				packed_per_task = {}

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)
	try:
		import joblib  # type: ignore
		payload = {"clf": clf, "keys": feature_keys}
		if packed_threshold is not None:
			payload["threshold"] = float(packed_threshold)
		if isinstance(packed_per_task, dict) and len(packed_per_task) > 0:
			payload["thresholds"] = dict(packed_per_task)
		joblib.dump(payload, out_dir / "router.joblib")
		# Write metadata.json
		try:
			from datetime import datetime as _dt  # type: ignore
			meta = {
				"trained_at": _dt.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
				"n_examples": int(len(rows)),
				"feature_keys": feature_keys,
				"tuned_threshold": float(packed_threshold) if packed_threshold is not None else None,
			}
			(out_dir / "metadata.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
		except Exception:
			pass
	except Exception:
		pass

	print(json.dumps({"n_examples": int(len(rows)), "keys": feature_keys, "threshold": packed_threshold, "thresholds": (packed_per_task if len(packed_per_task) > 0 else None)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


	"dom_url_www",
	"dom_selector_present",
	"dom_selector_len",
	"dom_selector_has_id",
	"dom_selector_has_class",
	"dom_selector_has_attr",
	"dom_selector_spaces",
	"dom_selector_commas",
	"dom_selector_child_ops",
	# DOM confidence features
	"dom_selector_entropy",
	"dom_selector_malformed",
	"dom_verifier_risk",
	# Optional auxiliaries from verifier/WM
	"verifier_risk",
	"wm_success_prob",
	"wm_risk",
]

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", default=str(root / "data" / "skills" / "router_ds.jsonl"))
	ap.add_argument("--negatives", default=str(root / "data" / "skills" / "router_negatives.jsonl"), help="Path to negative examples")
	ap.add_argument("--out", default=str(root / "models" / "router"))
	ap.add_argument("--tune-threshold", action="store_true", help="Tune decision threshold on a validation split and pack it")
	ap.add_argument("--feature-keys", nargs="*", default=DEFAULT_FEATURE_KEYS, help="Override feature keys used for training")
	ap.add_argument("--include-negatives", action="store_true", help="Include negative examples in training")
	args = ap.parse_args()

	rows = load_jsonl(Path(args.data))
	if not rows:
		print("No router data found.")
		return 0

	# Optionally include negative examples
	if args.include_negatives:


main

	"verifier_risk",
	"wm_success_prob",
	"wm_risk",
]

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", default=str(root / "data" / "skills" / "router_ds.jsonl"))
	ap.add_argument("--negatives", default=str(root / "data" / "skills" / "router_negatives.jsonl"), help="Path to negative examples")
	ap.add_argument("--out", default=str(root / "models" / "router"))
	ap.add_argument("--tune-threshold", action="store_true", help="Tune decision threshold on a validation split and pack it")
	ap.add_argument("--feature-keys", nargs="*", default=DEFAULT_FEATURE_KEYS, help="Override feature keys used for training")
	ap.add_argument("--include-negatives", action="store_true", help="Include negative examples in training")
	args = ap.parse_args()

	rows = load_jsonl(Path(args.data))
	if not rows:
		print("No router data found.")
		return 0

	# Optionally include negative examples
	if args.include_negatives:
		neg_path = Path(args.negatives)
		if neg_path.exists():
			neg_rows = load_jsonl(neg_path)
			rows.extend(neg_rows)
			print(f"Included {len(neg_rows)} negative examples")
		else:
			print(f"Warning: negative examples file not found: {neg_path}")

	X_list: List[List[float]] = []
	y_list: List[int] = []
	feature_keys: List[str] = list(args.feature_keys)
	for r in rows:
		f = r.get("features", {})
		X_list.append([float(f.get(k, 0.0)) for k in feature_keys])
		y_list.append(int(r.get("label", 0)))
		# Optional task label for per-task threshold tuning
		# If unavailable, use empty string
		r.setdefault("task", r.get("task", ""))

	X = np.asarray(X_list, dtype=float)
	y = np.asarray(y_list, dtype=int)

	# Handle degenerate case: single-class labels → cannot train logistic regression
	unique = np.unique(y)
	if unique.shape[0] < 2:
		# Save a stub pack with prior success probability; downstream code will fallback
		prior_p = float(np.mean(y)) if y.size > 0 else 0.5
		out_dir = Path(args.out)
		out_dir.mkdir(parents=True, exist_ok=True)
		try:
			import joblib  # type: ignore
			joblib.dump({"clf": None, "keys": feature_keys, "prior_p": prior_p}, out_dir / "router.joblib")
		except Exception:
			pass
		print(json.dumps({"n_examples": int(len(rows)), "note": "single_class", "prior_p": prior_p, "keys": feature_keys}))
		return 0

	clf = LogisticRegression(max_iter=1000, solver="liblinear")
	clf.fit(X, y)

	# Optional simple threshold tuning on a holdout split
	packed_threshold = None
	packed_per_task: Dict[str, float] = {}
	if args.tune_threshold and X.shape[0] > 10:
		# 80/20 split
		n = X.shape[0]
		split = int(0.8 * n)
		X_tr, X_va = X[:split], X[split:]
		y_tr, y_va = y[:split], y[split:]
		clf.fit(X_tr, y_tr)
		# If predict_proba not available, skip tuning
		if hasattr(clf, "predict_proba") and X_va.shape[0] > 0:
			probs = clf.predict_proba(X_va)[:, 1]
			# Sweep thresholds to maximize balanced accuracy-like utility
			ths = np.linspace(0.1, 0.9, 17)
			best_th, best_score = 0.5, -1.0
			for th in ths:
				preds = (probs >= th).astype(int)
				tp = float(((preds == 1) & (y_va == 1)).sum())
				tn = float(((preds == 0) & (y_va == 0)).sum())
				fp = float(((preds == 1) & (y_va == 0)).sum())
				fn = float(((preds == 0) & (y_va == 1)).sum())
				# Balanced accuracy
				tpr = tp / max(1.0, (tp + fn))
				tnr = tn / max(1.0, (tn + fp))
				score = 0.5 * (tpr + tnr)
				if score > best_score:
					best_score, best_th = score, float(th)
			packed_threshold = float(best_th)

			# Optional: per-task thresholds. Build task list from validation slice if available.
			try:
				tasks_va: List[str] = []
				# Recover tasks from original rows aligned to validation indices
				for i in range(split, n):
					t = rows[i].get("task", "")
					tasks_va.append(str(t) if isinstance(t, str) else "")
				# Compute per-task best threshold if enough examples
				from collections import defaultdict
				idx_by_task: Dict[str, List[int]] = defaultdict(list)
				for j, t in enumerate(tasks_va):
					idx_by_task[t].append(j)
				for t, idxs in idx_by_task.items():
					if len(idxs) < 5:
						continue
					p_t = probs[idxs]
					y_t = y_va[idxs]
					best_t, best_s = 0.5, -1.0
					for th in ths:
						preds_t = (p_t >= th).astype(int)
						tp = float(((preds_t == 1) & (y_t == 1)).sum())
						tn = float(((preds_t == 0) & (y_t == 0)).sum())
						fp = float(((preds_t == 1) & (y_t == 0)).sum())
						fn = float(((preds_t == 0) & (y_t == 1)).sum())
						tpr = tp / max(1.0, (tp + fn))
						tnr = tn / max(1.0, (tn + fp))
						sc = 0.5 * (tpr + tnr)
						if sc > best_s:
							best_s, best_t = sc, float(th)
					packed_per_task[str(t)] = float(best_t)
			except Exception:
				packed_per_task = {}

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)
	try:
		import joblib  # type: ignore
		payload = {"clf": clf, "keys": feature_keys}
		if packed_threshold is not None:
			payload["threshold"] = float(packed_threshold)
		if isinstance(packed_per_task, dict) and len(packed_per_task) > 0:
			payload["thresholds"] = dict(packed_per_task)
		joblib.dump(payload, out_dir / "router.joblib")
		# Write metadata.json
		try:
			from datetime import datetime as _dt  # type: ignore
			meta = {
				"trained_at": _dt.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
				"n_examples": int(len(rows)),
				"feature_keys": feature_keys,
				"tuned_threshold": float(packed_threshold) if packed_threshold is not None else None,
			}
			(out_dir / "metadata.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
		except Exception:
			pass
	except Exception:
		pass

	print(json.dumps({"n_examples": int(len(rows)), "keys": feature_keys, "threshold": packed_threshold, "thresholds": (packed_per_task if len(packed_per_task) > 0 else None)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Tuple

def _iter_rows(paths: List[str]) -> Tuple[int, List[Dict[str, Any]]]:
	try:
		from agi_dw.tools.ds_io import iter_rows as shared_iter  # type: ignore
		return shared_iter(paths)
	except Exception:
		rows: List[Dict[str, Any]] = []
		count = 0
		for p in (paths or []):
			dp = Path(p)
			if not dp.exists():
				continue
			with dp.open("r", encoding="utf-8") as f:
				for line in f:
					s = line.strip()
					if not s:
						continue
					try:
						row: Dict[str, Any] = json.loads(s)
					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		from agi_dw.tools.ds_io import safe_float as sf  # type: ignore
		return sf(x, default)
	except Exception:
		try:
			if x is None:
				return float(default)
			return float(x)
		except Exception:
			return float(default)

def _safe_int(x: Any, default: int = 0) -> int:
	try:
		from agi_dw.tools.ds_io import safe_int as si  # type: ignore
		return si(x, default)
	except Exception:
		try:
			if x is None:
				return int(default)
			return int(x)
		except Exception:
			return int(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_aux"))
	ap.add_argument("--log-metrics", action="store_true")
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	n_items, rows = _iter_rows(args.data or [])

	# Targeting head: label distributions for primary_path and target_symbols
	path_counts: Dict[str, int] = {}
	symbol_counts: Dict[str, int] = {}
	for r in rows:
		intent = r.get("intent") or {}
		pp = intent.get("primary_path")
		if isinstance(pp, str) and pp:
			path_counts[pp] = path_counts.get(pp, 0) + 1
		ts = intent.get("target_symbols") or []
		if isinstance(ts, list):
			for s in ts:
				try:
					name = str(s)
				except Exception:
					name = ""
				if name:
					symbol_counts[name] = symbol_counts.get(name, 0) + 1

	targeting = {
		"primary_path_vocab": sorted([{ "name": k, "count": int(v) } for k, v in path_counts.items()], key=lambda x: (-x["count"], x["name"]))[:512],
		"symbol_vocab": sorted([{ "name": k, "count": int(v) } for k, v in symbol_counts.items()], key=lambda x: (-x["count"], x["name"]))[:1024],
	}
	(out_dir / "targeting.json").write_text(json.dumps(targeting, ensure_ascii=False, indent=2), encoding="utf-8")

	# Size head: basic statistics on candidate size and files
	sizes = []
	files = []
	for r in rows:
		cand = r.get("candidate") or {}
		sizes.append(_safe_int((cand or {}).get("size"), 0))
		files.append(_safe_int((cand or {}).get("files"), 0))
	avg_size = (sum(sizes) / max(1, len(sizes))) if sizes else 0.0
	med_size = float(sorted(sizes)[len(sizes) // 2]) if sizes else 0.0
	avg_files = (sum(files) / max(1, len(files))) if files else 0.0
	med_files = float(sorted(files)[len(files) // 2]) if files else 0.0
	size_head = {
		"avg_size": float(round(avg_size, 4)),
		"median_size": float(round(med_size, 4)),
		"avg_files": float(round(avg_files, 4)),
		"median_files": float(round(med_files, 4)),
		"rule_threshold_size": float(round(med_size, 4)),
		"rule_threshold_files": float(round(med_files, 4)),
	}
	(out_dir / "size_head.json").write_text(json.dumps(size_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Risk head: summarize prior risk distribution
	risks = [_safe_float(r.get("wm_prior_risk"), 0.5) for r in rows if r.get("wm_prior_risk") is not None]
	avg_risk = (sum(risks) / max(1, len(risks))) if risks else 0.5
	med_risk = float(sorted(risks)[len(risks) // 2]) if risks else 0.5
	risk_head = {"avg_prior_risk": float(round(avg_risk, 4)), "median_prior_risk": float(round(med_risk, 4)), "coverage": int(len(risks))}
	(out_dir / "risk_head.json").write_text(json.dumps(risk_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Apply/stop head: baseline decision based on size threshold
	applied = [bool(r.get("applied_ok")) for r in rows]
	thr = med_size
	preds = []
	golds = []
	for r in rows:
		cand = r.get("candidate") or {}
		s = _safe_int(cand.get("size"), 0)
		pred = bool(s <= thr)
		preds.append(pred)
		golds.append(bool(r.get("applied_ok")))
	correct = sum(1 for p, g in zip(preds, golds) if bool(p) == bool(g))
	acc = float(correct / max(1, len(golds))) if golds else 0.0
	apply_head = {"rule": "size<=median", "threshold": float(round(thr, 4)), "accuracy": float(round(acc, 4)), "items": int(len(golds))}
	(out_dir / "apply_head.json").write_text(json.dumps(apply_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Reranker head: placeholder (single-candidate traces). Expose a simple size-then-risk preference.
	reranker = {
		"policy": "prefer smaller size; tiebreak by lower wm_prior_risk",
		"note": "Placeholder reranker; implement real scoring when multi-candidate traces are available",
	}
	(out_dir / "reranker.json").write_text(json.dumps(reranker, ensure_ascii=False, indent=2), encoding="utf-8")

	# Stop/continue head: simple placeholder using applied_ok as a proxy for continue; otherwise stop
	# Define: if recent attempts frequently fail, increase stop probability. Here, single-attempt traces → mirror applied_ok.
	stop_labels = [0 if bool(r.get("applied_ok")) else 1 for r in rows]
	stop_rate = float(sum(stop_labels) / max(1, len(stop_labels))) if stop_labels else 0.0
	stop_head = {
		"definition": "1=stop, 0=continue; proxy label uses applied_ok==False → stop",
		"stop_rate": float(round(stop_rate, 4)),
		"items": int(len(stop_labels)),
	}
	(out_dir / "stop_head.json").write_text(json.dumps(stop_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Summary artifact
	pack = {
		"kind": "coder_aux_placeholder",
		"items": int(n_items),
		"paths_vocab": int(len(path_counts)),
		"symbols_vocab": int(len(symbol_counts)),
		"avg_size": float(round(avg_size, 3)),
		"avg_files": float(round(avg_files, 3)),
		"avg_prior_risk": float(round(avg_risk, 3)),
	}
	(out_dir / "model.json").write_text(json.dumps(pack, ensure_ascii=False, indent=2), encoding="utf-8")

	if bool(getattr(args, "log-metrics", False)):
		metrics = {
			"items": int(n_items),
			"apply_rule_acc": float(round(acc, 4)),
			"risk_coverage": float(round(len(risks) / max(1, n_items), 4)),
		}
		(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

	print(json.dumps({"ok": True, "items": int(n_items), "out": str(out_dir)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List, Tuple

def _iter_rows(paths: List[str]) -> Tuple[int, List[Dict[str, Any]]]:
	try:
		from agi_dw.tools.ds_io import iter_rows as shared_iter  # type: ignore
		return shared_iter(paths)
	except Exception:
		rows: List[Dict[str, Any]] = []
		count = 0
		for p in (paths or []):
			dp = Path(p)
			if not dp.exists():
				continue
			with dp.open("r", encoding="utf-8") as f:
				for line in f:
					s = line.strip()
					if not s:
						continue
					try:
						row: Dict[str, Any] = json.loads(s)
					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		from agi_dw.tools.ds_io import safe_float as sf  # type: ignore
		return sf(x, default)
	except Exception:
		try:
			if x is None:
				return float(default)
			return float(x)
		except Exception:
			return float(default)

def _safe_int(x: Any, default: int = 0) -> int:
	try:
		from agi_dw.tools.ds_io import safe_int as si  # type: ignore
		return si(x, default)
	except Exception:


		rows: List[Dict[str, Any]] = []
		count = 0
		for p in (paths or []):
			dp = Path(p)
			if not dp.exists():
				continue
			with dp.open("r", encoding="utf-8") as f:
				for line in f:
					s = line.strip()
					if not s:
						continue
					try:
						row: Dict[str, Any] = json.loads(s)
					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		from agi_dw.tools.ds_io import safe_float as sf  # type: ignore
		return sf(x, default)
	except Exception:
		try:
			if x is None:
				return float(default)
			return float(x)
		except Exception:
			return float(default)

def _safe_int(x: Any, default: int = 0) -> int:
	try:
		from agi_dw.tools.ds_io import safe_int as si  # type: ignore
		return si(x, default)
	except Exception:
		try:
			if x is None:
				return int(default)
			return int(x)
		except Exception:
			return int(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_aux"))


					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		from agi_dw.tools.ds_io import safe_float as sf  # type: ignore
		return sf(x, default)
	except Exception:
		try:
			if x is None:
				return float(default)
			return float(x)
		except Exception:
			return float(default)

def _safe_int(x: Any, default: int = 0) -> int:
	try:
		from agi_dw.tools.ds_io import safe_int as si  # type: ignore
		return si(x, default)
	except Exception:
		try:
			if x is None:
				return int(default)
			return int(x)
		except Exception:
			return int(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_aux"))
	ap.add_argument("--log-metrics", action="store_true")
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	n_items, rows = _iter_rows(args.data or [])

	# Targeting head: label distributions for primary_path and target_symbols
	path_counts: Dict[str, int] = {}
	symbol_counts: Dict[str, int] = {}
	for r in rows:
		intent = r.get("intent") or {}


			if x is None:
				return float(default)
			return float(x)
		except Exception:
			return float(default)

def _safe_int(x: Any, default: int = 0) -> int:
	try:
		from agi_dw.tools.ds_io import safe_int as si  # type: ignore
		return si(x, default)
	except Exception:
		try:
			if x is None:
				return int(default)
			return int(x)
		except Exception:
			return int(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_aux"))
	ap.add_argument("--log-metrics", action="store_true")
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	n_items, rows = _iter_rows(args.data or [])

	# Targeting head: label distributions for primary_path and target_symbols
	path_counts: Dict[str, int] = {}
	symbol_counts: Dict[str, int] = {}
	for r in rows:
		intent = r.get("intent") or {}
		pp = intent.get("primary_path")
		if isinstance(pp, str) and pp:
			path_counts[pp] = path_counts.get(pp, 0) + 1
		ts = intent.get("target_symbols") or []
		if isinstance(ts, list):
			for s in ts:
				try:
					name = str(s)
				except Exception:
					name = ""
				if name:
					symbol_counts[name] = symbol_counts.get(name, 0) + 1

	targeting = {
		"primary_path_vocab": sorted([{ "name": k, "count": int(v) } for k, v in path_counts.items()], key=lambda x: (-x["count"], x["name"]))[:512],
		"symbol_vocab": sorted([{ "name": k, "count": int(v) } for k, v in symbol_counts.items()], key=lambda x: (-x["count"], x["name"]))[:1024],
	}
	(out_dir / "targeting.json").write_text(json.dumps(targeting, ensure_ascii=False, indent=2), encoding="utf-8")

	# Size head: basic statistics on candidate size and files
	sizes = []
	files = []
	for r in rows:
		cand = r.get("candidate") or {}
		sizes.append(_safe_int((cand or {}).get("size"), 0))
		files.append(_safe_int((cand or {}).get("files"), 0))
	avg_size = (sum(sizes) / max(1, len(sizes))) if sizes else 0.0
	med_size = float(sorted(sizes)[len(sizes) // 2]) if sizes else 0.0
	avg_files = (sum(files) / max(1, len(files))) if files else 0.0
	med_files = float(sorted(files)[len(files) // 2]) if files else 0.0
	size_head = {
		"avg_size": float(round(avg_size, 4)),
		"median_size": float(round(med_size, 4)),
		"avg_files": float(round(avg_files, 4)),
		"median_files": float(round(med_files, 4)),
		"rule_threshold_size": float(round(med_size, 4)),
		"rule_threshold_files": float(round(med_files, 4)),
	}
	(out_dir / "size_head.json").write_text(json.dumps(size_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Risk head: summarize prior risk distribution
	risks = [_safe_float(r.get("wm_prior_risk"), 0.5) for r in rows if r.get("wm_prior_risk") is not None]
	avg_risk = (sum(risks) / max(1, len(risks))) if risks else 0.5
	med_risk = float(sorted(risks)[len(risks) // 2]) if risks else 0.5
	risk_head = {"avg_prior_risk": float(round(avg_risk, 4)), "median_prior_risk": float(round(med_risk, 4)), "coverage": int(len(risks))}
	(out_dir / "risk_head.json").write_text(json.dumps(risk_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Apply/stop head: baseline decision based on size threshold
	applied = [bool(r.get("applied_ok")) for r in rows]
	thr = med_size
	preds = []
	golds = []
	for r in rows:
		cand = r.get("candidate") or {}
		s = _safe_int(cand.get("size"), 0)
		pred = bool(s <= thr)
		preds.append(pred)
		golds.append(bool(r.get("applied_ok")))
	correct = sum(1 for p, g in zip(preds, golds) if bool(p) == bool(g))
	acc = float(correct / max(1, len(golds))) if golds else 0.0
	apply_head = {"rule": "size<=median", "threshold": float(round(thr, 4)), "accuracy": float(round(acc, 4)), "items": int(len(golds))}
	(out_dir / "apply_head.json").write_text(json.dumps(apply_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Reranker head: placeholder (single-candidate traces). Expose a simple size-then-risk preference.
	reranker = {
		"policy": "prefer smaller size; tiebreak by lower wm_prior_risk",
		"note": "Placeholder reranker; implement real scoring when multi-candidate traces are available",
	}
	(out_dir / "reranker.json").write_text(json.dumps(reranker, ensure_ascii=False, indent=2), encoding="utf-8")

	# Stop/continue head: simple placeholder using applied_ok as a proxy for continue; otherwise stop
	# Define: if recent attempts frequently fail, increase stop probability. Here, single-attempt traces → mirror applied_ok.
	stop_labels = [0 if bool(r.get("applied_ok")) else 1 for r in rows]
	stop_rate = float(sum(stop_labels) / max(1, len(stop_labels))) if stop_labels else 0.0
	stop_head = {
		"definition": "1=stop, 0=continue; proxy label uses applied_ok==False → stop",
		"stop_rate": float(round(stop_rate, 4)),
		"items": int(len(stop_labels)),
	}
	(out_dir / "stop_head.json").write_text(json.dumps(stop_head, ensure_ascii=False, indent=2), encoding="utf-8")

	# Summary artifact
	pack = {
		"kind": "coder_aux_placeholder",
		"items": int(n_items),
		"paths_vocab": int(len(path_counts)),
		"symbols_vocab": int(len(symbol_counts)),
		"avg_size": float(round(avg_size, 3)),
		"avg_files": float(round(avg_files, 3)),
		"avg_prior_risk": float(round(avg_risk, 3)),
	}
	(out_dir / "model.json").write_text(json.dumps(pack, ensure_ascii=False, indent=2), encoding="utf-8")

	if bool(getattr(args, "log-metrics", False)):
		metrics = {
			"items": int(n_items),
			"apply_rule_acc": float(round(acc, 4)),
			"risk_coverage": float(round(len(risks) / max(1, n_items), 4)),
		}
		(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

	print(json.dumps({"ok": True, "items": int(n_items), "out": str(out_dir)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def _iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			s = line.strip()
			if not s:
				continue
			try:
				yield json.loads(s)
			except Exception:
				continue

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		return float(x)
	except Exception:
		return float(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--traces", nargs="*", default=[str(root / "data" / "traces" / "dev_loop.ci.jsonl"), str(root / "data" / "traces" / "dev_loop.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_rl"))
	# Reward weights: success − w_files*files − w_loc*loc − w_risk*risk
	ap.add_argument("--w-files", type=float, default=0.01)
	ap.add_argument("--w-loc", type=float, default=0.001)
	ap.add_argument("--w-risk", type=float, default=0.1)
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	total = 0
	rewards: List[float] = []
	succ = 0
	files_all: List[float] = []
	loc_all: List[float] = []
	risk_all: List[float] = []

	def _extract(rec: Dict[str, Any]) -> Dict[str, Any]:
		# Expect dev loop traces built via build_trace("devloop", ...)
		action = rec.get("action", {}) if isinstance(rec.get("action"), dict) else {}
		argsA = action.get("args", {}) if isinstance(action.get("args"), dict) else {}
		result = rec.get("result", {}) if isinstance(rec.get("result"), dict) else {}
		crit = rec.get("critique", {}) if isinstance(rec.get("critique"), dict) else {}
		ok = str(result.get("status", "")).lower() == "ok"
		size = float(argsA.get("size", 0.0) or 0.0)
		files = float(argsA.get("files", 0.0) or 0.0)
		risk = _safe_float(crit.get("risk", 0.0), 0.0)
		return {"ok": ok, "files": files, "size": size, "risk": risk}

	for p in [Path(x) for x in (args.traces or [])]:
		if not p.exists():
			continue
		for rec in _iter_jsonl(p):
			if not isinstance(rec, dict):
				continue
			# Only score attempts that have a code.patch.apply action
			action = rec.get("action")
			if not isinstance(action, dict) or str(action.get("tool", "")) != "code.patch.apply":
				continue
			info = _extract(rec)
			total += 1
			succ += 1 if info["ok"] else 0
			files_all.append(info["files"])
			loc_all.append(info["size"])  # use candidate size as LOC proxy (hunk changes)
			risk_all.append(info["risk"])
			r = (1.0 if info["ok"] else 0.0) - float(args.w_files) * info["files"] - float(args.w_loc) * info["size"] - float(args.w_risk) * info["risk"]
			rewards.append(float(r))

	avg_reward = (sum(rewards) / max(1, len(rewards))) if rewards else 0.0
	sr = (float(succ) / max(1.0, float(total))) if total else 0.0
	avg_files = (sum(files_all) / max(1, len(files_all))) if files_all else 0.0
	avg_loc = (sum(loc_all) / max(1, len(loc_all))) if loc_all else 0.0
	avg_risk = (sum(risk_all) / max(1, len(risk_all))) if risk_all else 0.0

	# Save simple metrics artifact to integrate with dashboards/CI
	metrics = {
		"items": int(total),
		"success_rate": float(round(sr, 4)),
		"avg_reward": float(round(avg_reward, 4)),
		"avg_files": float(round(avg_files, 3)),
		"avg_loc": float(round(avg_loc, 3)),
		"avg_risk": float(round(avg_risk, 3)),
		"weights": {"w_files": float(args.w_files), "w_loc": float(args.w_loc), "w_risk": float(args.w_risk)},
	}
	(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")
	# Minimal model.json to summarize run
	(out_dir / "model.json").write_text(json.dumps({"kind": "coder_rl_baseline", **metrics}, ensure_ascii=False, indent=2), encoding="utf-8")

	print(json.dumps({"ok": True, **metrics, "out": str(out_dir)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def _iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			s = line.strip()
			if not s:
				continue
			try:
				yield json.loads(s)
			except Exception:
				continue

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		return float(x)
	except Exception:
		return float(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--traces", nargs="*", default=[str(root / "data" / "traces" / "dev_loop.ci.jsonl"), str(root / "data" / "traces" / "dev_loop.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_rl"))
	# Reward weights: success − w_files*files − w_loc*loc − w_risk*risk
	ap.add_argument("--w-files", type=float, default=0.01)
	ap.add_argument("--w-loc", type=float, default=0.001)
	ap.add_argument("--w-risk", type=float, default=0.1)
	args = ap.parse_args()


import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def _iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			s = line.strip()
			if not s:
				continue
			try:
				yield json.loads(s)
			except Exception:
				continue

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		return float(x)
	except Exception:
		return float(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--traces", nargs="*", default=[str(root / "data" / "traces" / "dev_loop.ci.jsonl"), str(root / "data" / "traces" / "dev_loop.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_rl"))
	# Reward weights: success − w_files*files − w_loc*loc − w_risk*risk
	ap.add_argument("--w-files", type=float, default=0.01)
	ap.add_argument("--w-loc", type=float, default=0.001)
	ap.add_argument("--w-risk", type=float, default=0.1)
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	total = 0
	rewards: List[float] = []
	succ = 0
	files_all: List[float] = []


	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			s = line.strip()
			if not s:
				continue
			try:
				yield json.loads(s)
			except Exception:
				continue

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		return float(x)
	except Exception:
		return float(default)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--traces", nargs="*", default=[str(root / "data" / "traces" / "dev_loop.ci.jsonl"), str(root / "data" / "traces" / "dev_loop.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_rl"))
	# Reward weights: success − w_files*files − w_loc*loc − w_risk*risk
	ap.add_argument("--w-files", type=float, default=0.01)
	ap.add_argument("--w-loc", type=float, default=0.001)
	ap.add_argument("--w-risk", type=float, default=0.1)
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	total = 0
	rewards: List[float] = []
	succ = 0
	files_all: List[float] = []
	loc_all: List[float] = []
	risk_all: List[float] = []

	def _extract(rec: Dict[str, Any]) -> Dict[str, Any]:
		# Expect dev loop traces built via build_trace("devloop", ...)
		action = rec.get("action", {}) if isinstance(rec.get("action"), dict) else {}
		argsA = action.get("args", {}) if isinstance(action.get("args"), dict) else {}
		result = rec.get("result", {}) if isinstance(rec.get("result"), dict) else {}
		crit = rec.get("critique", {}) if isinstance(rec.get("critique"), dict) else {}
		ok = str(result.get("status", "")).lower() == "ok"
		size = float(argsA.get("size", 0.0) or 0.0)
		files = float(argsA.get("files", 0.0) or 0.0)
		risk = _safe_float(crit.get("risk", 0.0), 0.0)
		return {"ok": ok, "files": files, "size": size, "risk": risk}

	for p in [Path(x) for x in (args.traces or [])]:
		if not p.exists():
			continue
		for rec in _iter_jsonl(p):
			if not isinstance(rec, dict):
				continue
			# Only score attempts that have a code.patch.apply action
			action = rec.get("action")
			if not isinstance(action, dict) or str(action.get("tool", "")) != "code.patch.apply":
				continue
			info = _extract(rec)
			total += 1
			succ += 1 if info["ok"] else 0
			files_all.append(info["files"])
			loc_all.append(info["size"])  # use candidate size as LOC proxy (hunk changes)
			risk_all.append(info["risk"])
			r = (1.0 if info["ok"] else 0.0) - float(args.w_files) * info["files"] - float(args.w_loc) * info["size"] - float(args.w_risk) * info["risk"]
			rewards.append(float(r))

	avg_reward = (sum(rewards) / max(1, len(rewards))) if rewards else 0.0
	sr = (float(succ) / max(1.0, float(total))) if total else 0.0
	avg_files = (sum(files_all) / max(1, len(files_all))) if files_all else 0.0
	avg_loc = (sum(loc_all) / max(1, len(loc_all))) if loc_all else 0.0
	avg_risk = (sum(risk_all) / max(1, len(risk_all))) if risk_all else 0.0

	# Save simple metrics artifact to integrate with dashboards/CI
	metrics = {
		"items": int(total),
		"success_rate": float(round(sr, 4)),
		"avg_reward": float(round(avg_reward, 4)),
		"avg_files": float(round(avg_files, 3)),
		"avg_loc": float(round(avg_loc, 3)),
		"avg_risk": float(round(avg_risk, 3)),
		"weights": {"w_files": float(args.w_files), "w_loc": float(args.w_loc), "w_risk": float(args.w_risk)},
	}
	(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")
	# Minimal model.json to summarize run
	(out_dir / "model.json").write_text(json.dumps({"kind": "coder_rl_baseline", **metrics}, ensure_ascii=False, indent=2), encoding="utf-8")

	print(json.dumps({"ok": True, **metrics, "out": str(out_dir)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--traces", nargs="*", default=[str(root / "data" / "traces" / "dev_loop.ci.jsonl"), str(root / "data" / "traces" / "dev_loop.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_rl"))
	# Reward weights: success − w_files*files − w_loc*loc − w_risk*risk
	ap.add_argument("--w-files", type=float, default=0.01)
	ap.add_argument("--w-loc", type=float, default=0.001)
	ap.add_argument("--w-risk", type=float, default=0.1)
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	total = 0
	rewards: List[float] = []
	succ = 0
	files_all: List[float] = []
	loc_all: List[float] = []
	risk_all: List[float] = []

	def _extract(rec: Dict[str, Any]) -> Dict[str, Any]:
		# Expect dev loop traces built via build_trace("devloop", ...)
		action = rec.get("action", {}) if isinstance(rec.get("action"), dict) else {}
		argsA = action.get("args", {}) if isinstance(action.get("args"), dict) else {}
		result = rec.get("result", {}) if isinstance(rec.get("result"), dict) else {}
		crit = rec.get("critique", {}) if isinstance(rec.get("critique"), dict) else {}
		ok = str(result.get("status", "")).lower() == "ok"
		size = float(argsA.get("size", 0.0) or 0.0)
		files = float(argsA.get("files", 0.0) or 0.0)
		risk = _safe_float(crit.get("risk", 0.0), 0.0)
		return {"ok": ok, "files": files, "size": size, "risk": risk}

	for p in [Path(x) for x in (args.traces or [])]:
		if not p.exists():
			continue
		for rec in _iter_jsonl(p):
			if not isinstance(rec, dict):
				continue
			# Only score attempts that have a code.patch.apply action
			action = rec.get("action")
			if not isinstance(action, dict) or str(action.get("tool", "")) != "code.patch.apply":
				continue
			info = _extract(rec)
			total += 1
			succ += 1 if info["ok"] else 0
			files_all.append(info["files"])
			loc_all.append(info["size"])  # use candidate size as LOC proxy (hunk changes)
			risk_all.append(info["risk"])
			r = (1.0 if info["ok"] else 0.0) - float(args.w_files) * info["files"] - float(args.w_loc) * info["size"] - float(args.w_risk) * info["risk"]
			rewards.append(float(r))


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Dict, Any, List, Tuple

def _iter_rows(paths: List[str]) -> Tuple[int, List[Dict[str, Any]]]:
	# Back-compat shim; prefer shared IO
	try:
		from agi_dw.tools.ds_io import iter_rows as shared_iter  # type: ignore
		return shared_iter(paths)
	except Exception:
		rows: List[Dict[str, Any]] = []
		count = 0
		for p in (paths or []):
			dp = Path(p)
			if not dp.exists():
				continue
			with dp.open("r", encoding="utf-8") as f:
				for line in f:
					s = line.strip()
					if not s:
						continue
					try:
						row: Dict[str, Any] = json.loads(s)
					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_sft"))
	# Metrics logging
	ap.add_argument("--log-metrics", action="store_true")
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	n_items, rows = _iter_rows(args.data or [])

	# Minimal baseline trainer: compute simple label frequencies and average sizes
	applied_ok = sum(1 for r in rows if bool(r.get("applied_ok")))
	candidate_sizes = [int(r.get("candidate", {}).get("size", 0) or 0) for r in rows if isinstance(r.get("candidate"), dict)]
	candidate_files = [int(r.get("candidate", {}).get("files", 0) or 0) for r in rows if isinstance(r.get("candidate"), dict)]
	avg_size = (sum(candidate_sizes) / max(1, len(candidate_sizes))) if candidate_sizes else 0.0
	avg_files = (sum(candidate_files) / max(1, len(candidate_files))) if candidate_files else 0.0
	# Churn metrics
	churn_added = [int((r.get("churn") or {}).get("added", 0) or 0) for r in rows]
	churn_deleted = [int((r.get("churn") or {}).get("deleted", 0) or 0) for r in rows]
	avg_added = (sum(churn_added) / max(1, len(churn_added))) if churn_added else 0.0
	avg_deleted = (sum(churn_deleted) / max(1, len(churn_deleted))) if churn_deleted else 0.0
	# Presence metrics
	has_intent = sum(1 for r in rows if isinstance(r.get("intent"), dict) and (r.get("intent", {}).get("intent_summary") or r.get("intent", {}).get("primary_path")))
	has_wm = sum(1 for r in rows if r.get("wm_prior_risk") is not None)
	has_diff = sum(1 for r in rows if str(r.get("diff_text") or "").strip())

	# Save a metadata artifact
	pack = {
		"kind": "coder_sft_baseline",
		"items": int(n_items),
		"applied_ok": int(applied_ok),
		"avg_candidate_size": float(round(avg_size, 3)),
		"avg_candidate_files": float(round(avg_files, 3)),
		"avg_added": float(round(avg_added, 3)),
		"avg_deleted": float(round(avg_deleted, 3)),
		"has_intent_rows": int(has_intent),
		"has_wm_prior": int(has_wm),
		"has_diff_text": int(has_diff),
		"note": "Baseline SFT artifact; swap with full trainer when available",
	}
	(out_dir / "model.json").write_text(json.dumps(pack, ensure_ascii=False, indent=2), encoding="utf-8")

	if bool(getattr(args, "log-metrics", False)):
		metrics = {
			"items": int(n_items),
			"applied_rate": float(round(applied_ok / max(1, n_items), 4)),
			"avg_size": float(round(avg_size, 3)),
			"avg_files": float(round(avg_files, 3)),
			"avg_added": float(round(avg_added, 3)),
			"avg_deleted": float(round(avg_deleted, 3)),
			"intent_coverage": float(round(has_intent / max(1, n_items), 4)),
			"wm_prior_coverage": float(round(has_wm / max(1, n_items), 4)),
			"diff_coverage": float(round(has_diff / max(1, n_items), 4)),
		}
		(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

	print(json.dumps({"ok": True, "items": int(n_items), "out": str(out_dir)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Dict, Any, List, Tuple

def _iter_rows(paths: List[str]) -> Tuple[int, List[Dict[str, Any]]]:
	# Back-compat shim; prefer shared IO
	try:
		from agi_dw.tools.ds_io import iter_rows as shared_iter  # type: ignore
		return shared_iter(paths)
	except Exception:
		rows: List[Dict[str, Any]] = []
		count = 0
		for p in (paths or []):
			dp = Path(p)
			if not dp.exists():
				continue
			with dp.open("r", encoding="utf-8") as f:
				for line in f:
					s = line.strip()
					if not s:
						continue
					try:
						row: Dict[str, Any] = json.loads(s)
					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_sft"))
	# Metrics logging
	ap.add_argument("--log-metrics", action="store_true")
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	n_items, rows = _iter_rows(args.data or [])

	# Minimal baseline trainer: compute simple label frequencies and average sizes
	applied_ok = sum(1 for r in rows if bool(r.get("applied_ok")))
	candidate_sizes = [int(r.get("candidate", {}).get("size", 0) or 0) for r in rows if isinstance(r.get("candidate"), dict)]
	candidate_files = [int(r.get("candidate", {}).get("files", 0) or 0) for r in rows if isinstance(r.get("candidate"), dict)]


		rows: List[Dict[str, Any]] = []
		count = 0
		for p in (paths or []):
			dp = Path(p)
			if not dp.exists():
				continue
			with dp.open("r", encoding="utf-8") as f:
				for line in f:
					s = line.strip()
					if not s:
						continue
					try:
						row: Dict[str, Any] = json.loads(s)
					except Exception:
						continue
					rows.append(row)
					count += 1
		return count, rows

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", nargs="*", default=[str(root / "data" / "traces" / "coder_ds.jsonl")])
	ap.add_argument("--out", default=str(root / "models" / "coder_sft"))
	# Metrics logging
	ap.add_argument("--log-metrics", action="store_true")
	args = ap.parse_args()

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	n_items, rows = _iter_rows(args.data or [])

	# Minimal baseline trainer: compute simple label frequencies and average sizes
	applied_ok = sum(1 for r in rows if bool(r.get("applied_ok")))
	candidate_sizes = [int(r.get("candidate", {}).get("size", 0) or 0) for r in rows if isinstance(r.get("candidate"), dict)]
	candidate_files = [int(r.get("candidate", {}).get("files", 0) or 0) for r in rows if isinstance(r.get("candidate"), dict)]
	avg_size = (sum(candidate_sizes) / max(1, len(candidate_sizes))) if candidate_sizes else 0.0
	avg_files = (sum(candidate_files) / max(1, len(candidate_files))) if candidate_files else 0.0
	# Churn metrics
	churn_added = [int((r.get("churn") or {}).get("added", 0) or 0) for r in rows]
	churn_deleted = [int((r.get("churn") or {}).get("deleted", 0) or 0) for r in rows]
	avg_added = (sum(churn_added) / max(1, len(churn_added))) if churn_added else 0.0
	avg_deleted = (sum(churn_deleted) / max(1, len(churn_deleted))) if churn_deleted else 0.0
	# Presence metrics
	has_intent = sum(1 for r in rows if isinstance(r.get("intent"), dict) and (r.get("intent", {}).get("intent_summary") or r.get("intent", {}).get("primary_path")))
	has_wm = sum(1 for r in rows if r.get("wm_prior_risk") is not None)
	has_diff = sum(1 for r in rows if str(r.get("diff_text") or "").strip())

	# Save a metadata artifact
	pack = {
		"kind": "coder_sft_baseline",
		"items": int(n_items),
		"applied_ok": int(applied_ok),
		"avg_candidate_size": float(round(avg_size, 3)),
		"avg_candidate_files": float(round(avg_files, 3)),
		"avg_added": float(round(avg_added, 3)),
		"avg_deleted": float(round(avg_deleted, 3)),
		"has_intent_rows": int(has_intent),
		"has_wm_prior": int(has_wm),
		"has_diff_text": int(has_diff),
		"note": "Baseline SFT artifact; swap with full trainer when available",
	}
	(out_dir / "model.json").write_text(json.dumps(pack, ensure_ascii=False, indent=2), encoding="utf-8")

	if bool(getattr(args, "log-metrics", False)):
		metrics = {
			"items": int(n_items),
			"applied_rate": float(round(applied_ok / max(1, n_items), 4)),
			"avg_size": float(round(avg_size, 3)),
			"avg_files": float(round(avg_files, 3)),
			"avg_added": float(round(avg_added, 3)),
			"avg_deleted": float(round(avg_deleted, 3)),
			"intent_coverage": float(round(has_intent / max(1, n_items), 4)),
			"wm_prior_coverage": float(round(has_wm / max(1, n_items), 4)),
			"diff_coverage": float(round(has_diff / max(1, n_items), 4)),
		}
		(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

	print(json.dumps({"ok": True, "items": int(n_items), "out": str(out_dir)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


import logging
import argparse
from pathlib import Path

from datasets import load_dataset

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data", default=str(root / "data" / "skills" / "verifier_ds.jsonl"))
	ap.add_argument("--base-model", default="meta-llama/Llama-3.2-3B")
	ap.add_argument("--out", default=str(root / "models" / "verifier_qlora"))
	ap.add_argument("--epochs", type=int, default=1)
	ap.add_argument("--bsz", type=int, default=2)
	ap.add_argument("--max-len", type=int, default=512)
	ap.add_argument("--ewc-ref-adapter", default=None, help="Path to reference LoRA adapter to regularize towards")
	ap.add_argument("--ewc-lambda", type=float, default=0.0, help="Strength of EWC/L2 stability penalty (0 disables)")
	ap.add_argument("--ewc-fisher", default=None, help="Optional JSON mapping param_name->importance scalar for EWC")
	ap.add_argument("--metaopt", action="store_true", help="Enable meta-optimizer (MixtureMetaOptGrouped) during training")
	ap.add_argument("--meta-base-lr", type=float, default=3e-4)
	ap.add_argument("--meta-base-wd", type=float, default=0.01)
	args = ap.parse_args()

	try:
		import torch  # type: ignore
		from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments  # type: ignore
		from transformers import BitsAndBytesConfig  # type: ignore
		from peft import LoraConfig, get_peft_model  # type: ignore
		import bitsandbytes as bnb  # type: ignore
		from agi_dw.core.metaopt.hf_meta_trainer import MetaOptTrainer  # type: ignore
	except Exception as e:  # pragma: no cover
		print("[WARN] Missing deps for QLoRA (install: peft, bitsandbytes, accelerate). Skipping train.")
		print("Reason:", e)
		return 0

	tokenizer = AutoTokenizer.from_pretrained(args.base_model)
	if tokenizer.pad_token_id is None:
		tokenizer.pad_token_id = tokenizer.eos_token_id

	ds = load_dataset("json", data_files={"train": args.data})

	def preprocess(batch):
		input_ids_list = []
		attn_masks = []
		labels_list = []
		for x, y in zip(batch["input"], batch["output"]):
			prompt = (
				"You are a strict verifier. Given a trace snippet, output ONLY YAML with keys success_prob, risk, critique.\n"
				"Trace:\n" + x + "\n\nYAML:\n"
			)
			prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
			target_ids = tokenizer(y, add_special_tokens=False)["input_ids"]
			# combine and truncate
			ids = (prompt_ids + target_ids)[: args.max_len]
			# labels: ignore prompt tokens
			labels = ([-100] * min(len(prompt_ids), len(ids)) + ids[len(prompt_ids) :])
			# pad to max_len
			pad_len = args.max_len - len(ids)
			if pad_len > 0:
				ids = ids + [tokenizer.pad_token_id] * pad_len
				labels = labels + ([-100] * pad_len)
			attn = [1] * (args.max_len - pad_len) + [0] * pad_len
			input_ids_list.append(ids)
			attn_masks.append(attn)
			labels_list.append(labels)
		return {"input_ids": input_ids_list, "attention_mask": attn_masks, "labels": labels_list}

	proc = ds.map(preprocess, batched=True, remove_columns=["input", "output"])  # type: ignore

	_ = bnb  # ensure import
	quant_config = BitsAndBytesConfig(
		load_in_4bit=True,
		bnb_4bit_use_double_quant=True,
		bnb_4bit_quant_type="nf4",
		bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
	)
	model = AutoModelForCausalLM.from_pretrained(
		args.base_model,
		quantization_config=quant_config,
		device_map="auto",
	)
	lora_cfg = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], lora_dropout=0.05, bias="none")
	model = get_peft_model(model, lora_cfg)

	# Optional EWC stability guard
	ref_tensors = None
	fisher_importance = None
	if float(getattr(args, "ewc_lambda", 0.0) or 0.0) > 0.0 and args.ewc_ref_adapter:
		try:
			import os as _os
			import json as _json
			ref_dir = Path(args.ewc_ref_adapter)
			# Try common filenames
			cand = None
			for fn in ("adapter_model.bin", "pytorch_model.bin", "adapter_model.safetensors"):
				p = ref_dir / fn
				if p.exists():
					cand = p
					break
			if cand is not None:
				state = torch.load(str(cand), map_location="cuda") if cand.suffix != ".safetensors" else None
				if state is None and cand.suffix == ".safetensors":
					from safetensors.torch import load_file as _load_st  # type: ignore
					state = _load_st(str(cand))
				ref_tensors = {}
				# Intersect by exact parameter names when possible
				cur_sd = dict(model.named_parameters())
				for k, v in state.items():
					if k in cur_sd and isinstance(v, torch.Tensor):
						ref_tensors[k] = v.detach().clone().cpu()
			# Optional fisher importance
			if args.ewc_fisher and Path(args.ewc_fisher).exists():
				try:
					obj = _json.loads(Path(args.ewc_fisher).read_text(encoding="utf-8"))
					if isinstance(obj, dict):
						fisher_importance = {str(k): float(v) for k, v in obj.items()}
				except Exception:
					fisher_importance = None
		except Exception:
			ref_tensors = None

	training_args = TrainingArguments(
		output_dir=args.out,
		per_device_train_batch_size=args.bsz,
		num_train_epochs=args.epochs,
		logging_steps=10,
		save_strategy="no",
		report_to=[],
		gradient_accumulation_steps=1,
		bf16=True if torch.cuda.is_available() else False,
	)
	# Custom Trainer to add EWC penalty
	if float(getattr(args, "ewc_lambda", 0.0) or 0.0) > 0.0 and isinstance(ref_tensors, dict) and len(ref_tensors) > 0:
		class EWCTrainer(Trainer):  # type: ignore
			def compute_loss(self, model, inputs, return_outputs=False):  # type: ignore
				# Default loss
				out = model(**inputs)
				loss = out.get("loss") if isinstance(out, dict) else out[0]
				# Add EWC/L2 penalty on LoRA params present in ref
				try:
					pen = None
					lam = float(getattr(args, "ewc_lambda", 0.0) or 0.0)
					for name, param in model.named_parameters():
						if not param.requires_grad:
							continue
						ref = ref_tensors.get(name) if isinstance(ref_tensors, dict) else None
						if ref is None:
							continue
						w = float(fisher_importance.get(name, 1.0)) if isinstance(fisher_importance, dict) else 1.0
						d = (param - ref.to(param.device))
						term = (w * (d * d)).sum()
						pen = term if pen is None else pen + term
					if pen is not None and lam > 0.0:
						loss = loss + (lam * pen)
				except Exception:
					pass
				return (loss, out) if return_outputs else loss
		# If metaopt enabled, wrap EWC via composition: keep compute_loss override but step via MetaOptTrainer by delegating training_step.
		if bool(args.metaopt):
			base = EWCTrainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
			# Monkey-patch: replace optimizer_step with no-op and training_step with MetaOptTrainer.training_step behavior
			_meta = MetaOptTrainer(model=model, args=training_args, train_dataset=proc["train"], metaopt=True, meta_base_lr=float(args.meta_base_lr), meta_base_wd=float(args.meta_base_wd))  # type: ignore
			base.optimizer_step = lambda *a, **k: None  # type: ignore
			base.training_step = _meta.training_step  # type: ignore
			trainer = base  # type: ignore
		else:
			trainer = EWCTrainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
	else:
		if bool(args.metaopt):
			trainer = MetaOptTrainer(model=model, args=training_args, train_dataset=proc["train"], metaopt=True, meta_base_lr=float(args.meta_base_lr), meta_base_wd=float(args.meta_base_wd))  # type: ignore
		else:
			from transformers import Trainer  # type: ignore
			trainer = Trainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
	trainer.train()

	Path(args.out).mkdir(parents=True, exist_ok=True)
	model.save_pretrained(args.out)
	tokenizer.save_pretrained(args.out)
	print(f"Saved verifier LoRA -> {args.out}")
	# Write lightweight metadata
	try:
		from datetime import datetime as _dt  # type: ignore
		meta = {
			"trained_at": _dt.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
			"base_model": str(args.base_model),
			"epochs": int(args.epochs),
			"bsz": int(args.bsz),
			"max_len": int(args.max_len),
			"ewc_lambda": float(args.ewc_lambda),
		}
		(Path(args.out) / "metadata.json").write_text(__import__("json").dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
	except Exception:
		pass
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


preprocess

	args = ap.parse_args()

	try:
		import torch  # type: ignore
		from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments  # type: ignore
		from transformers import BitsAndBytesConfig  # type: ignore
		from peft import LoraConfig, get_peft_model  # type: ignore
		import bitsandbytes as bnb  # type: ignore
		from agi_dw.core.metaopt.hf_meta_trainer import MetaOptTrainer  # type: ignore
	except Exception as e:  # pragma: no cover
		print("[WARN] Missing deps for QLoRA (install: peft, bitsandbytes, accelerate). Skipping train.")
		print("Reason:", e)
		return 0

	tokenizer = AutoTokenizer.from_pretrained(args.base_model)
	if tokenizer.pad_token_id is None:
		tokenizer.pad_token_id = tokenizer.eos_token_id

	ds = load_dataset("json", data_files={"train": args.data})

	def preprocess(batch):
		input_ids_list = []
		attn_masks = []
		labels_list = []
		for x, y in zip(batch["input"], batch["output"]):
			prompt = (
				"You are a strict verifier. Given a trace snippet, output ONLY YAML with keys success_prob, risk, critique.\n"
				"Trace:\n" + x + "\n\nYAML:\n"
			)
			prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
			target_ids = tokenizer(y, add_special_tokens=False)["input_ids"]
			# combine and truncate
			ids = (prompt_ids + target_ids)[: args.max_len]
			# labels: ignore prompt tokens
			labels = ([-100] * min(len(prompt_ids), len(ids)) + ids[len(prompt_ids) :])
			# pad to max_len
			pad_len = args.max_len - len(ids)
			if pad_len > 0:
				ids = ids + [tokenizer.pad_token_id] * pad_len
				labels = labels + ([-100] * pad_len)
			attn = [1] * (args.max_len - pad_len) + [0] * pad_len
			input_ids_list.append(ids)
			attn_masks.append(attn)
			labels_list.append(labels)
		return {"input_ids": input_ids_list, "attention_mask": attn_masks, "labels": labels_list}

	proc = ds.map(preprocess, batched=True, remove_columns=["input", "output"])  # type: ignore

	_ = bnb  # ensure import
	quant_config = BitsAndBytesConfig(
		load_in_4bit=True,
		bnb_4bit_use_double_quant=True,
		bnb_4bit_quant_type="nf4",
		bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
	)
	model = AutoModelForCausalLM.from_pretrained(
		args.base_model,
		quantization_config=quant_config,
		device_map="auto",
	)
	lora_cfg = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], lora_dropout=0.05, bias="none")
	model = get_peft_model(model, lora_cfg)

	# Optional EWC stability guard
	ref_tensors = None


EWCTrainer

					obj = _json.loads(Path(args.ewc_fisher).read_text(encoding="utf-8"))
					if isinstance(obj, dict):
						fisher_importance = {str(k): float(v) for k, v in obj.items()}
				except Exception:
					fisher_importance = None
		except Exception:
			ref_tensors = None

	training_args = TrainingArguments(
		output_dir=args.out,
		per_device_train_batch_size=args.bsz,
		num_train_epochs=args.epochs,
		logging_steps=10,
		save_strategy="no",
		report_to=[],
		gradient_accumulation_steps=1,
		bf16=True if torch.cuda.is_available() else False,
	)
	# Custom Trainer to add EWC penalty
	if float(getattr(args, "ewc_lambda", 0.0) or 0.0) > 0.0 and isinstance(ref_tensors, dict) and len(ref_tensors) > 0:
		class EWCTrainer(Trainer):  # type: ignore
			def compute_loss(self, model, inputs, return_outputs=False):  # type: ignore
				# Default loss
				out = model(**inputs)
				loss = out.get("loss") if isinstance(out, dict) else out[0]
				# Add EWC/L2 penalty on LoRA params present in ref
				try:
					pen = None
					lam = float(getattr(args, "ewc_lambda", 0.0) or 0.0)
					for name, param in model.named_parameters():
						if not param.requires_grad:
							continue
						ref = ref_tensors.get(name) if isinstance(ref_tensors, dict) else None
						if ref is None:
							continue
						w = float(fisher_importance.get(name, 1.0)) if isinstance(fisher_importance, dict) else 1.0
						d = (param - ref.to(param.device))
						term = (w * (d * d)).sum()
						pen = term if pen is None else pen + term
					if pen is not None and lam > 0.0:
						loss = loss + (lam * pen)
				except Exception:
					pass
				return (loss, out) if return_outputs else loss
		# If metaopt enabled, wrap EWC via composition: keep compute_loss override but step via MetaOptTrainer by delegating training_step.
		if bool(args.metaopt):
			base = EWCTrainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
			# Monkey-patch: replace optimizer_step with no-op and training_step with MetaOptTrainer.training_step behavior
			_meta = MetaOptTrainer(model=model, args=training_args, train_dataset=proc["train"], metaopt=True, meta_base_lr=float(args.meta_base_lr), meta_base_wd=float(args.meta_base_wd))  # type: ignore
			base.optimizer_step = lambda *a, **k: None  # type: ignore
			base.training_step = _meta.training_step  # type: ignore
			trainer = base  # type: ignore
		else:
			trainer = EWCTrainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
	else:
		if bool(args.metaopt):
			trainer = MetaOptTrainer(model=model, args=training_args, train_dataset=proc["train"], metaopt=True, meta_base_lr=float(args.meta_base_lr), meta_base_wd=float(args.meta_base_wd))  # type: ignore
		else:
			from transformers import Trainer  # type: ignore
			trainer = Trainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
	trainer.train()

	Path(args.out).mkdir(parents=True, exist_ok=True)
	model.save_pretrained(args.out)


					if isinstance(obj, dict):
						fisher_importance = {str(k): float(v) for k, v in obj.items()}
				except Exception:
					fisher_importance = None
		except Exception:
			ref_tensors = None

	training_args = TrainingArguments(
		output_dir=args.out,
		per_device_train_batch_size=args.bsz,
		num_train_epochs=args.epochs,
		logging_steps=10,
		save_strategy="no",
		report_to=[],
		gradient_accumulation_steps=1,
		bf16=True if torch.cuda.is_available() else False,
	)
	# Custom Trainer to add EWC penalty
	if float(getattr(args, "ewc_lambda", 0.0) or 0.0) > 0.0 and isinstance(ref_tensors, dict) and len(ref_tensors) > 0:
		class EWCTrainer(Trainer):  # type: ignore
			def compute_loss(self, model, inputs, return_outputs=False):  # type: ignore
				# Default loss
				out = model(**inputs)
				loss = out.get("loss") if isinstance(out, dict) else out[0]
				# Add EWC/L2 penalty on LoRA params present in ref
				try:
					pen = None
					lam = float(getattr(args, "ewc_lambda", 0.0) or 0.0)
					for name, param in model.named_parameters():
						if not param.requires_grad:
							continue
						ref = ref_tensors.get(name) if isinstance(ref_tensors, dict) else None
						if ref is None:
							continue
						w = float(fisher_importance.get(name, 1.0)) if isinstance(fisher_importance, dict) else 1.0
						d = (param - ref.to(param.device))
						term = (w * (d * d)).sum()
						pen = term if pen is None else pen + term
					if pen is not None and lam > 0.0:
						loss = loss + (lam * pen)
				except Exception:
					pass
				return (loss, out) if return_outputs else loss
		# If metaopt enabled, wrap EWC via composition: keep compute_loss override but step via MetaOptTrainer by delegating training_step.
		if bool(args.metaopt):
			base = EWCTrainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
			# Monkey-patch: replace optimizer_step with no-op and training_step with MetaOptTrainer.training_step behavior
			_meta = MetaOptTrainer(model=model, args=training_args, train_dataset=proc["train"], metaopt=True, meta_base_lr=float(args.meta_base_lr), meta_base_wd=float(args.meta_base_wd))  # type: ignore
			base.optimizer_step = lambda *a, **k: None  # type: ignore
			base.training_step = _meta.training_step  # type: ignore
			trainer = base  # type: ignore
		else:
			trainer = EWCTrainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
	else:
		if bool(args.metaopt):
			trainer = MetaOptTrainer(model=model, args=training_args, train_dataset=proc["train"], metaopt=True, meta_base_lr=float(args.meta_base_lr), meta_base_wd=float(args.meta_base_wd))  # type: ignore
		else:
			from transformers import Trainer  # type: ignore
			trainer = Trainer(model=model, args=training_args, train_dataset=proc["train"])  # type: ignore
	trainer.train()

	Path(args.out).mkdir(parents=True, exist_ok=True)
	model.save_pretrained(args.out)


import logging
import argparse
import json
from pathlib import Path
from typing import List, Dict

import numpy as np
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import roc_auc_score

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[2]
	# Discover default traces: prefer data/traces under repo root; fall back to bench runs
	default_traces = root / "data" / "traces" / "seed_os_cli.verified.llm.jsonl"
	if not default_traces.exists():
		# Fallback: if a latest humaneval run exists with verifier sidecar, try that
		try:
			runs_dir = root / "data" / "bench" / "runs"
			cand: List[Path] = []
			if runs_dir.exists():
				for suite_dir in sorted(runs_dir.glob("humaneval/*")):
					if (suite_dir / "run.json").exists():
						cand.append(suite_dir / "run.json")
				# Pick the newest by mtime
				if cand:
					cand.sort(key=lambda p: p.stat().st_mtime, reverse=True)
					# Use the verbose sidecar if present
					try:
						meta = json.loads((cand[0]).read_text(encoding="utf-8"))
						paths = meta.get("paths", {}) if isinstance(meta, dict) else {}
						sidecar = paths.get("sidecar")
						if sidecar:
							default_traces = Path(sidecar)
					except Exception:
						pass
		except Exception:
			pass
	ap.add_argument("--traces", default=str(default_traces))
	ap.add_argument("--out", default=str(root / "models" / "verifier_calib"))
	args = ap.parse_args()

	tr_path = Path(args.traces)
	if not tr_path.exists():
		print(json.dumps({"ok": False, "error": f"traces not found: {tr_path}"}))
		return 2
	rows = load_jsonl(tr_path)
	if not rows:
		print("No verified traces found.")
		return 0

	risks: List[float] = []
	succ: List[int] = []
	for r in rows:
		status = r.get("result", {}).get("status", "")
		success = 1 if status == "ok" else 0
		risk = float(r.get("critique", {}).get("risk", 0.5))
		risks.append(risk)
		succ.append(success)

	risks_arr = np.array(risks, dtype=float)
	succ_arr = np.array(succ, dtype=int)

	base_success_prob = 1.0 - risks_arr

	base_auc = roc_auc_score(succ_arr, base_success_prob) if len(np.unique(succ_arr)) > 1 else 0.5
	base_ece = expected_calibration_error(base_success_prob, succ_arr, n_bins=10)

	y_fail = 1 - succ_arr
	iso = IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip")
	iso.fit(risks_arr, y_fail)

	calibrated_fail_prob = iso.predict(risks_arr)
	calibrated_success_prob = 1.0 - calibrated_fail_prob

	cal_auc = roc_auc_score(succ_arr, calibrated_success_prob) if len(np.unique(succ_arr)) > 1 else 0.5
	cal_ece = expected_calibration_error(calibrated_success_prob, succ_arr, n_bins=10)

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	try:
		import joblib  # type: ignore
		joblib.dump({"iso": iso}, out_dir / "calib.joblib")
	except Exception:
		pass

	metrics = {
		"base_auc": float(base_auc),
		"base_ece": float(base_ece),
		"cal_auc": float(cal_auc),
		"cal_ece": float(cal_ece),
		"n_examples": int(len(rows)),
	}
	# Persist metrics for dashboard aggregator and CI gates
	try:
		(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")
	except Exception:
		pass
	print(json.dumps(metrics))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


import logging
import argparse
import json
from pathlib import Path
from typing import List, Dict

import numpy as np
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import roc_auc_score

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[2]
	# Discover default traces: prefer data/traces under repo root; fall back to bench runs


import numpy as np
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import roc_auc_score

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[2]
	# Discover default traces: prefer data/traces under repo root; fall back to bench runs
	default_traces = root / "data" / "traces" / "seed_os_cli.verified.llm.jsonl"
	if not default_traces.exists():
		# Fallback: if a latest humaneval run exists with verifier sidecar, try that
		try:
			runs_dir = root / "data" / "bench" / "runs"
			cand: List[Path] = []
			if runs_dir.exists():
				for suite_dir in sorted(runs_dir.glob("humaneval/*")):
					if (suite_dir / "run.json").exists():
						cand.append(suite_dir / "run.json")
				# Pick the newest by mtime
				if cand:
					cand.sort(key=lambda p: p.stat().st_mtime, reverse=True)
					# Use the verbose sidecar if present


main

				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[2]
	# Discover default traces: prefer data/traces under repo root; fall back to bench runs
	default_traces = root / "data" / "traces" / "seed_os_cli.verified.llm.jsonl"
	if not default_traces.exists():
		# Fallback: if a latest humaneval run exists with verifier sidecar, try that
		try:
			runs_dir = root / "data" / "bench" / "runs"
			cand: List[Path] = []
			if runs_dir.exists():
				for suite_dir in sorted(runs_dir.glob("humaneval/*")):
					if (suite_dir / "run.json").exists():
						cand.append(suite_dir / "run.json")
				# Pick the newest by mtime
				if cand:
					cand.sort(key=lambda p: p.stat().st_mtime, reverse=True)
					# Use the verbose sidecar if present
					try:
						meta = json.loads((cand[0]).read_text(encoding="utf-8"))
						paths = meta.get("paths", {}) if isinstance(meta, dict) else {}
						sidecar = paths.get("sidecar")
						if sidecar:
							default_traces = Path(sidecar)
					except Exception:
						pass
		except Exception:
			pass
	ap.add_argument("--traces", default=str(default_traces))
	ap.add_argument("--out", default=str(root / "models" / "verifier_calib"))
	args = ap.parse_args()

	tr_path = Path(args.traces)
	if not tr_path.exists():
		print(json.dumps({"ok": False, "error": f"traces not found: {tr_path}"}))
		return 2
	rows = load_jsonl(tr_path)
	if not rows:
		print("No verified traces found.")
		return 0

	risks: List[float] = []
	succ: List[int] = []
	for r in rows:
		status = r.get("result", {}).get("status", "")
		success = 1 if status == "ok" else 0
		risk = float(r.get("critique", {}).get("risk", 0.5))
		risks.append(risk)
		succ.append(success)

	risks_arr = np.array(risks, dtype=float)
	succ_arr = np.array(succ, dtype=int)

	base_success_prob = 1.0 - risks_arr

	base_auc = roc_auc_score(succ_arr, base_success_prob) if len(np.unique(succ_arr)) > 1 else 0.5
	base_ece = expected_calibration_error(base_success_prob, succ_arr, n_bins=10)

	y_fail = 1 - succ_arr
	iso = IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip")
	iso.fit(risks_arr, y_fail)

	calibrated_fail_prob = iso.predict(risks_arr)
	calibrated_success_prob = 1.0 - calibrated_fail_prob

	cal_auc = roc_auc_score(succ_arr, calibrated_success_prob) if len(np.unique(succ_arr)) > 1 else 0.5
	cal_ece = expected_calibration_error(calibrated_success_prob, succ_arr, n_bins=10)

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)

	try:
		import joblib  # type: ignore
		joblib.dump({"iso": iso}, out_dir / "calib.joblib")
	except Exception:
		pass

	metrics = {
		"base_auc": float(base_auc),
		"base_ece": float(base_ece),
		"cal_auc": float(cal_auc),
		"cal_ece": float(cal_ece),
		"n_examples": int(len(rows)),
	}
	# Persist metrics for dashboard aggregator and CI gates
	try:
		(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")
	except Exception:
		pass
	print(json.dumps(metrics))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


import logging
import argparse
import json
from pathlib import Path
from typing import List, Dict, Tuple

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import roc_auc_score
from sklearn.isotonic import IsotonicRegression

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def brier_score(probs: np.ndarray, labels: np.ndarray) -> float:
	probs = probs.astype(float)
	labels = labels.astype(float)
	return float(np.mean((probs - labels) ** 2))

def train_val_split(n: int, val_ratio: float = 0.2, seed: int = 42) -> Tuple[np.ndarray, np.ndarray]:
	rng = np.random.RandomState(seed)
	idx = np.arange(n)
	rng.shuffle(idx)
	val_n = max(1, int(n * val_ratio)) if n > 1 else 1
	val_idx = idx[:val_n]
	train_idx = idx[val_n:]
	if train_idx.size == 0:
		train_idx = val_idx
	return train_idx, val_idx

def _enrich_text(inp: Dict[str, str], enable: bool) -> str:
	if not enable:
		return " \n ".join([inp.get("obs", ""), inp.get("plan", ""), inp.get("action", ""), inp.get("effects", "")])
	# Domain-aware tokens: extract hints from obs/plan/action JSON strings
	obs = inp.get("obs", "")
	plan = inp.get("plan", "")
	action = inp.get("action", "")
	effects = inp.get("effects", "")
	# Best-effort cheap tokens
	tokens: List[str] = []
	try:
		if "browser.read" in action or "url" in action:
			tokens.append("__DOM__")
			# URL and selector snippets
			import re  # type: ignore
			m = re.search(r"https?://[^\"'\s]+", action)
			if m:
				tokens.append("URL=" + m.group(0)[:64])
			m2 = re.search(r"selector\":\s*\"([^\"]+)\"", action)
			if m2:
				tokens.append("SEL=" + m2.group(1)[:64])
	except Exception:
		pass
	try:
		if "argv" in action or "tool" in action:
			tokens.append("__CLI__")
			# argv0
			import re  # type: ignore
			m = re.search(r"\[\s*\"([^\"]+)\"", action)
			if m:
				tokens.append("ARGV0=" + m.group(1)[:32])
	except Exception:
		pass
	return " \n ".join(tokens + [obs, plan, action, effects])

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[2]
	ap.add_argument("--data", default=str(root / "data" / "skills" / "wm_ds.jsonl"))
	ap.add_argument("--out", default=str(root / "models" / "wm_mlp"))
	ap.add_argument("--val-ratio", type=float, default=0.2)
	ap.add_argument("--domain-feats", action="store_true", help="Enrich training text with domain tokens (tool/url/selector/argv)")
	args = ap.parse_args()

	data_path = Path(args.data)
	rows = load_jsonl(data_path)
	if not rows:
		print("No data found for WM training.")
		return 0

	texts: List[str] = []
	y_success: List[int] = []
	y_risk: List[float] = []
	for r in rows:
		inp = r.get("input", {})
		# inp may also include effects
		txt = _enrich_text({
			"obs": inp.get("obs", ""),
			"plan": inp.get("plan", ""),
			"action": inp.get("action", ""),
			"effects": inp.get("effects", ""),
		}, enable=bool(args.domain_feats))
		texts.append(txt)
		y_success.append(int(r.get("success", 0)))
		y_risk.append(float(r.get("risk", 0.5)))

	X_all = texts
	y_s = np.array(y_success)
	y_r = np.array(y_risk)
	n = len(X_all)

	vec = TfidfVectorizer(max_features=5000, ngram_range=(1, 2))
	X = vec.fit_transform(X_all)

	# Split
	tr_idx, va_idx = train_val_split(n, val_ratio=float(args.val_ratio))
	X_tr, X_va = X[tr_idx], X[va_idx]
	y_s_tr, y_s_va = y_s[tr_idx], y_s[va_idx]
	y_r_tr, y_r_va = y_r[tr_idx], y_r[va_idx]

	# Classifier: success prob with single-class fallback
	clf = LogisticRegression(max_iter=1000, solver="liblinear")
	try:
		# mypy: sklearn types ignored
		clf.fit(X_tr, y_s_tr)
		s_probs_tr = clf.predict_proba(X_tr)[:, 1]
		s_probs_va = clf.predict_proba(X_va)[:, 1]
	except Exception:
		# If only a single class present, fall back to a constant-prob model
		p_const = float(np.mean(y_s_tr)) if y_s_tr.size > 0 else 0.5
		s_probs_tr = np.full(shape=(X_tr.shape[0],), fill_value=p_const, dtype=float)
		s_probs_va = np.full(shape=(X_va.shape[0],), fill_value=p_const, dtype=float)
	auc_va = roc_auc_score(y_s_va, s_probs_va) if len(np.unique(y_s_va)) > 1 else 0.5
	ece_va = expected_calibration_error(s_probs_va, y_s_va, n_bins=10)
	brier_va = brier_score(s_probs_va, y_s_va)

	# Optional isotonic calibration on validation
	cal: IsotonicRegression | None = None
	try:
		if len(np.unique(y_s_va)) > 1:
			cal = IsotonicRegression(out_of_bounds="clip")
			cal.fit(s_probs_va, y_s_va)
			# Evaluate post-calibration ECE
			cal_probs = cal.predict(s_probs_va)
			ece_va = expected_calibration_error(cal_probs, y_s_va, n_bins=10)
			brier_va = brier_score(cal_probs, y_s_va)
		else:
			cal = None
	except Exception:
		cal = None

	# Regressor: risk
	reg = Ridge(alpha=1.0)
	reg.fit(X_tr, y_r_tr)
	r_pred_va = reg.predict(X_va)
	r_mse_va = float(np.mean((r_pred_va - y_r_va) ** 2))
	# Simple risk uncertainty: std of residuals on validation set
	risk_std = float(np.std(r_pred_va - y_r_va)) if r_pred_va.size > 1 else 0.0

	out_dir = Path(args.out)
	out_dir.mkdir(parents=True, exist_ok=True)
	# Save pack including optional calibration and risk_std
	try:
		import joblib  # type: ignore
		joblib.dump({"vec": vec, "clf": clf, "reg": reg, "cal": cal, "risk_std": risk_std}, out_dir / "wm_mlp.joblib")
		# Also maintain latest/best copies and registry
		import shutil
		from datetime import datetime
		(latest_path := out_dir / "latest.joblib").write_bytes((out_dir / "wm_mlp.joblib").read_bytes())
		# Load existing best metrics if any
		metrics_path = out_dir / "metrics.json"
		registry_path = out_dir / "registry.json"
		prev_best: Dict[str, float] = {}
		try:
			if registry_path.exists():
				reg_doc = json.loads(registry_path.read_text(encoding="utf-8"))
				prev_best = reg_doc.get("best", {}) if isinstance(reg_doc, dict) else {}
		except Exception:
			prev_best = {}
		# Save metrics next (computed below)
	except Exception:
		pass

	# Save metrics
	metrics = {
		"wm_auc": float(auc_va),
		"wm_ece": float(ece_va),
		"wm_brier": float(brier_va),
		"wm_risk_mse": float(r_mse_va),
		"n_examples": int(n),
		"val_ratio": float(args.val_ratio),
		"domain_feats": bool(args.domain_feats),
	}
	(out_dir / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

	# Update registry and best model
	try:
		reg_doc = {"history": [], "best": {}}
		if (out_dir / "registry.json").exists():
			try:
				reg_doc = json.loads((out_dir / "registry.json").read_text(encoding="utf-8"))
			except Exception:
				reg_doc = {"history": [], "best": {}}
		reg_doc.setdefault("history", []).append({"ts": __import__("datetime").datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"), "metrics": metrics})
		best = reg_doc.get("best", {})
		def is_better(new: Dict[str, float], old: Dict[str, float]) -> bool:
			try:
				new_auc = float(new.get("wm_auc", 0.0))
				old_auc = float(old.get("wm_auc", 0.0))
				if new_auc != old_auc:
					return new_auc > old_auc
				# tie-breaker: lower ECE
				return float(new.get("wm_ece", 1e9)) < float(old.get("wm_ece", 1e9))
			except Exception:
				return True
		if (not best) or is_better(metrics, best):
			# Promote to best
			try:
				(out_dir / "best.joblib").write_bytes((out_dir / "wm_mlp.joblib").read_bytes())
			except Exception:
				pass
			reg_doc["best"] = metrics
		(out_dir / "registry.json").write_text(json.dumps(reg_doc, ensure_ascii=False, indent=2), encoding="utf-8")
	except Exception:
		pass

	# Write lightweight training metadata
	try:
		from datetime import datetime as _dt  # type: ignore
		meta = {
			"trained_at": _dt.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
			"val_ratio": float(args.val_ratio),
			"domain_feats": bool(args.domain_feats),
		}
		(out_dir / "metadata.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
	except Exception:
		pass

	print(json.dumps(metrics))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


import logging
import argparse
import json
from pathlib import Path
from typing import List, Dict, Tuple

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import roc_auc_score
from sklearn.isotonic import IsotonicRegression

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def brier_score(probs: np.ndarray, labels: np.ndarray) -> float:
	probs = probs.astype(float)
	labels = labels.astype(float)
	return float(np.mean((probs - labels) ** 2))


from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import roc_auc_score
from sklearn.isotonic import IsotonicRegression

def load_jsonl(path: Path) -> List[Dict]:
	rows: List[Dict] = []
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def brier_score(probs: np.ndarray, labels: np.ndarray) -> float:
	probs = probs.astype(float)
	labels = labels.astype(float)
	return float(np.mean((probs - labels) ** 2))

def train_val_split(n: int, val_ratio: float = 0.2, seed: int = 42) -> Tuple[np.ndarray, np.ndarray]:
	rng = np.random.RandomState(seed)
	idx = np.arange(n)
	rng.shuffle(idx)
	val_n = max(1, int(n * val_ratio)) if n > 1 else 1
	val_idx = idx[:val_n]
	train_idx = idx[val_n:]
	if train_idx.size == 0:
		train_idx = val_idx
	return train_idx, val_idx


				rows.append(json.loads(line))
			except Exception:
				continue
	return rows

def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
	bins = np.linspace(0.0, 1.0, n_bins + 1)
	ece = 0.0
	for i in range(n_bins):
		lo, hi = bins[i], bins[i + 1]
		mask = (probs >= lo) & (probs < hi)
		if not np.any(mask):
			continue
		conf = probs[mask].mean()
		acc = labels[mask].mean()
		ece += (mask.mean()) * abs(acc - conf)
	return float(ece)

def brier_score(probs: np.ndarray, labels: np.ndarray) -> float:
	probs = probs.astype(float)
	labels = labels.astype(float)
	return float(np.mean((probs - labels) ** 2))

def train_val_split(n: int, val_ratio: float = 0.2, seed: int = 42) -> Tuple[np.ndarray, np.ndarray]:
	rng = np.random.RandomState(seed)
	idx = np.arange(n)
	rng.shuffle(idx)
	val_n = max(1, int(n * val_ratio)) if n > 1 else 1
	val_idx = idx[:val_n]
	train_idx = idx[val_n:]
	if train_idx.size == 0:
		train_idx = val_idx
	return train_idx, val_idx

def _enrich_text(inp: Dict[str, str], enable: bool) -> str:
	if not enable:
		return " \n ".join([inp.get("obs", ""), inp.get("plan", ""), inp.get("action", ""), inp.get("effects", "")])
	# Domain-aware tokens: extract hints from obs/plan/action JSON strings
	obs = inp.get("obs", "")
	plan = inp.get("plan", "")