v18_johanna_curriculum/trainer.py

"""
Johanna-Tiny Full Battery Diagnostic
======================================
Comprehensive analysis of the curriculum-trained 16-type noise model.

Tests:
  1. Per-type MSE (100 samples each, full eval)
  2. Per-type byte accuracy (discrete reconstruction precision)
  3. Geometric fingerprint per noise type (S₀, ratio, erank, CV)
  4. Cross-type omega token similarity (cosine distance matrix)
  5. Spectrum analysis per type (which modes carry which distributions)
  6. Reconstruction visualization grid (all 16 types)
  7. Zero-shot transfer: real images through noise-trained model
  8. Zero-shot transfer: text bytes through noise-trained model
  9. Piecemeal 256→64: can tiny do tiled reconstruction?
  10. Noise-to-noise: encode type A, does it look like type A?
  11. Effective capacity: what percentage of the signal survives?
  12. Alpha profile: what did the cross-attention learn?
"""

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

# ── Load model ───────────────────────────────────────────────────

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

# Option 1: Load from local checkpoint
CHECKPOINT = '/content/checkpoints/best.pt'
# Option 2: Load from HuggingFace
HF_CHECKPOINT = 'AbstractPhil/geolip-SVAE'
HF_FILE = 'v18_johanna_curriculum/checkpoints/epoch_0300.pt'


def load_model():
    """Load model from local or HF checkpoint."""
    from huggingface_hub import hf_hub_download

    # Try local first
    if os.path.exists(CHECKPOINT):
        path = CHECKPOINT
        print(f"  Loading local: {path}")
    else:
        path = hf_hub_download(repo_id=HF_CHECKPOINT, filename=HF_FILE, repo_type="model")
        print(f"  Loading HF: {HF_FILE}")

    ckpt = torch.load(path, map_location='cpu', weights_only=False)
    cfg = ckpt['config']
    print(f"  Epoch: {ckpt.get('epoch')}, MSE: {ckpt.get('test_mse', '?')}")
    print(f"  Config: {cfg}")

    # Build model inline (same architecture)
    from types import SimpleNamespace

    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, 'gh': gh, 'gw': gw}

        @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()

    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"  Loaded {sum(p.numel() for p in model.parameters()):,} params")
    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=64):
    """Generate n samples of a given noise type."""
    imgs = []
    rng = np.random.RandomState(42)
    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)


# ════════════════════════════════════════════════════════════════
# DIAGNOSTIC TESTS
# ════════════════════════════════════════════════════════════════

def test_1_per_type_mse(model, n=100, s=64):
    """Per-type reconstruction MSE."""
    print(f"\n{'='*70}")
    print("TEST 1: Per-Type Reconstruction MSE (100 samples each)")
    print(f"{'='*70}")
    results = {}
    model.eval()
    with torch.no_grad():
        for t in range(16):
            imgs = generate_noise(t, n, s).to(DEVICE)
            out = model(imgs)
            mse = F.mse_loss(out['recon'], imgs, reduction='none').mean(dim=(1,2,3))
            results[NOISE_NAMES[t]] = {
                'mean': mse.mean().item(),
                'std': mse.std().item(),
                'min': mse.min().item(),
                'max': mse.max().item(),
            }
            print(f"  {NOISE_NAMES[t]:18s}: {mse.mean():.6f} ± {mse.std():.6f}  "
                  f"[{mse.min():.6f} — {mse.max():.6f}]")
    return results


def test_2_byte_accuracy(model, n=100, s=64):
    """Byte-level reconstruction accuracy per type."""
    print(f"\n{'='*70}")
    print("TEST 2: Byte-Level Accuracy (quantized to 256 levels)")
    print(f"{'='*70}")
    results = {}
    model.eval()
    with torch.no_grad():
        for t in range(16):
            imgs = generate_noise(t, n, s).to(DEVICE)
            out = model(imgs)
            # Quantize to 256 levels
            orig_q = ((imgs + 4) / 8 * 255).round().clamp(0, 255).long()
            recon_q = ((out['recon'] + 4) / 8 * 255).round().clamp(0, 255).long()
            acc = (orig_q == recon_q).float().mean().item()
            # Within-1 accuracy
            acc1 = ((orig_q - recon_q).abs() <= 1).float().mean().item()
            results[NOISE_NAMES[t]] = {'exact': acc, 'within_1': acc1}
            print(f"  {NOISE_NAMES[t]:18s}: exact={acc*100:5.1f}%  ±1={acc1*100:5.1f}%")
    return results


