analysis/scripts/generate_figures.py

"""
Generate all figures for the ASCAD MTL Preliminary Research Report.
All data is sourced from actual experiment logs, HuggingFace results, and model analysis.
"""
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import json
import os

out_dir = '/home/ubuntu/figures'
os.makedirs(out_dir, exist_ok=True)

plt.rcParams.update({
    'font.size': 11,
    'font.family': 'serif',
    'axes.labelsize': 12,
    'axes.titlesize': 13,
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'legend.fontsize': 10,
    'figure.dpi': 300,
})

# =============================================================================
# Figure 1: HPS Per-Byte Results (from WandB run keu5xeqn + HF results.json)
# =============================================================================
bytes_idx = list(range(16))
hps_ranks = [0, 0, 2, 66, 57, 86, 27, 19, 41, 91, 129, 14, 39, 88, 109, 28]
hps_acc = [24.03, 19.07, 1.86, 1.82, 1.84, 1.92, 1.87, 1.92, 1.90, 1.89, 1.87, 1.88, 1.89, 1.93, 1.75, 1.88]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))

colors_rank = ['#2ecc71' if r == 0 else '#e74c3c' for r in hps_ranks]
bars1 = ax1.bar(bytes_idx, hps_ranks, color=colors_rank, edgecolor='black', linewidth=0.5)
ax1.set_xlabel('Key Byte Index')
ax1.set_ylabel('Final Rank')
ax1.set_title('(a) HPS Final Rank per Byte')
ax1.set_xticks(bytes_idx)
ax1.axhline(y=0, color='green', linestyle='--', alpha=0.5, label='Rank 0 (success)')
ax1.legend()
for i, v in enumerate(hps_ranks):
    if v > 0:
        ax1.text(i, v + 2, str(v), ha='center', va='bottom', fontsize=7)

colors_acc = ['#2ecc71' if a > 5 else '#e74c3c' for a in hps_acc]
bars2 = ax2.bar(bytes_idx, hps_acc, color=colors_acc, edgecolor='black', linewidth=0.5)
ax2.set_xlabel('Key Byte Index')
ax2.set_ylabel('Training Accuracy (%)')
ax2.set_title('(b) HPS Training Accuracy per Byte')
ax2.set_xticks(bytes_idx)
ax2.axhline(y=1/256*100, color='gray', linestyle=':', alpha=0.7, label=f'Random guessing ({1/256*100:.2f}%)')
ax2.legend()
for i, v in enumerate(hps_acc):
    ax2.text(i, v + 0.3, f'{v:.1f}', ha='center', va='bottom', fontsize=6)

plt.tight_layout()
plt.savefig(f'{out_dir}/fig1_hps_baseline.png', bbox_inches='tight')
plt.close()

# =============================================================================
# Figure 2: GAP Weight Analysis (from gradient_analysis_v4.py results)
# =============================================================================
# Per-byte weight norms from the trained HPS model
weight_norms = [65.03, 61.06, 32.45, 32.28, 32.33, 32.16, 32.38, 32.41, 32.29, 32.35, 32.22, 32.40, 32.37, 32.19, 32.31, 32.44]

fig, ax = plt.subplots(figsize=(10, 5))
colors_w = ['#2ecc71' if i < 2 else '#e74c3c' for i in range(16)]
bars = ax.bar(bytes_idx, weight_norms, color=colors_w, edgecolor='black', linewidth=0.5)
ax.set_xlabel('Key Byte Index')
ax.set_ylabel('Weight Matrix Frobenius Norm')
ax.set_title('HPS Per-Byte Head Weight Norms (Trained Model)')
ax.set_xticks(bytes_idx)

# Add annotations
ax.axhline(y=np.mean(weight_norms[2:]), color='red', linestyle='--', alpha=0.5, 
           label=f'Failed bytes mean: {np.mean(weight_norms[2:]):.2f}')
ax.axhline(y=np.mean(weight_norms[:2]), color='green', linestyle='--', alpha=0.5,
           label=f'Succeeded bytes mean: {np.mean(weight_norms[:2]):.2f}')
ax.legend()

for i, v in enumerate(weight_norms):
    ax.text(i, v + 0.5, f'{v:.1f}', ha='center', va='bottom', fontsize=7)

# Add ratio annotation
ax.annotate(f'2.0x gap', xy=(1, 61.06), xytext=(4, 55),
            arrowprops=dict(arrowstyle='->', color='black'),
            fontsize=10, fontweight='bold')

