# --- Running Experiment 1: Sweeping Noise (Layers=2) ---
# Noise p=0.000 | Mean Grad: 1.481 | Mean QNG Step: 32.78
# Noise p=0.025 | Mean Grad: 0.979 | Mean QNG Step: 50.43
# Noise p=0.050 | Mean Grad: 0.658 | Mean QNG Step: 74.52
# Noise p=0.075 | Mean Grad: 0.429 | Mean QNG Step: 116.95
# Noise p=0.100 | Mean Grad: 0.274 | Mean QNG Step: 182.59

# --- Running Experiment 2: Sweeping Layers (Noise=0.05) ---
# Layers L=1 | Mean Grad: 1.163 | Mean QNG Step: 32.30
# Layers L=2 | Mean Grad: 0.669 | Mean QNG Step: 76.24
# Layers L=3 | Mean Grad: 0.441 | Mean QNG Step: 169.57
# Layers L=4 | Mean Grad: 0.317 | Mean QNG Step: 312.07
# Layers L=5 | Mean Grad: 0.205 | Mean QNG Step: 556.86


import matplotlib.pyplot as plt

# --- 1. Extracted Data ---
noise_levels = [0.000, 0.025, 0.050, 0.075, 0.100]
exp1_grad = [1.481, 0.979, 0.658, 0.429, 0.274]
exp1_qng = [32.78, 50.43, 74.52, 116.95, 182.59]

layer_counts = [1, 2, 3, 4, 5]
exp2_grad = [1.163, 0.669, 0.441, 0.317, 0.205]
exp2_qng = [32.30, 76.24, 169.57, 312.07, 556.86]

# --- 2. Plot Configuration ---
# 14x3.5 creates a very wide, short aspect ratio perfect for saving vertical space
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 2.2))

def style_axes(ax, ax_twin):
    ax.tick_params(axis='both', which='major', labelsize=14)
    ax_twin.tick_params(axis='y', which='major', labelsize=14)
    ax.grid(True, linestyle=':', alpha=0.8, linewidth=1.2, color='gray')

# --- Plot 1: Sweeping Noise ---
ax1.plot(noise_levels, exp1_grad, color='navy', linestyle='-', linewidth=3, marker='o', label='Euclidean Gradient Norm')
ax1.set_xlabel('Depolarizing Noise Probability (p)', fontsize=14)

ax1_twin = ax1.twinx()
ax1_twin.plot(noise_levels, exp1_qng, color='crimson', linestyle='--', linewidth=3, marker='x', markersize=8, label='Riemannian Gradient Norm')

style_axes(ax1, ax1_twin)

lines_1, labels_1 = ax1.get_legend_handles_labels()
lines_2, labels_2 = ax1_twin.get_legend_handles_labels()
ax1_twin.legend(lines_1 + lines_2, labels_1 + labels_2, loc='best', fontsize=13, framealpha=1.0, edgecolor='black')

# --- Plot 2: Sweeping Layers ---
ax2.plot(layer_counts, exp2_grad, color='navy', linestyle='-', linewidth=3, marker='o', label='Euclidean Gradient Norm')
ax2.set_xlabel('Number of Layers (L)', fontsize=14)
ax2.set_xticks(layer_counts)

ax2_twin = ax2.twinx()
ax2_twin.plot(layer_counts, exp2_qng, color='crimson', linestyle='--', linewidth=3, marker='x', markersize=8, label='Riemannian Gradient Norm')

style_axes(ax2, ax2_twin)

lines_3, labels_3 = ax2.get_legend_handles_labels()
lines_4, labels_4 = ax2_twin.get_legend_handles_labels()
# ax2_twin.legend(lines_3 + lines_4, labels_3 + labels_4, loc='best', fontsize=13, framealpha=1.0, edgecolor='black')

# --- Shared y-axis labels in the center ---
fig.text(0.5, 0.5, 'Euclidean Norm', ha='center', va='center',
         fontsize=12, color='navy', rotation=90,
         transform=fig.transFigure)
fig.text(0.515, 0.5, 'Riemannian Norm', ha='center', va='center',
         fontsize=12, color='crimson', rotation=90,
         transform=fig.transFigure)

# --- Render ---
fig.subplots_adjust(wspace=0.45)  # widen gap to fit the shared labels
fig.tight_layout(rect=[0, 0, 1, 1], w_pad=5.0)  # w_pad controls inter-subplot spacing
plt.show()