def test_3_geometric_fingerprint(model, n=64, s=64):
    """Geometric properties per noise type."""
    print(f"\n{'='*70}")
    print("TEST 3: Geometric Fingerprint Per Type")
    print(f"{'='*70}")
    D = model.D
    results = {}
    model.eval()
    with torch.no_grad():
        for t in range(16):
            imgs = generate_noise(t, n, s).to(DEVICE)
            out = model(imgs)
            S = out['svd']['S']  # (B, N, D)
            S_mean = S.mean(dim=(0, 1))
            ratio = (S_mean[0] / (S_mean[-1] + 1e-8)).item()
            erank = model.effective_rank(S.reshape(-1, D)).mean().item()
            s0 = S_mean[0].item()
            sd = S_mean[-1].item()
            results[NOISE_NAMES[t]] = {'S0': s0, 'SD': sd, 'ratio': ratio, 'erank': erank}
            print(f"  {NOISE_NAMES[t]:18s}: S₀={s0:.3f}  SD={sd:.3f}  "
                  f"ratio={ratio:.2f}  erank={erank:.2f}")
    return results


def test_4_omega_similarity(model, n=32, s=64):
    """Cross-type omega token cosine similarity matrix."""
    print(f"\n{'='*70}")
    print("TEST 4: Cross-Type Omega Token Similarity")
    print(f"{'='*70}")
    D = model.D
    type_centroids = {}
    model.eval()
    with torch.no_grad():
        for t in range(16):
            imgs = generate_noise(t, n, s).to(DEVICE)
            out = model(imgs)
            # Average omega token per type: (D,)
            omega = out['svd']['S'].mean(dim=(0, 1))
            type_centroids[t] = omega

    # Cosine similarity matrix
    keys = sorted(type_centroids.keys())
    centroids = torch.stack([type_centroids[k] for k in keys])
    centroids_norm = F.normalize(centroids, dim=-1)
    sim_matrix = centroids_norm @ centroids_norm.T

    # Print matrix
    header = "         " + " ".join([f"{NOISE_NAMES[k][:5]:>5s}" for k in keys])
    print(f"  {header}")
    for i, ki in enumerate(keys):
        row = f"  {NOISE_NAMES[ki]:8s}"
        for j, kj in enumerate(keys):
            v = sim_matrix[i, j].item()
            row += f" {v:5.2f}"
        print(row)
    return sim_matrix.cpu()


def test_5_spectrum_per_type(model, n=64, s=64):
    """Singular value spectrum analysis per type."""
    print(f"\n{'='*70}")
    print("TEST 5: Spectrum Profile Per Type")
    print(f"{'='*70}")
    D = model.D
    results = {}
    model.eval()
    with torch.no_grad():
        for t in range(16):
            imgs = generate_noise(t, n, s).to(DEVICE)
            out = model(imgs)
            S_mean = out['svd']['S'].mean(dim=(0, 1))
            total = (S_mean**2).sum()
            cum = 0
            spectrum = []
            for d in range(D):
                e = (S_mean[d]**2).item()
                cum += e
                spectrum.append({'value': S_mean[d].item(), 'energy_pct': cum/total.item()*100})
            results[NOISE_NAMES[t]] = spectrum

    # Print top-3 and bottom-3 modes per type
    for t in range(16):
        name = NOISE_NAMES[t]
        sp = results[name]
        top = f"S0={sp[0]['value']:.3f}({sp[0]['energy_pct']:.1f}%)"
        mid = f"S7={sp[7]['value']:.3f}({sp[7]['energy_pct']:.1f}%)"
        bot = f"S15={sp[15]['value']:.3f}(100%)"
        print(f"  {name:18s}: {top}  {mid}  {bot}")
    return results


