import json
import numpy as np
from gs1_reference import reference_baselines
from rx_gs1 import ROOT, kernel_baselines

rng = np.random.default_rng(0)
out = {}
n, cur = 256, 64
x = np.column_stack([rng.normal(4.5, 0.4, n), rng.normal(-1.5, 0.3, n)])
r = (rng.random(n) < 1 / (1 + np.exp(-3 * (x[:, 0] - 4.5)))).astype(np.float64)
for bw in (0.5, 1.0):
    b, ne = kernel_baselines(x, r, cur, bandwidth=bw)
    rb, rn = reference_baselines(x, r, cur, bandwidth=bw)
    out[f"activation_bw{bw}"] = dict(baseline_std=float(b.std()), nonconstant=bool(b.std() > 1e-6),
        max_abs_vs_reference=float(np.abs(b - rb).max()), neighbors_max_abs_vs_reference=float(np.abs(ne - rn).max()),
        neff_p10=float(np.percentile(ne, 10)))
b, ne = kernel_baselines(x, r, cur, bandwidth=1e6)
loo = (r.sum() - r[n - cur:]) / (n - 1)
out["nonactivation_bw1e6"] = dict(max_abs_vs_window_loo_mean=float(np.abs(b - loo).max()),
    neff_min=float(ne.min()), equals_batch_mean_baseline=bool(np.abs(b - loo).max() < 1e-9))
xf = x.copy(); xf[:, 0] = np.arange(n) * 100.0
b, ne = kernel_baselines(xf, r, cur, bandwidth=1e-3)
out["fallback_low_mass"] = dict(max_abs_vs_window_loo_mean=float(np.abs(b - loo).max()),
    uses_fallback=bool(np.abs(b - loo).max() < 1e-12))
ok = (all(out[f"activation_bw{bw}"]["nonconstant"] and out[f"activation_bw{bw}"]["max_abs_vs_reference"] < 1e-12 for bw in (0.5, 1.0))
      and out["nonactivation_bw1e6"]["equals_batch_mean_baseline"] and out["fallback_low_mass"]["uses_fallback"])
out["passed"] = bool(ok)
(ROOT / "offline-advantage-tests.json").write_text(json.dumps(out, indent=2))
print(json.dumps(out, indent=2))