"""Activation gate + offline B_eff diagnostic (trajectory.ai App. A) from pilot_results.json."""
import json
import numpy as np
from datasets import load_dataset
from common import *
import gs1

R = json.load(open(f"{OUT}/pilot_results.json"))
info = json.load(open(f"{OUT}/restoration.json"))
g, b = R["gs1_seed0"]["steps"], R["grpo_gs8_seed0"]["steps"]
n_eff = np.concatenate([s["n_eff"] for s in g])
absd = np.concatenate([s["abs_baseline_minus_uniform_loo"] for s in g])
vals = {
    "gs1_executed_group_size": float(max(max(s["executed_group_sizes"]) for s in g)) if all(s["executed_group_sizes"] == [1] for s in g) else float("nan"),
    "rollouts_per_step_both_arms": float(min(s["rollouts"] for s in g + b)) if len({s["rollouts"] for s in g + b}) == 1 else float("nan"),
    "gs1_effective_contributors_median": float(np.median(n_eff)),
    "gs1_mean_abs_baseline_minus_uniform_loo": float(absd.mean()),
    "batch_mean_reward_pilot": float(np.mean([s["mean_reward"] for s in g])),
    "executed_base_model": info["arms"][1]["server_info"]["model_name"],
}
ranges = {"gs1_executed_group_size": (1, 1), "rollouts_per_step_both_arms": (64, 64), "gs1_effective_contributors_median": (2, 40),
          "gs1_mean_abs_baseline_minus_uniform_loo": (0.02, None), "batch_mean_reward_pilot": (0.1, 0.95)}
checks = {k: (lo is None or vals[k] >= lo) and (hi is None or vals[k] <= hi) for k, (lo, hi) in ranges.items()}
checks["executed_base_model"] = vals["executed_base_model"] == MODEL
extra = {"grpo_batch_mean_reward_pilot": float(np.mean([s["mean_reward"] for s in b])),
         "gs1_n_eff_per_step_median": [float(np.median(s["n_eff"])) for s in g],
         "gs1_frac_correct_per_step": [s["frac_correct"] for s in g], "grpo_frac_correct_per_step": [s["frac_correct"] for s in b]}

# B_eff: GRPO pilot groups as calibration tasks (8 rollouts each)
corpus = list(load_dataset("openai/gsm8k", "main", split="train")["question"])
kern = gs1.TfidfKernel(corpus, TAU)
recs = [r for s in b for r in s["rollout_records"]]
prompts = [r["prompt"] for r in recs]
Rw = np.array([r["rewards"] for r in recs], np.float64)  # T x 8
T, G = Rw.shape
sigma2 = float(Rw.var(axis=1, ddof=1).mean())
W = kern.weights(prompts)
rng = np.random.default_rng(0)
err = {"uniform_loo": [], "gs1_tfidf": []}
noise = []
for _ in range(2000):
    pick = rng.integers(0, G, size=T)
    r = Rw[np.arange(T), pick]
    mask = np.ones_like(Rw, bool); mask[np.arange(T), pick] = False
    mu = (Rw * mask).sum(1) / (G - 1)
    err["uniform_loo"].append((((r.sum() - r) / (T - 1)) - mu) ** 2)
    err["gs1_tfidf"].append((W @ r - mu) ** 2)
mu_noise = sigma2 / (G - 1)
beff = {}
for k, e in err.items():
    mse = max(float(np.mean(e)) - mu_noise, 0.0)
    beff[k] = {"mse_corrected": mse, "B_eff": float(T * sigma2 / (sigma2 + mse)) if sigma2 > 0 else None}
out = {"values": vals, "checks": checks, "passed": all(checks.values()), "ranges": ranges, "extra": extra,
       "b_eff_diagnostic": {"n_tasks": T, "rollouts_per_task": G, "sigma2_within": sigma2, "replays": 2000, "seed": 0, **beff,
                            "note": "batch B=T calibration tasks, one rollout each; target = mean of other 7 same-task rollouts; noise sigma2/7 subtracted"}}
dump("activation_calibration.json", out)
print(json.dumps({k: out[k] for k in ("values", "checks", "passed", "extra", "b_eff_diagnostic")}, indent=1))