plt.tight_layout()
plt.savefig(f'{out_dir}/fig2_gap_weight_analysis.png', bbox_inches='tight')
plt.close()

# =============================================================================
# Figure 3: GradNorm Weight Explosion (from gradnorm_epochs.txt raw logs)
# =============================================================================
# Exact data from GradNorm log lines
gn_epochs = [1, 2, 3, 4, 9, 14, 19]
gn_min_w = [0.982, 0.970, 0.953, 0.942, 0.887, 0.749, 0.184]
gn_max_w = [1.024, 1.039, 1.053, 1.055, 1.168, 1.943, 10.850]
gn_min_byte = ['byte 8', 'byte 8', 'byte 1', 'byte 14', 'byte 12', 'byte 12', 'byte 5']
gn_max_byte = ['byte 10', 'byte 13', 'byte 13', 'byte 13', 'byte 7', 'byte 0', 'byte 1']

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))

# Left: min/max weight over epochs
ax1.fill_between(gn_epochs, gn_min_w, gn_max_w, alpha=0.3, color='#3498db')
ax1.plot(gn_epochs, gn_max_w, 'o-', color='#e74c3c', linewidth=2, markersize=6, label='Max weight')
ax1.plot(gn_epochs, gn_min_w, 's-', color='#2ecc71', linewidth=2, markersize=6, label='Min weight')
ax1.axhline(y=1.0, color='gray', linestyle=':', alpha=0.5, label='Initial weight (1.0)')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('GradNorm Task Weight')
ax1.set_title('(a) GradNorm Weight Range Over Training')
ax1.legend()
ax1.set_xticks(gn_epochs)

# Annotate key points
ax1.annotate(f'10.850\n({gn_max_byte[-1]})', xy=(19, 10.850), xytext=(15, 9),
            arrowprops=dict(arrowstyle='->', color='red'),
            fontsize=9, color='red', fontweight='bold')
ax1.annotate(f'0.184\n({gn_min_byte[-1]})', xy=(19, 0.184), xytext=(15, 2),
            arrowprops=dict(arrowstyle='->', color='green'),
            fontsize=9, color='green', fontweight='bold')

# Right: weight range (max - min)
gn_range = [mx - mn for mx, mn in zip(gn_max_w, gn_min_w)]
ax2.bar(range(len(gn_epochs)), gn_range, tick_label=[str(e) for e in gn_epochs],
        color=['#2ecc71' if r < 1 else '#f39c12' if r < 5 else '#e74c3c' for r in gn_range],
        edgecolor='black', linewidth=0.5)
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Weight Range (max - min)')
ax2.set_title('(b) GradNorm Weight Disparity Growth')
for i, v in enumerate(gn_range):
    ax2.text(i, v + 0.1, f'{v:.3f}', ha='center', va='bottom', fontsize=8)

plt.tight_layout()
plt.savefig(f'{out_dir}/fig3_gradnorm_explosion.png', bbox_inches='tight')
plt.close()

# =============================================================================
# Figure 4: GradNorm Per-Byte Accuracy at Epoch 19 (from raw log)
# =============================================================================
# Epoch 19 per-byte accuracies from the raw log
epoch19_acc = {
    0: 0.0645, 1: 0.0689, 2: 0.0052, 3: 0.0060, 4: 0.0051,
    5: 0.0045, 6: 0.0049, 7: 0.0050, 8: 0.0049, 9: 0.0062,
    10: 0.0054, 11: 0.0046, 12: 0.0056, 13: 0.0057, 14: 0.0047, 15: 0.0055
}
epoch19_loss = {
    0: 4.2776, 1: 4.2300, 2: 5.5398, 3: 5.5398, 4: 5.5400,
    5: 5.5401, 6: 5.5406, 7: 5.5406, 8: 5.5411, 9: 5.5403,
    10: 5.5400, 11: 5.5408, 12: 5.5402, 13: 5.5405, 14: 5.5403, 15: 5.5415
}

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))

accs = [epoch19_acc[i]*100 for i in range(16)]
losses = [epoch19_loss[i] for i in range(16)]
colors_a = ['#2ecc71' if i < 2 else '#e74c3c' for i in range(16)]

ax1.bar(bytes_idx, accs, color=colors_a, edgecolor='black', linewidth=0.5)
ax1.set_xlabel('Key Byte Index')
ax1.set_ylabel('Training Accuracy (%)')
ax1.set_title('(a) Per-Byte Accuracy at Epoch 19 (GradNorm)')
ax1.set_xticks(bytes_idx)
ax1.axhline(y=1/256*100, color='gray', linestyle=':', alpha=0.7, label='Random guessing')
ax1.legend()
for i, v in enumerate(accs):
    ax1.text(i, v + 0.05, f'{v:.2f}', ha='center', va='bottom', fontsize=6, rotation=45)

