#@title MobiusNet Aggregate Analysis - 1024 Samples
!pip install -q datasets safetensors huggingface_hub

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file as load_safetensors
import matplotlib.pyplot as plt
import numpy as np
from sklearn.decomposition import PCA
from typing import Tuple
from collections import defaultdict
import math
import json

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

# ============================================================================
# MOBIUSNET (compact)
# ============================================================================

class MobiusLens(nn.Module):
    def __init__(self, dim, layer_idx, total_layers, scale_range=(1.0, 9.0)):
        super().__init__()
        self.t = layer_idx / max(total_layers - 1, 1)
        scale_span = scale_range[1] - scale_range[0]
        step = scale_span / max(total_layers, 1)
        self.register_buffer('scales', torch.tensor([scale_range[0] + self.t * scale_span,
                                                      scale_range[0] + self.t * scale_span + step]))
        self.twist_in_angle = nn.Parameter(torch.tensor(self.t * math.pi))
        self.twist_in_proj = nn.Linear(dim, dim, bias=False)
        self.omega = nn.Parameter(torch.tensor(math.pi))
        self.alpha = nn.Parameter(torch.tensor(1.5))
        self.phase_l, self.drift_l = nn.Parameter(torch.zeros(2)), nn.Parameter(torch.ones(2))
        self.phase_m, self.drift_m = nn.Parameter(torch.zeros(2)), nn.Parameter(torch.zeros(2))
        self.phase_r, self.drift_r = nn.Parameter(torch.zeros(2)), nn.Parameter(-torch.ones(2))
        self.accum_weights = nn.Parameter(torch.tensor([0.4, 0.2, 0.4]))
        self.xor_weight = nn.Parameter(torch.tensor(0.7))
        self.gate_norm = nn.LayerNorm(dim)
        self.twist_out_angle = nn.Parameter(torch.tensor(-self.t * math.pi))
        self.twist_out_proj = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        cos_t, sin_t = torch.cos(self.twist_in_angle), torch.sin(self.twist_in_angle)
        x = x * cos_t + self.twist_in_proj(x) * sin_t
        x_norm = torch.tanh(x)
        t = x_norm.abs().mean(dim=-1, keepdim=True).unsqueeze(-2)
        x_exp = x_norm.unsqueeze(-2)
        s = self.scales.view(-1, 1)
        def wave(phase, drift):
            a = self.alpha.abs() + 0.1
            pos = s * self.omega * (x_exp + drift.view(-1, 1) * t) + phase.view(-1, 1)
            return torch.exp(-a * torch.sin(pos).pow(2)).prod(dim=-2)
        L, M, R = wave(self.phase_l, self.drift_l), wave(self.phase_m, self.drift_m), wave(self.phase_r, self.drift_r)
        w = torch.softmax(self.accum_weights, dim=0)
        xor_w = torch.sigmoid(self.xor_weight)
        lr = xor_w * (L + R - 2*L*R).abs() + (1 - xor_w) * L * R
        gate = torch.sigmoid(self.gate_norm((w[0]*L + w[1]*M + w[2]*R) * (0.5 + 0.5*lr)))
        x = x * gate
        cos_t, sin_t = torch.cos(self.twist_out_angle), torch.sin(self.twist_out_angle)
        return x * cos_t + self.twist_out_proj(x) * sin_t, gate

class MobiusConvBlock(nn.Module):
    def __init__(self, channels, layer_idx, total_layers, scale_range=(1.0, 9.0), reduction=0.5):
        super().__init__()
        self.conv = nn.Sequential(nn.Conv2d(channels, channels, 3, padding=1, groups=channels, bias=False),
                                   nn.Conv2d(channels, channels, 1, bias=False), nn.BatchNorm2d(channels))
        self.lens = MobiusLens(channels, layer_idx, total_layers, scale_range)
        third = channels // 3
        which_third = layer_idx % 3
        mask = torch.ones(channels)
        mask[which_third*third : which_third*third + third + (channels%3 if which_third==2 else 0)] = reduction
        self.register_buffer('thirds_mask', mask.view(1, -1, 1, 1))
        self.residual_weight = nn.Parameter(torch.tensor(0.9))

    def forward(self, x):
        identity = x
        h = self.conv(x).permute(0, 2, 3, 1)
        h, gate = self.lens(h)
        h = h.permute(0, 3, 1, 2) * self.thirds_mask
        rw = torch.sigmoid(self.residual_weight)
        return rw * identity + (1 - rw) * h, gate