def test_6_reconstruction_grid(model, s=64):
    """Visual reconstruction grid — all 16 types."""
    print(f"\n{'='*70}")
    print("TEST 6: Reconstruction Grid (saved to johanna_diagnostic_grid.png)")
    print(f"{'='*70}")
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    model.eval()
    fig, axes = plt.subplots(16, 3, figsize=(9, 48))

    with torch.no_grad():
        for t in range(16):
            img = generate_noise(t, 1, s).to(DEVICE)
            out = model(img)
            recon = out['recon']
            mse = F.mse_loss(recon, img).item()

            orig_np = img[0].cpu().clamp(-3, 3).permute(1, 2, 0).numpy()
            recon_np = recon[0].cpu().clamp(-3, 3).permute(1, 2, 0).numpy()
            diff_np = (img[0] - recon[0]).abs().cpu().clamp(0, 2).permute(1, 2, 0).numpy()

            # Normalize for display
            for arr in [orig_np, recon_np]:
                arr -= arr.min(); arr /= (arr.max() + 1e-8)
            diff_np /= (diff_np.max() + 1e-8)

            axes[t, 0].imshow(orig_np); axes[t, 0].set_ylabel(NOISE_NAMES[t], fontsize=8)
            axes[t, 1].imshow(recon_np)
            axes[t, 2].imshow(diff_np)
            for j in range(3):
                axes[t, j].axis('off')

    axes[0, 0].set_title('Original', fontsize=9)
    axes[0, 1].set_title('Recon', fontsize=9)
    axes[0, 2].set_title('|Error|', fontsize=9)
    plt.tight_layout()
    plt.savefig('johanna_diagnostic_grid.png', dpi=150, bbox_inches='tight')
    print("  Saved: johanna_diagnostic_grid.png")
    plt.close()


def test_7_real_images(model, s=64):
    """Zero-shot: real images through noise-trained model."""
    print(f"\n{'='*70}")
    print("TEST 7: Zero-Shot Real Image Reconstruction")
    print(f"{'='*70}")
    from datasets import load_dataset

    ds = load_dataset('zh-plus/tiny-imagenet', split='valid', streaming=True)
    transform = T.Compose([T.ToTensor(), T.Normalize((0.4802,0.4481,0.3975),(0.2770,0.2691,0.2821))])

    imgs = []
    for i, sample in enumerate(ds):
        img = sample['image'].convert('RGB')
        imgs.append(transform(img))
        if i >= 99:
            break

    batch = torch.stack(imgs).to(DEVICE)
    model.eval()
    with torch.no_grad():
        out = model(batch)
        mse = F.mse_loss(out['recon'], batch, reduction='none').mean(dim=(1,2,3))

    print(f"  TinyImageNet (100 images, {s}×{s}):")
    print(f"    Mean MSE: {mse.mean():.6f}")
    print(f"    Std:      {mse.std():.6f}")
    print(f"    Min/Max:  {mse.min():.6f} / {mse.max():.6f}")
    print(f"    Fidelity: {(1 - mse.mean())*100:.3f}%")
    return {'mean': mse.mean().item(), 'std': mse.std().item()}


def test_8_text_bytes(model, s=64):
    """Zero-shot: text through noise-trained model."""
    print(f"\n{'='*70}")
    print("TEST 8: Zero-Shot Text Byte Reconstruction")
    print(f"{'='*70}")

    texts = [
        "Hello, world! This is a test of the Johanna geometric encoder.",
        "The quick brown fox jumps over the lazy dog. 0123456789 ABCDEF",
        "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",
    ]

    n_bytes = 3 * s * s
    model.eval()

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

        arr = np.frombuffer(raw, 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()

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

        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()

        print(f"\n  Input:    '{text[:60]}'")
        print(f"  Output:   '{recovered[:60]}'")
        print(f"  MSE:      {mse:.6f}")
        print(f"  Byte acc: {exact_acc*100:.1f}%")


def test_9_piecemeal(model, s=64):
    """Piecemeal: tile 256×256 noise into 64×64 tiles."""
    print(f"\n{'='*70}")
    print(f"TEST 9: Piecemeal 256→{s} Tiled Reconstruction")
    print(f"{'='*70}")
    model.eval()

    results = {}
    with torch.no_grad():
        for t in [0, 4, 6, 13]:  # Gaussian, Pink, Salt-pepper, Cauchy
            img_256 = generate_noise(t, 1, 256).squeeze(0)  # (3, 256, 256)
            tiles = []
            gh, gw = 256 // s, 256 // s
            for gy in range(gh):
                for gx in range(gw):
                    tile = img_256[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s]
                    tiles.append(tile)
            tile_batch = torch.stack(tiles).to(DEVICE)
            out = model(tile_batch)
            recon_tiles = out['recon'].cpu()

            # Stitch
            recon_full = torch.zeros(3, 256, 256)
            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_256).item()
            results[NOISE_NAMES[t]] = mse
            n_tiles = gh * gw
            print(f"  {NOISE_NAMES[t]:18s}: {n_tiles} tiles, MSE={mse:.6f}")
    return results


