v16_johanna_omega/diagnostic.py

"""
Universal SVAE Diagnostic Battery
===================================
One script. Any checkpoint. Every dataset.

Usage:
  python universal_diagnostic.py                          # local best.pt
  python universal_diagnostic.py --hf v13_imagenet256     # HF version
  python universal_diagnostic.py --checkpoint /path/to.pt # explicit path

Tests across:
  - CIFAR-10 (32×32, resized to model native)
  - MNIST (28×28, resized, grayscale→RGB)
  - TinyImageNet (64×64)
  - ImageNet-128 (128×128)
  - ImageNet-256 (256×256)
  - 16 noise types (native resolution)
  - Text bytes (5 sentences)
  - Piecemeal tiling (4× resolution)
  - Geometric fingerprint per dataset
  - Spectrum analysis
  - Alpha profile
  - Compression metrics
"""

import os, sys, json, math, time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as T
from collections import defaultdict

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

# ── Model ────────────────────────────────────────────────────────

class BoundarySmooth(nn.Module):
    def __init__(self, channels=3, mid=16):
        super().__init__()
        self.net = nn.Sequential(nn.Conv2d(channels, mid, 3, padding=1), nn.GELU(),
                                  nn.Conv2d(mid, channels, 3, padding=1))
        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
    def forward(self, x): return x + self.net(x)

