"""Measure untrained preferences and ReLU bends on all 361 grid poses.

Run from the repository root:
python3 classes/ml/labs/robot-arm/tools/probe_initialization.py
No training is performed. Every initialization uses the same seeds 0–9.
"""
import sys
from pathlib import Path
from statistics import mean

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from q_network import make_network
from world import HOME, encode, reference


def measure(initialization, seed, poses):
    torch.manual_seed(seed)
    network = make_network(initialization)
    with torch.no_grad():
        q = network(poses)
        top_two = q.topk(2, dim=1).values
        hidden = network[0](poses)
        # A ReLU bends at zero. It must see both signs to bend on this map.
        bends = (hidden.min(dim=0).values < 0) & (hidden.max(dim=0).values > 0)
    return (q.abs().max().item(),
            (top_two[:, 0] - top_two[:, 1]).mean().item(),
            int((~bends).sum()))


if __name__ == "__main__":
    torch.set_num_threads(1)
    poses = torch.tensor(encode([[i, j] for i in range(19) for j in range(19)]))
    print("Seeds 0–9; all 361 poses; before any training.")
    print("Columns: mean per-seed max |Q|; mean best-minus-second Q over poses/seeds; mean non-bending ReLUs / 16.")
    print(f"Real full-model home return: {reference()[HOME].max():.6f}")
    for initialization in ("default", "tiny-output", "scaled-hidden", "custom"):
        rows = [measure(initialization, seed, poses) for seed in range(10)]
        averages = [mean(row[k] for row in rows) for k in range(3)]
        print(f"{initialization:14s} {averages[0]:.6f} {averages[1]:.6f} {averages[2]:.1f}")
        for seed, (maximum, gap, unbent) in enumerate(rows):
            print(f"  seed {seed}: {maximum:.6f} {gap:.6f} {unbent}")