ax2.bar(bytes_idx, losses, color=colors_a, edgecolor='black', linewidth=0.5)
ax2.set_xlabel('Key Byte Index')
ax2.set_ylabel('Training Loss')
ax2.set_title('(b) Per-Byte Loss at Epoch 19 (GradNorm)')
ax2.set_xticks(bytes_idx)
ax2.axhline(y=np.log(256), color='gray', linestyle=':', alpha=0.7, label='ln(256) = 5.545')
ax2.legend()

plt.tight_layout()
plt.savefig(f'{out_dir}/fig4_gradnorm_perbyte_epoch19.png', bbox_inches='tight')
plt.close()

# =============================================================================
# Figure 5: Trace Efficiency Curves (from trace_efficiency_results_v2.json)
# =============================================================================
with open('/home/ubuntu/trace_efficiency_results_v2.json') as f:
    te_data = json.load(f)

# Extract num_rank0 for each desync and trace count
trace_counts = [50, 100, 200, 500, 1000, 2000, 5000, 10000]
desync_labels = ['desync_0', 'desync_50', 'desync_100']
desync_display = ['Desync 0', 'Desync 50', 'Desync 100']
colors_d = ['#2ecc71', '#3498db', '#e74c3c']
markers = ['o', 's', '^']

fig, ax = plt.subplots(figsize=(9, 5))
for idx, (dk, dl) in enumerate(zip(desync_labels, desync_display)):
    rank0_counts = []
    for tc in trace_counts:
        tc_str = str(tc)
        if tc_str in te_data.get(dk, {}):
            rank0_counts.append(te_data[dk][tc_str]['num_rank0'])
        else:
            rank0_counts.append(0)
    ax.plot(trace_counts, rank0_counts, f'{markers[idx]}-', color=colors_d[idx],
            linewidth=2, markersize=8, label=dl)

ax.set_xlabel('Number of Attack Traces')
ax.set_ylabel('Bytes at Rank 0 (out of 16)')
ax.set_title('V7b Trace Efficiency Across Desynchronization Levels')
ax.set_xticks(trace_counts)
ax.set_yticks(range(0, 17, 2))
ax.axhline(y=16, color='green', linestyle='--', alpha=0.3, label='Full recovery (16/16)')
ax.legend()
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig(f'{out_dir}/fig5_trace_efficiency.png', bbox_inches='tight')
plt.close()

# =============================================================================
# Figure 6: Architecture Evolution Summary
# =============================================================================
models = ['HPS\n(Baseline)', 'MTAN-Lite', 'LMIC\n(v1-v3)', 'LMIC-TSBN\n(V5e)', 'LMIC-TSBN\n(V7b)']
rank0_counts = [2, 0, 0, 14, 16]
params = [5250688, None, None, None, 1023232]  # Only HPS and V7b known
colors_m = ['#e74c3c', '#e74c3c', '#e74c3c', '#f39c12', '#2ecc71']

fig, ax = plt.subplots(figsize=(10, 5))
bars = ax.bar(range(len(models)), rank0_counts, color=colors_m, edgecolor='black', linewidth=0.5)
ax.set_xlabel('Architecture')
ax.set_ylabel('Bytes at Rank 0 (out of 16)')
ax.set_title('Architecture Evolution: Key Recovery Progress')
ax.set_xticks(range(len(models)))
ax.set_xticklabels(models)
ax.axhline(y=16, color='green', linestyle='--', alpha=0.3)
ax.set_ylim(0, 18)

for i, v in enumerate(rank0_counts):
    ax.text(i, v + 0.3, f'{v}/16', ha='center', va='bottom', fontsize=11, fontweight='bold')

# Add annotations for key changes
ax.annotate('+ Localized\nwindows', xy=(2, 0.5), fontsize=8, ha='center', color='gray')
ax.annotate('+ TSBN\n+ DTP', xy=(3, 14.5), fontsize=8, ha='center', color='gray')
ax.annotate('+ Multi-bit\n(8 binary)', xy=(4, 16.5), fontsize=8, ha='center', color='gray')

plt.tight_layout()
plt.savefig(f'{out_dir}/fig6_architecture_evolution.png', bbox_inches='tight')
plt.close()

print("All 6 figures generated successfully:")
for f in sorted(os.listdir(out_dir)):
    print(f"  {out_dir}/{f}")