class SpectralCrossAttention(nn.Module):
    def __init__(self, D, n_heads=4, max_alpha=0.2, alpha_init=-2.0):
        super().__init__()
        self.n_heads = n_heads; self.head_dim = D // n_heads
        self.max_alpha = max_alpha
        self.qkv = nn.Linear(D, 3*D); self.out_proj = nn.Linear(D, D)
        self.norm = nn.LayerNorm(D); self.scale = self.head_dim**-0.5
        self.alpha_logits = nn.Parameter(torch.full((D,), alpha_init))
    @property
    def alpha(self): return self.max_alpha * torch.sigmoid(self.alpha_logits)
    def forward(self, S):
        B, N, D = S.shape; S_n = self.norm(S)
        qkv = self.qkv(S_n).reshape(B,N,3,self.n_heads,self.head_dim).permute(2,0,3,1,4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        out = (((q @ k.transpose(-2,-1))*self.scale).softmax(-1) @ v).transpose(1,2).reshape(B,N,D)
        return S * (1.0 + self.alpha.unsqueeze(0).unsqueeze(0) * torch.tanh(self.out_proj(out)))

class PatchSVAE(nn.Module):
    def __init__(self, V=256, D=16, ps=16, hidden=768, depth=4, n_cross=2):
        super().__init__()
        self.matrix_v, self.D, self.patch_size = V, D, ps
        self.patch_dim = 3*ps*ps; self.mat_dim = V*D
        self.enc_in = nn.Linear(self.patch_dim, hidden)
        self.enc_blocks = nn.ModuleList([nn.Sequential(
            nn.LayerNorm(hidden), nn.Linear(hidden, hidden),
            nn.GELU(), nn.Linear(hidden, hidden)) for _ in range(depth)])
        self.enc_out = nn.Linear(hidden, self.mat_dim)
        self.dec_in = nn.Linear(self.mat_dim, hidden)
        self.dec_blocks = nn.ModuleList([nn.Sequential(
            nn.LayerNorm(hidden), nn.Linear(hidden, hidden),
            nn.GELU(), nn.Linear(hidden, hidden)) for _ in range(depth)])
        self.dec_out = nn.Linear(hidden, self.patch_dim)
        nn.init.orthogonal_(self.enc_out.weight)
        self.cross_attn = nn.ModuleList([
            SpectralCrossAttention(D, n_heads=min(4,D)) for _ in range(n_cross)])
        self.boundary_smooth = BoundarySmooth(channels=3, mid=16)

    def _svd(self, A):
        orig = A.dtype
        with torch.amp.autocast('cuda', enabled=False):
            A_d = A.double()
            G = torch.bmm(A_d.transpose(1,2), A_d)
            G.diagonal(dim1=-2, dim2=-1).add_(1e-12)
            eig, V = torch.linalg.eigh(G)
            eig = eig.flip(-1); V = V.flip(-1)
            S = torch.sqrt(eig.clamp(min=1e-24))
            U = torch.bmm(A_d, V) / S.unsqueeze(1).clamp(min=1e-16)
            Vh = V.transpose(-2,-1).contiguous()
        return U.to(orig), S.to(orig), Vh.to(orig)

    def encode_patches(self, patches):
        B, N, _ = patches.shape
        h = F.gelu(self.enc_in(patches.reshape(B*N,-1)))
        for block in self.enc_blocks: h = h + block(h)
        M = F.normalize(self.enc_out(h).reshape(B*N, self.matrix_v, self.D), dim=-1)
        U, S, Vt = self._svd(M)
        U = U.reshape(B,N,self.matrix_v,self.D); S = S.reshape(B,N,self.D)
        Vt = Vt.reshape(B,N,self.D,self.D); M = M.reshape(B,N,self.matrix_v,self.D)
        S_c = S
        for layer in self.cross_attn: S_c = layer(S_c)
        return {'U':U, 'S_orig':S, 'S':S_c, 'Vt':Vt, 'M':M}

    def decode_patches(self, U, S, Vt):
        B, N, V, D = U.shape
        M_hat = torch.bmm(U.reshape(B*N,V,D)*S.reshape(B*N,D).unsqueeze(1), Vt.reshape(B*N,D,D))
        h = F.gelu(self.dec_in(M_hat.reshape(B*N,-1)))
        for block in self.dec_blocks: h = h + block(h)
        return self.dec_out(h).reshape(B, N, -1)

    def forward(self, images):
        B, C, H, W = images.shape
        ps = self.patch_size
        gh, gw = H//ps, W//ps
        p = images.reshape(B,C,gh,ps,gw,ps).permute(0,2,4,1,3,5).reshape(B,gh*gw,C*ps*ps)
        svd = self.encode_patches(p)
        dec = self.decode_patches(svd['U'], svd['S'], svd['Vt'])
        dec = dec.reshape(B,gh,gw,3,ps,ps).permute(0,3,1,4,2,5).reshape(B,3,gh*ps,gw*ps)
        return {'recon': self.boundary_smooth(dec), 'svd': svd}

    @staticmethod
    def effective_rank(S):
        p = S / (S.sum(-1, keepdim=True)+1e-8); p = p.clamp(min=1e-8)
        return (-(p * p.log()).sum(-1)).exp()


def load_model(hf_version=None, checkpoint_path=None):
    from huggingface_hub import hf_hub_download
    if checkpoint_path and os.path.exists(checkpoint_path):
        path = checkpoint_path
    elif hf_version:
        path = hf_hub_download(repo_id='AbstractPhil/geolip-SVAE',
                                filename=f'{hf_version}/checkpoints/best.pt', repo_type='model')
    else:
        path = '/content/checkpoints/best.pt'
    print(f"  Loading: {path}")
    ckpt = torch.load(path, map_location='cpu', weights_only=False)
    cfg = ckpt['config']
    print(f"  Epoch: {ckpt.get('epoch')}, MSE: {ckpt.get('test_mse','?'):.6f}")
    print(f"  Config: {cfg}")
    model = PatchSVAE(V=cfg['V'], D=cfg['D'], ps=cfg['patch_size'],
                       hidden=cfg['hidden'], depth=cfg['depth'],
                       n_cross=cfg['n_cross_layers'])
    model.load_state_dict(ckpt['model_state_dict'], strict=True)
    model = model.to(DEVICE).eval()
    print(f"  Params: {sum(p.numel() for p in model.parameters()):,}")
    return model, cfg


# ── Noise Generators ─────────────────────────────────────────────

NOISE_NAMES = {
    0:'gaussian', 1:'uniform', 2:'uniform_scaled', 3:'poisson',
    4:'pink', 5:'brown', 6:'salt_pepper', 7:'sparse',
    8:'block', 9:'gradient', 10:'checkerboard', 11:'mixed',
    12:'structural', 13:'cauchy', 14:'exponential', 15:'laplace',
}

def _pink(shape):
    w = torch.randn(shape); S = torch.fft.rfft2(w)
    h, ww = shape[-2], shape[-1]
    fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, ww//2+1)
    fx = torch.fft.rfftfreq(ww).unsqueeze(0).expand(h, -1)
    return torch.fft.irfft2(S / torch.sqrt(fx**2 + fy**2).clamp(min=1e-8), s=(h, ww))

def _brown(shape):
    w = torch.randn(shape); S = torch.fft.rfft2(w)
    h, ww = shape[-2], shape[-1]
    fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, ww//2+1)
    fx = torch.fft.rfftfreq(ww).unsqueeze(0).expand(h, -1)
    return torch.fft.irfft2(S / (fx**2 + fy**2).clamp(min=1e-8), s=(h, ww))

def generate_noise(noise_type, n, s):
    rng = np.random.RandomState(42)
    imgs = []
    for _ in range(n):
        if noise_type == 0: img = torch.randn(3,s,s)
        elif noise_type == 1: img = torch.rand(3,s,s)*2-1
        elif noise_type == 2: img = (torch.rand(3,s,s)-0.5)*4
        elif noise_type == 3:
            lam = rng.uniform(0.5,20.0)
            img = torch.poisson(torch.full((3,s,s),lam))/lam-1.0
        elif noise_type == 4: img = _pink((3,s,s)); img = img/(img.std()+1e-8)
        elif noise_type == 5: img = _brown((3,s,s)); img = img/(img.std()+1e-8)
        elif noise_type == 6:
            img = torch.where(torch.rand(3,s,s)>0.5, torch.ones(3,s,s)*2, -torch.ones(3,s,s)*2)
            img = img + torch.randn(3,s,s)*0.1
        elif noise_type == 7: img = torch.randn(3,s,s)*(torch.rand(3,s,s)>0.9).float()*3
        elif noise_type == 8:
            b = rng.randint(2,16); sm = torch.randn(3,s//b+1,s//b+1)
            img = F.interpolate(sm.unsqueeze(0), size=s, mode='nearest').squeeze(0)
        elif noise_type == 9:
            gy = torch.linspace(-2,2,s).unsqueeze(1).expand(s,s)
            gx = torch.linspace(-2,2,s).unsqueeze(0).expand(s,s)
            a = rng.uniform(0,2*math.pi)
            img = (math.cos(a)*gx+math.sin(a)*gy).unsqueeze(0).expand(3,-1,-1)+torch.randn(3,s,s)*0.5
        elif noise_type == 10:
            cs = rng.randint(2,16); cy = torch.arange(s)//cs; cx = torch.arange(s)//cs
            img = ((cy.unsqueeze(1)+cx.unsqueeze(0))%2).float().unsqueeze(0).expand(3,-1,-1)*2-1+torch.randn(3,s,s)*0.3
        elif noise_type == 11:
            alpha = rng.uniform(0.2,0.8)
            img = alpha*torch.randn(3,s,s)+(1-alpha)*(torch.rand(3,s,s)*2-1)
        elif noise_type == 12:
            img = torch.zeros(3,s,s); h2 = s//2
            img[:,:h2,:h2] = torch.randn(3,h2,h2)
            img[:,:h2,h2:] = torch.rand(3,h2,h2)*2-1
            img[:,h2:,:h2] = _pink((3,h2,h2))/2
            img[:,h2:,h2:] = torch.where(torch.rand(3,h2,h2)>0.5, torch.ones(3,h2,h2), -torch.ones(3,h2,h2))
        elif noise_type == 13: img = torch.tan(math.pi*(torch.rand(3,s,s)-0.5)).clamp(-3,3)
        elif noise_type == 14: img = torch.empty(3,s,s).exponential_(1.0)-1.0
        elif noise_type == 15:
            u = torch.rand(3,s,s)-0.5; img = -torch.sign(u)*torch.log1p(-2*u.abs())
        else: img = torch.randn(3,s,s)
        imgs.append(img.clamp(-4,4))
    return torch.stack(imgs)


# ── Batched Forward ──────────────────────────────────────────────

def batched_forward(model, images, max_batch=16):
    """Forward pass in chunks to avoid OOM."""
    all_recon = []; all_S = []; all_S_orig = []; all_M = []
    model.eval()
    with torch.no_grad():
        for i in range(0, len(images), max_batch):
            batch = images[i:i+max_batch].to(DEVICE)
            out = model(batch)
            all_recon.append(out['recon'].cpu())
            all_S.append(out['svd']['S'].cpu())
            all_S_orig.append(out['svd']['S_orig'].cpu())
            all_M.append(out['svd']['M'].cpu())
    return {
        'recon': torch.cat(all_recon),
        'S': torch.cat(all_S),
        'S_orig': torch.cat(all_S_orig),
        'M': torch.cat(all_M),
    }


# ── Dataset Loaders ──────────────────────────────────────────────

def load_dataset_batch(name, s, n=100):
    """Load n images from a dataset, resized to s×s, normalized.
    Returns: (tensor [N,3,H,W], mean, std, ds_name)
    """
    from datasets import load_dataset as hf_load

    if name == 'cifar10':
        transform = T.Compose([T.Resize(s), T.ToTensor(),
                                T.Normalize((0.4914,0.4822,0.4465),(0.2470,0.2435,0.2616))])
        ds = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
        imgs = [ds[i][0] for i in range(min(n, len(ds)))]
        return torch.stack(imgs), (0.4914,0.4822,0.4465), (0.2470,0.2435,0.2616), f'CIFAR-10→{s}'

    elif name == 'mnist':
        transform = T.Compose([T.Resize(s), T.Grayscale(3), T.ToTensor(),
                                T.Normalize((0.1307,0.1307,0.1307),(0.3081,0.3081,0.3081))])
        ds = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
        imgs = [ds[i][0] for i in range(min(n, len(ds)))]
        return torch.stack(imgs), (0.1307,0.1307,0.1307), (0.3081,0.3081,0.3081), f'MNIST→{s}'

    elif name == 'tiny_imagenet':
        ds = hf_load('zh-plus/tiny-imagenet', split='valid', streaming=True)
        transform = T.Compose([T.Resize(s), T.CenterCrop(s), T.ToTensor(),
                                T.Normalize((0.4802,0.4481,0.3975),(0.2770,0.2691,0.2821))])
        imgs = []
        for i, sample in enumerate(ds):
            imgs.append(transform(sample['image'].convert('RGB')))
            if i >= n-1: break
        return torch.stack(imgs), (0.4802,0.4481,0.3975), (0.2770,0.2691,0.2821), f'TinyImageNet→{s}'

    elif name == 'imagenet128':
        ds = hf_load('benjamin-paine/imagenet-1k-128x128', split='validation', streaming=True)
        transform = T.Compose([T.Resize(s), T.CenterCrop(s), T.ToTensor(),
                                T.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))])
        imgs = []
        for i, sample in enumerate(ds):
            imgs.append(transform(sample['image'].convert('RGB')))
            if i >= n-1: break
        return torch.stack(imgs), (0.485,0.456,0.406), (0.229,0.224,0.225), f'ImageNet-128→{s}'

    elif name == 'imagenet256':
        ds = hf_load('benjamin-paine/imagenet-1k-256x256', split='validation', streaming=True)
        transform = T.Compose([T.Resize(s), T.CenterCrop(s), T.ToTensor(),
                                T.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))])
        imgs = []
        for i, sample in enumerate(ds):
            imgs.append(transform(sample['image'].convert('RGB')))
            if i >= n-1: break
        return torch.stack(imgs), (0.485,0.456,0.406), (0.229,0.224,0.225), f'ImageNet-256→{s}'

    raise ValueError(f"Unknown dataset: {name}")