class MobiusNet(nn.Module):
    def __init__(self, in_chans=1, num_classes=1000, channels=(64,128,256),
                 depths=(2,2,2), scale_range=(0.5,2.5), use_integrator=True):
        super().__init__()
        total_layers = sum(depths)
        channels = list(channels)
        self.stem = nn.Sequential(nn.Conv2d(in_chans, channels[0], 3, padding=1, bias=False), nn.BatchNorm2d(channels[0]))
        self.stages = nn.ModuleList()
        self.downsamples = nn.ModuleList()
        layer_idx = 0
        for si, d in enumerate(depths):
            self.stages.append(nn.ModuleList([MobiusConvBlock(channels[si], layer_idx+i, total_layers, scale_range) for i in range(d)]))
            layer_idx += d
            if si < len(depths)-1:
                self.downsamples.append(nn.Sequential(nn.Conv2d(channels[si], channels[si+1], 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(channels[si+1])))
        self.integrator = nn.Sequential(nn.Conv2d(channels[-1], channels[-1], 3, padding=1, bias=False), nn.BatchNorm2d(channels[-1]), nn.GELU()) if use_integrator else nn.Identity()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Linear(channels[-1], num_classes)

    def forward_with_intermediates(self, x):
        out = {'input': x, 'stem': None, 'stages': [], 'gates': [], 'final': None}
        x = self.stem(x)
        out['stem'] = x
        for i, stage in enumerate(self.stages):
            acts, gates = [], []
            for block in stage:
                x, g = block(x)
                acts.append(x)
                gates.append(g)
            out['stages'].append(acts)
            out['gates'].append(gates)
            if i < len(self.downsamples):
                x = self.downsamples[i](x)
        x = self.integrator(x)
        out['final'] = x
        return self.head(self.pool(x).flatten(1)), out

# ============================================================================
# LOAD MODEL
# ============================================================================

print("Loading model...")
config_path = hf_hub_download("AbstractPhil/mobiusnet-distillations",
    "checkpoints/mobius_tiny_s_imagenet_clip_vit_l14/20260111_000512/config.json")
with open(config_path) as f:
    config = json.load(f)
model_path = hf_hub_download("AbstractPhil/mobiusnet-distillations",
    "checkpoints/mobius_tiny_s_imagenet_clip_vit_l14/20260111_000512/checkpoints/best_model.safetensors")

cfg = config['model']
model = MobiusNet(cfg['in_chans'], cfg['num_classes'], tuple(cfg['channels']),
                  tuple(cfg['depths']), tuple(cfg['scale_range']), cfg['use_integrator']).to(device)
model.load_state_dict(load_safetensors(model_path))
model.eval()
print(f"✓ Loaded MobiusNet")

# ============================================================================
# AGGREGATE OVER 1024 SAMPLES
# ============================================================================

print("\nProcessing 1024 samples...")
ds = load_dataset("AbstractPhil/imagenet-clip-features-orderly", "clip_vit_l14",
                  split="validation", streaming=True).with_format("torch")
loader = DataLoader(ds, batch_size=64)

# Accumulators
n_samples = 0
total_correct = 0
agg = {
    'input': {'sum': None, 'sum_sq': None},
    'stem': {'sum': None, 'sum_sq': None},
    'final': {'sum': None, 'sum_sq': None},
}
gate_stats = defaultdict(lambda: {'sum': 0, 'sum_sq': 0, 'min': float('inf'), 'max': float('-inf'), 'count': 0})
stage_stats = defaultdict(lambda: {'sum': None, 'sum_sq': None})

# Class-wise gate means
class_gate_means = defaultdict(lambda: defaultdict(list))

for batch_idx, batch in enumerate(loader):
    if n_samples >= 1024:
        break
    
    features = batch['clip_features'].view(-1, 1, 24, 32).to(device)
    labels = batch['label'].to(device)
    bs = features.shape[0]
    
    with torch.no_grad():
        logits, intermediates = model.forward_with_intermediates(features)
        preds = logits.argmax(dim=-1)
        total_correct += (preds == labels).sum().item()
    
    # Aggregate inputs, stem, final
    for key in ['input', 'stem', 'final']:
        tensor = intermediates[key].detach()
        if agg[key]['sum'] is None:
            agg[key]['sum'] = tensor.sum(dim=0)
            agg[key]['sum_sq'] = (tensor ** 2).sum(dim=0)
        else:
            agg[key]['sum'] += tensor.sum(dim=0)
            agg[key]['sum_sq'] += (tensor ** 2).sum(dim=0)
    
    # Aggregate gates and stages
    for si, (acts, gates) in enumerate(zip(intermediates['stages'], intermediates['gates'])):
        for bi, (act, gate) in enumerate(zip(acts, gates)):
            key = f"S{si}B{bi}"
            
            # Gate stats
            g = gate.detach()
            gate_stats[key]['sum'] += g.mean().item() * bs
            gate_stats[key]['sum_sq'] += (g.mean(dim=(1,2,3)) ** 2).sum().item()
            gate_stats[key]['min'] = min(gate_stats[key]['min'], g.min().item())
            gate_stats[key]['max'] = max(gate_stats[key]['max'], g.max().item())
            gate_stats[key]['count'] += bs
            
            # Stage activation stats
            a = act.detach()
            if stage_stats[key]['sum'] is None:
                stage_stats[key]['sum'] = a.sum(dim=0)
                stage_stats[key]['sum_sq'] = (a ** 2).sum(dim=0)
            else:
                stage_stats[key]['sum'] += a.sum(dim=0)
                stage_stats[key]['sum_sq'] += (a ** 2).sum(dim=0)
            
            # Per-class gate means
            for i in range(bs):
                lbl = labels[i].item()
                class_gate_means[key][lbl].append(g[i].mean().item())
    
    n_samples += bs
    if (batch_idx + 1) % 4 == 0:
        print(f"  Processed {n_samples} samples...")

print(f"\n✓ Processed {n_samples} samples")
print(f"✓ Overall accuracy: {total_correct / n_samples:.2%}")

# ============================================================================
# COMPUTE MEANS AND STDS
# ============================================================================

def compute_mean_std(agg_dict, n):
    mean = agg_dict['sum'] / n
    std = torch.sqrt(agg_dict['sum_sq'] / n - mean ** 2 + 1e-8)
    return mean.cpu().numpy(), std.cpu().numpy()

input_mean, input_std = compute_mean_std(agg['input'], n_samples)
stem_mean, stem_std = compute_mean_std(agg['stem'], n_samples)
final_mean, final_std = compute_mean_std(agg['final'], n_samples)

stage_means, stage_stds = {}, {}
for key in stage_stats:
    stage_means[key], stage_stds[key] = compute_mean_std(stage_stats[key], n_samples)

# ============================================================================
# VISUALIZATION
# ============================================================================

fig = plt.figure(figsize=(24, 18))

# Row 1: Input/Stem aggregates
ax = fig.add_subplot(4, 6, 1)
ax.imshow(input_mean[0], cmap='viridis', aspect='auto')
ax.set_title(f"Input Mean\n[1×24×32]", fontsize=10)
ax.axis('off')

ax = fig.add_subplot(4, 6, 2)
ax.imshow(input_std[0], cmap='magma', aspect='auto')
ax.set_title(f"Input Std\nσ={input_std.mean():.4f}", fontsize=10)
ax.axis('off')

ax = fig.add_subplot(4, 6, 3)
pca = PCA(n_components=3)
stem_pca = pca.fit_transform(stem_mean.reshape(64, -1).T).reshape(24, 32, 3)
stem_pca = (stem_pca - stem_pca.min()) / (stem_pca.max() - stem_pca.min() + 1e-8)
ax.imshow(stem_pca, aspect='auto')
ax.set_title(f"Stem Mean PCA\n({pca.explained_variance_ratio_.sum()*100:.1f}%)", fontsize=10)
ax.axis('off')

ax = fig.add_subplot(4, 6, 4)
ax.imshow(stem_std.mean(axis=0), cmap='magma', aspect='auto')
ax.set_title(f"Stem Std (avg ch)\nσ={stem_std.mean():.4f}", fontsize=10)
ax.axis('off')

ax = fig.add_subplot(4, 6, 5)
pca = PCA(n_components=3)
final_pca = pca.fit_transform(final_mean.reshape(256, -1).T).reshape(6, 8, 3)
final_pca = (final_pca - final_pca.min()) / (final_pca.max() - final_pca.min() + 1e-8)
ax.imshow(final_pca, aspect='auto')
ax.set_title(f"Final Mean PCA\n({pca.explained_variance_ratio_.sum()*100:.1f}%)", fontsize=10)
ax.axis('off')

ax = fig.add_subplot(4, 6, 6)
ax.imshow(np.linalg.norm(final_mean, axis=0), cmap='hot', aspect='auto')
ax.set_title(f"Final Mean L2\nμ={np.linalg.norm(final_mean, axis=0).mean():.2f}", fontsize=10)
ax.axis('off')

# Row 2: Stage means
for i, key in enumerate(['S0B0', 'S0B1', 'S1B0', 'S1B1', 'S2B0', 'S2B1']):
    ax = fig.add_subplot(4, 6, 7 + i)
    m = stage_means[key]
    pca = PCA(n_components=3)
    m_pca = pca.fit_transform(m.reshape(m.shape[0], -1).T).reshape(m.shape[1], m.shape[2], 3)
    m_pca = (m_pca - m_pca.min()) / (m_pca.max() - m_pca.min() + 1e-8)
    ax.imshow(m_pca, aspect='auto')
    ax.set_title(f"{key} Mean\n[{m.shape[0]}×{m.shape[1]}×{m.shape[2]}]", fontsize=10)
    ax.axis('off')

# Row 3: Stage stds + Gate summary
for i, key in enumerate(['S0B0', 'S0B1', 'S1B0', 'S1B1', 'S2B0', 'S2B1']):
    ax = fig.add_subplot(4, 6, 13 + i)
    s = stage_stds[key]
    ax.imshow(s.mean(axis=0), cmap='magma', aspect='auto')
    gs = gate_stats[key]
    gate_mean = gs['sum'] / gs['count']
    ax.set_title(f"{key} Std | Gate μ={gate_mean:.3f}\nrange=[{gs['min']:.2f}, {gs['max']:.2f}]", fontsize=9)
    ax.axis('off')

# Row 4: Analysis plots
ax = fig.add_subplot(4, 6, 19)
keys = ['S0B0', 'S0B1', 'S1B0', 'S1B1', 'S2B0', 'S2B1']
gate_means_plot = [gate_stats[k]['sum'] / gate_stats[k]['count'] for k in keys]
gate_mins = [gate_stats[k]['min'] for k in keys]
gate_maxs = [gate_stats[k]['max'] for k in keys]
x = np.arange(len(keys))
ax.bar(x, gate_means_plot, color='steelblue', alpha=0.7)
ax.errorbar(x, gate_means_plot, yerr=[np.array(gate_means_plot) - gate_mins, np.array(gate_maxs) - gate_means_plot],
            fmt='none', color='black', capsize=3)
ax.set_xticks(x)
ax.set_xticklabels(keys, fontsize=9)
ax.set_ylabel('Gate Activation')
ax.set_title('Gate Means (± range)', fontsize=10)
ax.set_ylim(0, 1)

# Lens parameters
ax = fig.add_subplot(4, 6, 20)
lens_data = []
for si, stage in enumerate(model.stages):
    for bi, block in enumerate(stage):
        L = block.lens
        lens_data.append({
            'name': f"S{si}B{bi}",
            'omega': L.omega.item(),
            'alpha': L.alpha.item(),
            'xor': torch.sigmoid(L.xor_weight).item()
        })
omegas = [d['omega'] for d in lens_data]
alphas = [d['alpha'] for d in lens_data]
xors = [d['xor'] for d in lens_data]
x = np.arange(len(lens_data))
width = 0.25
ax.bar(x - width, omegas, width, label='ω', alpha=0.7)
ax.bar(x, alphas, width, label='α', alpha=0.7)
ax.bar(x + width, xors, width, label='xor', alpha=0.7)
ax.set_xticks(x)
ax.set_xticklabels([d['name'] for d in lens_data], fontsize=9)
ax.legend(fontsize=8)
ax.set_title('Lens Parameters', fontsize=10)

# Distribution flow
ax = fig.add_subplot(4, 6, 21)
flow_labels = ['Input', 'Stem'] + keys + ['Final']
flow_stds = [input_std.std(), stem_std.std()] + [stage_stds[k].std() for k in keys] + [final_std.std()]
ax.plot(flow_labels, flow_stds, 'o-', color='purple', linewidth=2, markersize=8)
ax.set_ylabel('Std of Std')
ax.set_title('Activation Variance Flow', fontsize=10)
ax.tick_params(axis='x', rotation=45)

# Per-class gate variance (are gates class-specific?)
ax = fig.add_subplot(4, 6, 22)
class_variance = []
for key in keys:
    per_class = class_gate_means[key]
    class_means = [np.mean(v) for v in per_class.values() if len(v) > 0]
    class_variance.append(np.std(class_means) if len(class_means) > 1 else 0)
ax.bar(keys, class_variance, color='coral', alpha=0.7)
ax.set_ylabel('Std of Class Gate Means')
ax.set_title('Gate Class-Specificity', fontsize=10)
ax.tick_params(axis='x', rotation=45)

# Accuracy bar
ax = fig.add_subplot(4, 6, 23)
ax.bar(['Accuracy'], [total_correct / n_samples], color='green', alpha=0.7)
ax.set_ylim(0, 1)
ax.set_title(f'Overall: {total_correct / n_samples:.1%}\n({total_correct}/{n_samples})', fontsize=10)

# Summary text
ax = fig.add_subplot(4, 6, 24)
summary = f"""AGGREGATE STATS (n={n_samples})

Input:  μ={input_mean.mean():.6f}, σ={input_std.mean():.6f}
Stem:   μ={stem_mean.mean():.6f}, σ={stem_std.mean():.4f}
Final:  μ={final_mean.mean():.4f}, σ={final_std.mean():.4f}

Gate Progression:
  S0: {gate_means_plot[0]:.3f} → {gate_means_plot[1]:.3f}
  S1: {gate_means_plot[2]:.3f} → {gate_means_plot[3]:.3f}
  S2: {gate_means_plot[4]:.3f} → {gate_means_plot[5]:.3f}

Accuracy: {total_correct / n_samples:.2%}
"""
ax.text(0.05, 0.95, summary, transform=ax.transAxes, fontsize=9,
        va='top', family='monospace')
ax.set_title('Summary', fontsize=10)
ax.axis('off')

plt.suptitle(f'MobiusNet Aggregate Analysis: {n_samples} Samples from CLIP-ViT-L14 Features', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.savefig("mobiusnet_aggregate_1024.png", dpi=150, bbox_inches="tight")
plt.show()

# ============================================================================
# PRINT DETAILED STATS
# ============================================================================

print(f"\n{'='*70}")
print(f"AGGREGATE STATISTICS ({n_samples} samples)")
print(f"{'='*70}")

print(f"\nActivation Flow:")
print(f"  Input:  μ={input_mean.mean():.6f}, σ={input_std.mean():.6f}")
print(f"  Stem:   μ={stem_mean.mean():.6f}, σ={stem_std.mean():.6f}")
for key in keys:
    m, s = stage_means[key], stage_stds[key]
    print(f"  {key}:   μ={m.mean():.6f}, σ={s.mean():.6f}")
print(f"  Final:  μ={final_mean.mean():.6f}, σ={final_std.mean():.6f}")

print(f"\nGate Statistics:")
for key in keys:
    gs = gate_stats[key]
    print(f"  {key}: μ={gs['sum']/gs['count']:.4f}, range=[{gs['min']:.3f}, {gs['max']:.3f}]")

print(f"\nClass-Specificity (std of per-class gate means):")
for i, key in enumerate(keys):
    print(f"  {key}: {class_variance[i]:.6f}")