def test_10_signal_survival(model, n=100, s=64):
    """What percentage of the original signal energy survives reconstruction?"""
    print(f"\n{'='*70}")
    print("TEST 10: Signal Energy Survival Rate")
    print(f"{'='*70}")
    model.eval()
    with torch.no_grad():
        for t in range(16):
            imgs = generate_noise(t, n, s).to(DEVICE)
            out = model(imgs)
            recon = out['recon']
            orig_energy = (imgs**2).mean().item()
            recon_energy = (recon**2).mean().item()
            error_energy = ((imgs - recon)**2).mean().item()
            survival = recon_energy / (orig_energy + 1e-8) * 100
            snr = 10 * math.log10(orig_energy / (error_energy + 1e-8))
            print(f"  {NOISE_NAMES[t]:18s}: survival={survival:6.1f}%  SNR={snr:5.1f}dB  "
                  f"orig_E={orig_energy:.3f}  recon_E={recon_energy:.3f}")


def test_11_alpha_profile(model):
    """Cross-attention alpha analysis."""
    print(f"\n{'='*70}")
    print("TEST 11: Cross-Attention Alpha Profile")
    print(f"{'='*70}")
    for li, layer in enumerate(model.cross_attn):
        alpha = layer.alpha.detach().cpu()
        print(f"\n  Layer {li}: mean={alpha.mean():.4f}  max={alpha.max():.4f}  "
              f"min={alpha.min():.4f}  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}")


def test_12_compression_ratio(model, s=64):
    """Actual compression metrics."""
    print(f"\n{'='*70}")
    print("TEST 12: Compression Metrics")
    print(f"{'='*70}")
    D = model.D
    ps = model.patch_size
    n_patches = (s // ps) ** 2
    input_values = 3 * s * s
    latent_values = D * n_patches
    ratio = input_values / latent_values
    print(f"  Input:   {s}×{s}×3 = {input_values:,} values")
    print(f"  Latent:  {D}×{n_patches} = {latent_values:,} values (omega tokens)")
    print(f"  Ratio:   {ratio:.1f}:1 compression")
    print(f"  Patches: {n_patches} of {ps}×{ps}")
    print(f"  Omega shape: ({D}, {s//ps}, {s//ps})")

    # Bits per value at different quantization levels
    for bits in [8, 16, 32]:
        input_bytes = input_values * (bits // 8)
        latent_bytes = latent_values * (bits // 8)
        print(f"  At {bits}-bit: input={input_bytes/1024:.1f}KB  latent={latent_bytes/1024:.1f}KB  "
              f"ratio={input_bytes/latent_bytes:.1f}:1")


# ════════════════════════════════════════════════════════════════
# RUN ALL
# ════════════════════════════════════════════════════════════════

def run_all():
    print("=" * 70)
    print("JOHANNA-TINY FULL BATTERY DIAGNOSTIC")
    print("=" * 70)

    model, cfg = load_model()
    s = cfg.get('img_size', 64)

    all_results = {}
    all_results['config'] = cfg

    all_results['per_type_mse'] = test_1_per_type_mse(model, n=100, s=s)
    all_results['byte_accuracy'] = test_2_byte_accuracy(model, n=100, s=s)
    all_results['geometry'] = test_3_geometric_fingerprint(model, n=64, s=s)
    sim_matrix = test_4_omega_similarity(model, n=32, s=s)
    all_results['spectrum'] = test_5_spectrum_per_type(model, n=64, s=s)
    test_6_reconstruction_grid(model, s=s)
    all_results['real_images'] = test_7_real_images(model, s=s)
    test_8_text_bytes(model, s=s)
    all_results['piecemeal'] = test_9_piecemeal(model, s=s)
    test_10_signal_survival(model, n=100, s=s)
    test_11_alpha_profile(model)
    test_12_compression_ratio(model, s=s)

    # Save results
    out_path = 'johanna_diagnostic_results.json'
    with open(out_path, 'w') as f:
        json.dump(all_results, f, indent=2, default=str)
    print(f"\n  Results saved: {out_path}")

    print(f"\n{'='*70}")
    print("DIAGNOSTIC COMPLETE")
    print(f"{'='*70}")


if __name__ == "__main__":
    run_all()