# ════════════════════════════════════════════════════════════════
# TESTS
# ════════════════════════════════════════════════════════════════

IMAGE_DATASETS = ['cifar10', 'mnist', 'tiny_imagenet', 'imagenet128', 'imagenet256']


def test_image_datasets(model, cfg, n=100):
    """Reconstruction MSE + geometry across all image datasets."""
    s = cfg['img_size']
    D = cfg['D']
    bs = max(4, 64 // max(1, (s // 64) ** 2))
    print(f"\n{'='*80}")
    print(f"IMAGE DATASET BATTERY ({s}×{s}, n={n})")
    print(f"{'='*80}")
    print(f"  {'dataset':22s} {'MSE':>10s} {'std':>10s} {'min':>10s} {'max':>10s} | "
          f"{'S0':>6s} {'SD':>6s} {'ratio':>6s} {'erank':>6s}")
    print("-" * 100)

    results = {}
    for ds_name in IMAGE_DATASETS:
        try:
            imgs, mean, std, label = load_dataset_batch(ds_name, s, n)
            out = batched_forward(model, imgs, max_batch=bs)
            mse = F.mse_loss(out['recon'], imgs, reduction='none').mean(dim=(1,2,3))
            S_mean = out['S'].mean(dim=(0,1))
            ratio = (S_mean[0] / (S_mean[-1]+1e-8)).item()
            erank = model.effective_rank(out['S'].reshape(-1, D)).mean().item()

            results[ds_name] = {
                'label': label,
                'mse_mean': mse.mean().item(), 'mse_std': mse.std().item(),
                'mse_min': mse.min().item(), 'mse_max': mse.max().item(),
                'S0': S_mean[0].item(), 'SD': S_mean[-1].item(),
                'ratio': ratio, 'erank': erank,
                'fidelity': (1 - mse.mean()).item() * 100,
            }
            print(f"  {label:22s} {mse.mean():10.6f} {mse.std():10.6f} "
                  f"{mse.min():10.6f} {mse.max():10.6f} | "
                  f"{S_mean[0]:6.3f} {S_mean[-1]:6.3f} {ratio:6.2f} {erank:6.2f}")
        except Exception as e:
            print(f"  {ds_name:22s} FAILED: {e}")
            results[ds_name] = {'error': str(e)}
    return results


def test_noise_types(model, cfg, n=64):
    """Per-type noise reconstruction + geometry."""
    s = cfg['img_size']
    D = cfg['D']
    bs = max(4, 64 // max(1, (s // 64) ** 2))
    print(f"\n{'='*80}")
    print(f"NOISE TYPE BATTERY ({s}×{s}, n={n})")
    print(f"{'='*80}")
    print(f"  {'type':18s} {'MSE':>10s} {'std':>10s} | "
          f"{'S0':>6s} {'SD':>6s} {'ratio':>6s} {'erank':>6s} | "
          f"{'byte_acc':>8s} {'±1_acc':>8s}")
    print("-" * 100)

    results = {}
    for t in range(16):
        name = NOISE_NAMES[t]
        imgs = generate_noise(t, n, s)
        out = batched_forward(model, imgs, max_batch=bs)
        mse = F.mse_loss(out['recon'], imgs, reduction='none').mean(dim=(1,2,3))
        S_mean = out['S'].mean(dim=(0,1))
        ratio = (S_mean[0] / (S_mean[-1]+1e-8)).item()
        erank = model.effective_rank(out['S'].reshape(-1, D)).mean().item()

        # Byte accuracy
        orig_q = ((imgs + 4) / 8 * 255).round().clamp(0,255).long()
        recon_q = ((out['recon'] + 4) / 8 * 255).round().clamp(0,255).long()
        byte_acc = (orig_q == recon_q).float().mean().item()
        byte_1 = ((orig_q - recon_q).abs() <= 1).float().mean().item()

        results[name] = {
            'mse_mean': mse.mean().item(), 'mse_std': mse.std().item(),
            'S0': S_mean[0].item(), 'SD': S_mean[-1].item(),
            'ratio': ratio, 'erank': erank,
            'byte_exact': byte_acc, 'byte_within1': byte_1,
        }
        print(f"  {name:18s} {mse.mean():10.6f} {mse.std():10.6f} | "
              f"{S_mean[0]:6.3f} {S_mean[-1]:6.3f} {ratio:6.2f} {erank:6.2f} | "
              f"{byte_acc*100:7.2f}% {byte_1*100:7.2f}%")
    return results


def test_text_bytes(model, cfg):
    """Text-as-bytes reconstruction."""
    s = cfg['img_size']
    print(f"\n{'='*80}")
    print(f"TEXT BYTE RECONSTRUCTION ({s}×{s})")
    print(f"{'='*80}")

    texts = [
        "Hello, world! This is a test of the geometric encoder.",
        "The quick brown fox jumps over the lazy dog. 0123456789",
        "import torch; model = PatchSVAE(); output = model(x)",
        "E = mc² — Albert Einstein, theoretical physicist, 1905",
        "To be, or not to be, that is the question. — Shakespeare",
        "∫₀^∞ e^(-x²) dx = √π/2 — Gaussian integral",
        "01101000 01100101 01101100 01101100 01101111 — binary hello",
        "SELECT * FROM models WHERE cv BETWEEN 0.20 AND 0.23;",
    ]

    n_bytes = 3 * s * s
    results = {}
    model.eval()

    for text in texts:
        raw = text.encode('utf-8')
        actual_len = min(len(raw), n_bytes)
        padded = (raw + b'\x00' * n_bytes)[:n_bytes]

        arr = np.frombuffer(padded, dtype=np.uint8).copy()
        tensor = torch.from_numpy(arr).float()
        tensor = (tensor / 127.5) - 1.0
        tensor = tensor.reshape(1, 3, s, s).to(DEVICE)

        with torch.no_grad():
            out = model(tensor)
            recon = out['recon']
            mse = F.mse_loss(recon, tensor).item()

        orig_b = ((tensor.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0,255).byte()
        recon_b = ((recon.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0,255).byte()
        exact_acc = (orig_b[:actual_len] == recon_b[:actual_len]).float().mean().item()
        recovered = recon_b[:actual_len].numpy().tobytes().decode('utf-8', errors='replace')

        results[text[:40]] = {'mse': mse, 'byte_acc': exact_acc}
        print(f"\n  In:  '{text[:60]}'")
        print(f"  Out: '{recovered[:60]}'")
        print(f"  MSE: {mse:.6f}  Byte: {exact_acc*100:.1f}%")
    return results


def test_piecemeal(model, cfg):
    """Piecemeal tiling at 4× resolution."""
    s = cfg['img_size']
    src = max(256, s * 4)
    bs = max(2, 32 // max(1, (s // 64) ** 2))
    print(f"\n{'='*80}")
    print(f"PIECEMEAL {src}→{s} TILED RECONSTRUCTION")
    print(f"{'='*80}")
    model.eval()
    results = {}

    test_types = [0, 1, 4, 6, 13]  # Gaussian, Uniform, Pink, Salt-pepper, Cauchy
    with torch.no_grad():
        for t in test_types:
            img_src = generate_noise(t, 1, src).squeeze(0)
            tiles = []
            gh, gw = src // s, src // s
            for gy in range(gh):
                for gx in range(gw):
                    tiles.append(img_src[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s])

            # Batch tiles through model
            all_recon = []
            tile_t = torch.stack(tiles)
            for i in range(0, len(tile_t), bs):
                batch = tile_t[i:i+bs].to(DEVICE)
                out = model(batch)
                all_recon.append(out['recon'].cpu())
            recon_tiles = torch.cat(all_recon)

            recon_full = torch.zeros(3, src, src)
            idx = 0
            for gy in range(gh):
                for gx in range(gw):
                    recon_full[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s] = recon_tiles[idx]
                    idx += 1

            mse = F.mse_loss(recon_full, img_src).item()
            results[NOISE_NAMES[t]] = mse
            print(f"  {NOISE_NAMES[t]:18s}: {gh*gw} tiles, MSE={mse:.6f}")
    return results


def test_signal_survival(model, cfg, n=32):
    """Signal energy survival and SNR per dataset."""
    s = cfg['img_size']
    bs = max(4, 64 // max(1, (s // 64) ** 2))
    print(f"\n{'='*80}")
    print(f"SIGNAL ENERGY SURVIVAL")
    print(f"{'='*80}")
    print(f"  {'source':22s} {'survival':>10s} {'SNR_dB':>10s} {'orig_E':>10s} {'recon_E':>10s}")
    print("-" * 70)

    results = {}

    # Image datasets
    for ds_name in IMAGE_DATASETS:
        try:
            imgs, _, _, label = load_dataset_batch(ds_name, s, n)
            out = batched_forward(model, imgs, max_batch=bs)
            orig_E = (imgs**2).mean().item()
            recon_E = (out['recon']**2).mean().item()
            err_E = ((imgs - out['recon'])**2).mean().item()
            survival = recon_E / (orig_E + 1e-8) * 100
            snr = 10 * math.log10(orig_E / (err_E + 1e-8))
            results[ds_name] = {'survival': survival, 'snr': snr}
            print(f"  {label:22s} {survival:9.1f}% {snr:9.1f}dB {orig_E:10.4f} {recon_E:10.4f}")
        except:
            pass

    # Key noise types
    for t in [0, 4, 6, 13]:
        imgs = generate_noise(t, n, s)
        out = batched_forward(model, imgs, max_batch=bs)
        orig_E = (imgs**2).mean().item()
        recon_E = (out['recon']**2).mean().item()
        err_E = ((imgs - out['recon'])**2).mean().item()
        survival = recon_E / (orig_E + 1e-8) * 100
        snr = 10 * math.log10(orig_E / (err_E + 1e-8))
        results[NOISE_NAMES[t]] = {'survival': survival, 'snr': snr}
        print(f"  noise/{NOISE_NAMES[t]:17s} {survival:9.1f}% {snr:9.1f}dB {orig_E:10.4f} {recon_E:10.4f}")

    return results


def test_alpha_profile(model):
    """Cross-attention alpha analysis."""
    print(f"\n{'='*80}")
    print("ALPHA PROFILE")
    print(f"{'='*80}")
    results = {}
    for li, layer in enumerate(model.cross_attn):
        alpha = layer.alpha.detach().cpu()
        results[f'layer_{li}'] = {
            'mean': alpha.mean().item(), 'max': alpha.max().item(),
            'min': alpha.min().item(), 'std': alpha.std().item(),
            'values': alpha.tolist(),
        }
        print(f"\n  Layer {li}: mean={alpha.mean():.5f}  max={alpha.max():.5f}  "
              f"min={alpha.min():.5f}  std={alpha.std():.6f}")
        bar_scale = 50 / (alpha.max().item() + 1e-8)
        for d in range(len(alpha)):
            bar = "█" * int(alpha[d].item() * bar_scale)
            print(f"    α[{d:2d}]: {alpha[d]:.5f}  {bar}")
    return results


def test_compression(model, cfg):
    """Compression metrics."""
    s = cfg['img_size']
    D = cfg['D']; ps = cfg['patch_size']
    n_patches = (s // ps) ** 2
    input_vals = 3 * s * s
    latent_vals = D * n_patches
    ratio = input_vals / latent_vals
    print(f"\n{'='*80}")
    print("COMPRESSION METRICS")
    print(f"{'='*80}")
    print(f"  Input:   {s}×{s}×3 = {input_vals:,} values")
    print(f"  Latent:  {D}×{n_patches} = {latent_vals:,} omega tokens")
    print(f"  Ratio:   {ratio:.1f}:1")
    for bits in [8, 16, 32]:
        ib = input_vals * (bits//8)
        lb = latent_vals * (bits//8)
        print(f"  {bits}-bit: input={ib/1024:.1f}KB  latent={lb/1024:.1f}KB  ratio={ib/lb:.1f}:1")
    return {'input_values': input_vals, 'latent_values': latent_vals, 'ratio': ratio}


def test_reconstruction_grid(model, cfg):
    """Save visual grid: 5 image datasets + 4 key noise types."""
    s = cfg['img_size']
    print(f"\n{'='*80}")
    print("RECONSTRUCTION GRID")
    print(f"{'='*80}")
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    rows = []
    labels = []

    # Image datasets
    for ds_name in IMAGE_DATASETS:
        try:
            imgs, mean, std, label = load_dataset_batch(ds_name, s, 2)
            mean_t = torch.tensor(mean).reshape(1,3,1,1)
            std_t = torch.tensor(std).reshape(1,3,1,1)
            out = batched_forward(model, imgs[:1], max_batch=1)
            orig_vis = (imgs[:1] * std_t + mean_t).clamp(0,1)
            recon_vis = (out['recon'][:1] * std_t + mean_t).clamp(0,1)
            rows.append((orig_vis[0], recon_vis[0]))
            labels.append(label)
        except:
            pass

    # Key noise types
    for t in [0, 6, 13, 4]:
        imgs = generate_noise(t, 1, s)
        out = batched_forward(model, imgs, max_batch=1)
        o = imgs[0].clamp(-3,3); r = out['recon'][0].clamp(-3,3)
        o = (o - o.min())/(o.max()-o.min()+1e-8)
        r = (r - r.min())/(r.max()-r.min()+1e-8)
        rows.append((o, r))
        labels.append(f'noise/{NOISE_NAMES[t]}')

    n_rows = len(rows)
    fig, axes = plt.subplots(n_rows, 3, figsize=(9, n_rows*3))
    if n_rows == 1: axes = axes.reshape(1, -1)
    for i, (orig, recon) in enumerate(rows):
        diff = (orig - recon).abs().clamp(0,1)
        axes[i,0].imshow(orig.permute(1,2,0).numpy())
        axes[i,1].imshow(recon.permute(1,2,0).numpy())
        axes[i,2].imshow((diff * 5).clamp(0,1).permute(1,2,0).numpy())
        axes[i,0].set_ylabel(labels[i], fontsize=8)
        for j in range(3): axes[i,j].axis('off')
    axes[0,0].set_title('Original', fontsize=9)
    axes[0,1].set_title('Recon', fontsize=9)
    axes[0,2].set_title('|Err|×5', fontsize=9)
    plt.tight_layout()
    fname = 'universal_diagnostic_grid.png'
    plt.savefig(fname, dpi=150, bbox_inches='tight')
    print(f"  Saved: {fname}")
    plt.close()


# ════════════════════════════════════════════════════════════════
# MAIN
# ════════════════════════════════════════════════════════════════

def run(hf_version=None, checkpoint_path=None, n_samples=64):
    """
    Run full diagnostic battery.

    Usage in Colab cell:
        run(hf_version='v13_imagenet256')
        run(hf_version='v16_johanna_omega')
        run(hf_version='v18_johanna_curriculum')
        run(checkpoint_path='/content/checkpoints/best.pt')
    """
    print("=" * 80)
    print("UNIVERSAL SVAE DIAGNOSTIC BATTERY")
    print("=" * 80)

    model, cfg = load_model(hf_version=hf_version, checkpoint_path=checkpoint_path)

    # Infer img_size if not in config
    if 'img_size' not in cfg:
        ds = cfg.get('dataset', '')
        if '256' in ds: cfg['img_size'] = 256
        elif '128' in ds: cfg['img_size'] = 128
        elif 'tiny' in ds: cfg['img_size'] = 64
        elif 'cifar' in ds: cfg['img_size'] = 32
        else: cfg['img_size'] = 64
        print(f"  Inferred img_size={cfg['img_size']} from dataset='{ds}'")

    s = cfg['img_size']
    n = min(n_samples, max(16, 100 // max(1, (s // 64) ** 2)))
    print(f"  Resolution: {s}×{s}, samples_per_test: {n}")

    results = {'config': cfg}

    results['image_datasets'] = test_image_datasets(model, cfg, n=n)
    results['noise_types'] = test_noise_types(model, cfg, n=n)
    results['text'] = test_text_bytes(model, cfg)
    results['piecemeal'] = test_piecemeal(model, cfg)
    results['signal_survival'] = test_signal_survival(model, cfg, n=n)
    results['alpha'] = test_alpha_profile(model)
    results['compression'] = test_compression(model, cfg)
    test_reconstruction_grid(model, cfg)

    tag = hf_version or 'local'
    out_path = f'diagnostic_{tag.replace("/","_")}.json'
    with open(out_path, 'w') as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\n  Results: {out_path}")
    print(f"\n{'='*80}")
    print("DIAGNOSTIC COMPLETE")
    print(f"{'='*80}")
    return results


# ── CONFIG: Change this per run ──────────────────────────────────
# Uncomment the model you want to diagnose:

# HF_VERSION = 'v13_imagenet256'        # Fresnel-base 256
HF_VERSION = 'v16_johanna_omega'    # Johanna-small 128
# HF_VERSION = 'v18_johanna_curriculum'  # Johanna-tiny 64
# HF_VERSION = None; CHECKPOINT = '/content/checkpoints/best.pt'

if __name__ == "__main__":
    run(hf_version=HF_VERSION)