"""feat-R probe: R's construction (object-toggle signal / temporal+spatial clutter)
measured in the UNTRAINED ENCODER's feature space instead of 8x8 pooled pixels.
Mirrors compute_R.py line-for-line, swapping pool() -> enc(). Trainless.
Tests: does the random conv encoder preserve the pixel signal/clutter ratio (R)
into the feature space the linear head learns from, and does that predict learning?
Pre-registration: PREREGISTRATION_why_gradientSNR.md.
"""
import sys, json, glob, re
sys.path.insert(0, ".")
import numpy as np, cv2, torch
from triclock.model import LogLensNet

DEV = "cuda" if torch.cuda.is_available() else "cpu"
B = np.load("simreal/frames64.npy", mmap_mode="r")
NSAMP = 80

def make_enc(seed):
    torch.manual_seed(seed)
    enc = LogLensNet().enc.to(DEV).eval()
    @torch.no_grad()
    def f(imgs):                      # imgs: list of HxWx3 float -> [len,64]
        x = torch.from_numpy(np.stack(imgs)).permute(0, 3, 1, 2).float().to(DEV)
        return enc(x).cpu().numpy()
    return f

def draw(img, pos, c, size, ct):
    img = img.copy(); x, y = pos
    img[y:y+size, x:x+size] = c*ct + img[y:y+size, x:x+size].mean(axis=(0, 1))*(1-ct)
    return img

def feat_R(mode, size, ct, enc, rng):
    sig, den = [], []
    for _ in range(NSAMP):
        idx = rng.integers(len(B))
        f0 = np.asarray(B[idx], np.float32) / 255.0
        frames = [f0]*9 if mode == "static" else [np.asarray(B[rng.integers(len(B))], np.float32)/255.0 for _ in range(9)]
        pos = (int(rng.integers(2, 60-size)), int(rng.integers(2, 60-size)))
        c = rng.uniform(0.15, 0.95, 3).astype(np.float32)
        # signal: object color-toggle delta in feature space
        a, b = enc([draw(f0, pos, c, size, ct), draw(f0, pos, 1-c, size, ct)])
        sig.append(np.sum((a - b)**2) / 4)
        # clutter: temporal (across frames) + spatial (shifted probes) feature variance
        P = enc(frames)
        probes = enc([np.roll(np.roll(f0, int(rng.integers(64)), 0), int(rng.integers(64)), 1) for _ in range(6)])
        den.append(P.var(0).sum() + probes.var(0).sum())
    return float(np.mean(sig) / (np.mean(den) + 1e-9))

def parse_cell(fn):
    b = "static" if "_static_" in fn else "moving"
    m = re.search(r"_o(\d+)_ct([\d.]+)_", fn)
    return (b, int(m.group(1)), float(m.group(2))) if m else None

if __name__ == "__main__":
    smoke = "--smoke" in sys.argv
    R = json.load(open("results_compiler_v2/R_grid.json"))
    if smoke:
        enc = make_enc(0)
        for c in [("static", 14, 1.0), ("static", 5, 0.4), ("static", 10, 0.4)]:
            print(f"{c}: featR={feat_R(*c, enc, np.random.default_rng(0)):.4f}  R={R[f'{c[0]}_o{c[1]}_ct{c[2]:g}']:.4f}", flush=True)
        sys.exit(0)
    SIZES, CTS, SEEDS = [5, 6, 8, 10, 12, 14, 16], [1.0, 0.4], [0, 1, 2]
    grid = {}
    for bg in ["static", "moving"]:
        for ct in CTS:
            for s in SIZES:
                vals = [feat_R(bg, s, ct, make_enc(sd), np.random.default_rng(100+sd)) for sd in SEEDS]
                key = f"{bg}_o{s}_ct{ct:g}"
                grid[key] = {"featR_mean": float(np.mean(vals)), "featR_seeds": vals, "R": R.get(key)}
                print(f"{key:16s} R={R.get(key):.4f}  featR={np.mean(vals):.4f}", flush=True)
    outcomes = {}
    for fn in glob.glob("results_*/*an0c0_s0.json") + glob.glob("results_*/*_an0_s0.json"):
        c = parse_cell(fn)
        if c:
            try: outcomes[f"{c[0]}_o{c[1]}_ct{c[2]:g}"] = json.load(open(fn))["heldout_mean"]
            except Exception: pass
    p3 = {sd: feat_R("static", 10, 0.4, make_enc(sd), np.random.default_rng(999)) for sd in range(6)}
    json.dump({"grid": grid, "direct_outcomes": outcomes, "p3_static10ct0.4_featR_by_seed": p3},
              open("results_compiler_v2/featR.json", "w"), indent=2)
    from scipy.stats import spearmanr
    keys = [k for k in grid if grid[k]["R"] is not None]
    rho1, pv1 = spearmanr([grid[k]["R"] for k in keys], [grid[k]["featR_mean"] for k in keys])
    ok = [k for k in keys if k in outcomes]
    rho2 = spearmanr([grid[k]["featR_mean"] for k in ok], [outcomes[k] for k in ok])[0] if len(ok) >= 4 else None
    print(f"\nP1 Spearman(R, featR) over {len(keys)} cells = {rho1:.3f} (p={pv1:.2e})")
    print(f"P2 Spearman(featR, trained_acc) over {len(ok)} cells = {rho2}")
    print(f"P3 static10@0.4 featR by seed: {p3}")
    print("controls: static14ct1 featR=%.4f (HIGH?)  static5ct0.4 featR=%.4f (low?)"
          % (grid['static_o14_ct1']['featR_mean'], grid['static_o5_ct0.4']['featR_mean']))