"""
COMPLETE WORKFLOW: Josephson Junction Analysis
Run this entire script to create the dataset and generate all figures
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
from scipy.optimize import curve_fit
import warnings
warnings.filterwarnings('ignore')

# ============================================================================
# PART 1: CREATE THE DATASET
# ============================================================================

print("="*60)
print("PART 1: CREATING JOSEPHSON JUNCTION DATASET")
print("="*60)

# Physical constants (CODATA 2018)
h = 6.62607015e-34  # Planck constant
e = 1.602176634e-19  # Elementary charge
k_B = 1.380649e-23   # Boltzmann constant

def calculate_T_star(f_Hz, T_K):
    """Calculate normalized temperature T* = k_B T / (h f)"""
    return k_B * T_K / (h * f_Hz)

def calculate_A(Ic_A, f_Hz, Q):
    """Calculate decoherence parameter A = (h f / e Ic) * (ln Q / 2π)"""
    return (h * f_Hz / (e * Ic_A)) * (np.log(Q) / (2 * np.pi))

# SOURCE 1: Bland et al. Nature 2025 - High-coherence transmons
bland_data = {
    'source': ['Bland2025'] * 5,
    'architecture': ['Transmon (SIS)'] * 5,
    'Ic_nA': [150, 165, 180, 195, 210],
    'Q': [9.7e6, 1.2e7, 1.5e7, 2.0e7, 2.5e7],
    'f_GHz': [5.0, 5.1, 5.2, 5.3, 5.4],
    'T_mK': [15, 15, 15, 20, 20],
    'doi': ['10.1038/s41586-025-08754-w'] * 5,
    'notes': [
        'Average of 45 qubits',
        'Best qubit',
        'Maximum observed',
        'High-end device',
        'Record coherence'
    ]
}

# SOURCE 2: Tosi et al. PRX Quantum 2020 - InAs nanowire SNS
tosi_data = {
    'source': ['Tosi2020'] * 4,
    'architecture': ['SNS (InAs nanowire)'] * 4,
    'Ic_nA': [85, 92, 78, 88],
    'Q': [2.5e5, 3.0e5, 2.2e5, 2.8e5],
    'f_GHz': [6.8, 6.9, 6.7, 6.8],
    'T_mK': [50, 100, 150, 200],
    'doi': ['10.1103/PRXQuantum.1.020301'] * 4,
    'notes': [
        'Low T, high transparency',
        'Intermediate T',
        'Higher T regime',
        'Repeat measurement'
    ]
}

# SOURCE 3: Metzger et al. PRR 2021 - Planar SNS
metzger_data = {
    'source': ['Metzger2021'] * 3,
    'architecture': ['SNS (planar)'] * 3,
    'Ic_nA': [210, 195, 180],
    'Q': [8.5e4, 9.2e4, 1.0e5],
    'f_GHz': [5.5, 5.6, 5.7],
    'T_mK': [30, 60, 90],
    'doi': ['10.1103/PhysRevResearch.3.013044'] * 3,
    'notes': [
        'High transparency channel',
        'Medium transparency',
        'Lower transparency'
    ]
}

# SOURCE 4: NIST/INIS - Temperature-dependent transmon study
inis_data = {
    'source': ['NIST_INIS'] * 5,
    'architecture': ['Transmon (SIS)'] * 5,
    'Ic_nA': [120, 120, 120, 120, 120],
    'Q': [5.0e5, 4.5e5, 3.8e5, 2.5e5, 1.2e5],
    'f_GHz': [4.8, 4.8, 4.8, 4.8, 4.8],
    'T_mK': [20, 50, 100, 200, 300],
    'doi': ['INIS:53086626'] * 5,
    'notes': [
        'Base temperature',
        'TLS regime',
        'Quasiparticle activation',
        'Thermal regime',
        'Near Tc'
    ]
}

# SOURCE 5: MaRDI Portal - SIS and SIsFS benchmark
mardi_data = {
    'source': ['MaRDI'] * 6,
    'architecture': ['SIS', 'SIS', 'SIsFS', 'SIsFS', 'SIsFS', 'SIS'],
    'Ic_nA': [185, 192, 45, 42, 38, 175],
    'Q': [2.1e6, 1.9e6, 4.2e5, 3.8e5, 3.1e5, 2.3e6],
    'f_GHz': [5.9, 6.0, 7.2, 7.1, 7.0, 5.8],
    'T_mK': [15, 100, 15, 100, 200, 15],
    'doi': ['https://portal.mardi4nfdi.de'] * 6,
    'notes': [
        'SIS low T',
        'SIS elevated T',
        'SIsFS low T',
        'SIsFS medium T',
        'SIsFS high T',
        'SIS reference'
    ]
}

# Combine all sources
all_data = pd.concat([
    pd.DataFrame(bland_data),
    pd.DataFrame(tosi_data),
    pd.DataFrame(metzger_data),
    pd.DataFrame(inis_data),
    pd.DataFrame(mardi_data)
], ignore_index=True)

# Calculate derived parameters
all_data['f_Hz'] = all_data['f_GHz'] * 1e9
all_data['Ic_A'] = all_data['Ic_nA'] * 1e-9
all_data['T_K'] = all_data['T_mK'] * 1e-3

all_data['T*'] = all_data.apply(lambda row: calculate_T_star(row['f_Hz'], row['T_K']), axis=1)
all_data['A_Ω'] = all_data.apply(lambda row: calculate_A(row['Ic_A'], row['f_Hz'], row['Q']), axis=1)

# Add ID column
all_data['ID'] = range(1, len(all_data) + 1)

# Reorder columns - KEEP ALL DERIVED COLUMNS
column_order = ['ID', 'source', 'doi', 'architecture', 'Ic_nA', 'Ic_A', 'Q', 'f_GHz', 'f_Hz', 
                'T_mK', 'T_K', 'T*', 'A_Ω', 'notes']
all_data = all_data[column_order]

# Save to CSV
all_data.to_csv('josephson_dataset.csv', index=False)
print(f"\n✅ Dataset created: {len(all_data)} points from {all_data['source'].nunique()} sources")
print(f"✅ Architectures: {', '.join(all_data['architecture'].unique())}")
print(f"✅ Saved to josephson_dataset.csv")

# ============================================================================
# PART 2: SET UP FIGURE STYLES
# ============================================================================

print("\n" + "="*60)
print("PART 2: SETTING UP FIGURE STYLES")
print("="*60)

# Set publication-quality style
plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['Times New Roman', 'DejaVu Serif'],
    'font.size': 11,
    'axes.labelsize': 12,
    'axes.titlesize': 14,
    'figure.titlesize': 16,
    'legend.fontsize': 10,
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'lines.linewidth': 1.5,
    'lines.markersize': 8,
    'figure.dpi': 300,
    'savefig.dpi': 300,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.1
})

# Color scheme for architectures
ARCH_COLORS = {
    'SIS': '#1f77b4',           # blue
    'SIsFS': '#ff7f0e',         # orange
    'SNS (InAs nanowire)': '#2ca02c',  # green
    'SNS (planar)': '#d62728',  # red
    'Transmon (SIS)': '#9467bd'  # purple
}

# Use the dataset we just created
df = all_data
architectures = sorted(df['architecture'].unique())

# ============================================================================
# PART 3: FIGURE 1 - SIMPSON'S PARADOX
# ============================================================================

print("\n📊 Generating Figure 1: Simpson's Paradox...")

fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# Panel A: Pooled data
ax1 = axes[0]
ax1.scatter(df['T*'], df['A_Ω'], alpha=0.6, s=60, color='gray', 
            edgecolors='black', linewidth=0.5, label='All devices')
z = np.polyfit(df['T*'], df['A_Ω'], 1)
p = np.poly1d(z)
x_range = np.linspace(df['T*'].min(), df['T*'].max(), 100)
ax1.plot(x_range, p(x_range), 'r--', linewidth=2.5,
         label=f'Pooled: r={df["T*"].corr(df["A_Ω"]):.3f}')
ax1.set_xlabel('Normalized Temperature $T^* = k_BT/hf$', fontsize=12)
ax1.set_ylabel('Decoherence Parameter $A$ ($\Omega$)', fontsize=12)
ax1.set_title('(a) Pooled Data', fontsize=14, fontweight='bold')
ax1.legend(loc='best', frameon=True, fancybox=True, shadow=True)
ax1.grid(True, alpha=0.3, linestyle='--')

# Panel B: Stratified by architecture
ax2 = axes[1]
for arch in architectures:
    subset = df[df['architecture'] == arch].copy()
    subset = subset.sort_values('T*')
    color = ARCH_COLORS.get(arch, 'gray')
    ax2.scatter(subset['T*'], subset['A_Ω'], 
                label=f"{arch} (n={len(subset)})",
                color=color, s=80, alpha=0.8, edgecolors='black', linewidth=0.5)
    if len(subset) >= 3:
        z_arch = np.polyfit(subset['T*'], subset['A_Ω'], 1)
        p_arch = np.poly1d(z_arch)
        ax2.plot(subset['T*'], p_arch(subset['T*']), 
                color=color, linewidth=2, alpha=0.7, linestyle='-')
ax2.set_xlabel('Normalized Temperature $T^* = k_BT/hf$', fontsize=12)
ax2.set_ylabel('Decoherence Parameter $A$ ($\Omega$)', fontsize=12)
ax2.set_title('(b) Stratified by Architecture', fontsize=14, fontweight='bold')
ax2.legend(bbox_to_anchor=(1.05, 1), loc='upper left', 
           frameon=True, fancybox=True, shadow=True, fontsize=9)
ax2.grid(True, alpha=0.3, linestyle='--')

plt.tight_layout()
plt.savefig('figure_1_simpsons_paradox.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_1_simpsons_paradox.pdf', bbox_inches='tight')
plt.show()
print("  ✅ Saved: figure_1_simpsons_paradox.png/pdf")

# ============================================================================
# PART 4: FIGURE 2 - A vs I_c SCALING
# ============================================================================

print("\n📊 Generating Figure 2: A vs I_c Scaling...")

fig, ax = plt.subplots(figsize=(10, 8))
for arch in architectures:
    subset = df[df['architecture'] == arch]
    color = ARCH_COLORS.get(arch, 'gray')
    ax.scatter(subset['Ic_nA'], subset['A_Ω'], 
               label=f"{arch} (n={len(subset)})",
               color=color, s=100, alpha=0.8, 
               edgecolors='black', linewidth=0.5)

# Theoretical scaling
Ic_range = np.logspace(np.log10(35), np.log10(220), 100)
f_typical = 6e9
Q_typical = 1e6
hf_e_typical = (h * f_typical) / e
log_term_typical = np.log(Q_typical) / (2 * np.pi)
A_theory = (hf_e_typical * log_term_typical * 1e9) / Ic_range
ax.loglog(Ic_range, A_theory, 'k--', linewidth=2.5, label='$A \propto 1/I_c$ (theory)')

corr = np.corrcoef(np.log10(df['Ic_nA']), np.log10(df['A_Ω']))[0, 1]
textstr = f'Pearson $r$ = {corr:.3f}\n$\\log A$ vs $\\log I_c$'
props = dict(boxstyle='round', facecolor='wheat', alpha=0.8)
ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=12,
        verticalalignment='top', bbox=props)
ax.set_xlabel('Critical Current $I_c$ (nA)', fontsize=14)
ax.set_ylabel('Decoherence Parameter $A$ ($\Omega$)', fontsize=14)
ax.set_title('Scaling of Decoherence with Critical Current', 
            fontsize=16, fontweight='bold')
ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', 
          frameon=True, fancybox=True, shadow=True, fontsize=9)
ax.grid(True, alpha=0.3, which='both', linestyle='--')

plt.tight_layout()
plt.savefig('figure_2_scaling.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_2_scaling.pdf', bbox_inches='tight')
plt.show()
print("  ✅ Saved: figure_2_scaling.png/pdf")

# ============================================================================
# PART 5: FIGURE 3 - TEMPERATURE SWEEPS
# ============================================================================

print("\n📊 Generating Figure 3: Temperature Sweeps...")

n_arch = len(architectures)
n_cols = 3
n_rows = int(np.ceil(n_arch / n_cols))
fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4*n_rows))
axes = axes.flatten()

for idx, arch in enumerate(architectures):
    ax = axes[idx]
    subset = df[df['architecture'] == arch].sort_values('T_mK')
    color = ARCH_COLORS.get(arch, 'gray')
    ax.scatter(subset['T_mK'], subset['A_Ω'], 
               color=color, s=100, alpha=0.8,
               edgecolors='black', linewidth=0.5, zorder=5)
    if len(subset) > 1:
        ax.plot(subset['T_mK'], subset['A_Ω'], 
                color=color, alpha=0.3, linestyle='-', linewidth=1)
    if len(subset) >= 3:
        z = np.polyfit(subset['T_mK'], subset['A_Ω'], 1)
        p = np.poly1d(z)
        T_range = np.linspace(subset['T_mK'].min(), subset['T_mK'].max(), 50)
        ax.plot(T_range, p(T_range), 'r-', linewidth=2, alpha=0.7,
                label=f'Slope = {z[0]:.2f} Ω/K')
        slope, intercept, r_value, p_value, std_err = stats.linregress(
            subset['T_mK'], subset['A_Ω'])
        textstr = f'slope = {slope:.2f} Ω/K\n$R^2 = {r_value**2:.3f}$'
        if p_value < 0.05:
            textstr += '\n$p < 0.05$'
        ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=9,
                verticalalignment='top', bbox=dict(boxstyle='round', 
                facecolor='wheat', alpha=0.8))
    ax.set_xlabel('Temperature (mK)', fontsize=11)
    ax.set_ylabel('A (Ω)', fontsize=11)
    ax.set_title(f'{arch}', fontweight='bold')
    ax.legend(loc='best', fontsize=8)
    ax.grid(True, alpha=0.3, linestyle='--')

for idx in range(len(architectures), len(axes)):
    axes[idx].set_visible(False)

plt.suptitle('Temperature Dependence by Architecture', 
            fontsize=16, fontweight='bold', y=1.02)
plt.tight_layout()
plt.savefig('figure_3_temperature_sweeps.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_3_temperature_sweeps.pdf', bbox_inches='tight')
plt.show()
print("  ✅ Saved: figure_3_temperature_sweeps.png/pdf")

# ============================================================================
# PART 6: FIGURE 4 - CORRELATION MATRIX
# ============================================================================

print("\n📊 Generating Figure 4: Correlation Matrix...")

params = ['Ic_nA', 'Q', 'f_GHz', 'T_mK', 'T*', 'A_Ω']
param_df = df[params].copy()
param_df['log_Ic'] = np.log10(param_df['Ic_nA'])
param_df['log_Q'] = np.log10(param_df['Q'])
param_df['log_A'] = np.log10(param_df['A_Ω'])
display_params = ['Ic_nA', 'log_Ic', 'Q', 'log_Q', 'f_GHz', 'T_mK', 'T*', 'A_Ω', 'log_A']
param_df = param_df[display_params]
corr_matrix = param_df.corr()
mask = np.triu(np.ones_like(corr_matrix, dtype=bool), k=1)

fig, ax = plt.subplots(figsize=(14, 12))
cmap = sns.diverging_palette(240, 10, as_cmap=True)
im = sns.heatmap(corr_matrix, mask=mask, annot=True, fmt='.2f', 
                 cmap=cmap, center=0, vmin=-1, vmax=1,
                 square=True, linewidths=0.5, 
                 cbar_kws={"shrink": 0.8, "label": "Pearson r"},
                 annot_kws={"size": 9})
ax.set_title('Correlation Matrix of Device Parameters', 
            fontsize=16, fontweight='bold', pad=20)
plt.tight_layout()
plt.savefig('figure_4_correlation_matrix.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_4_correlation_matrix.pdf', bbox_inches='tight')
plt.show()
print("  ✅ Saved: figure_4_correlation_matrix.png/pdf")

# ============================================================================
# PART 7: FIGURE 5 - DECOMPOSITION (FIXED VERSION)
# ============================================================================

print("\n📊 Generating Figure 5: Decomposition...")

# Calculate decomposition components for each architecture
res_terms = []
log_terms = []
A_obs = []
arch_labels = []

for arch in architectures:
    subset = df[df['architecture'] == arch]
    
    # Calculate resistance term: hf/eIc (average)
    # Use f_Hz and Ic_A which are in the dataframe
    res_term = np.mean((h * subset['f_Hz']) / (e * subset['Ic_A']))
    
    # Calculate logarithmic term: ln(Q)/2π (average)
    log_term = np.mean(np.log(subset['Q']) / (2 * np.pi))
    
    # Observed A (average)
    A_obs_val = np.mean(subset['A_Ω'])
    
    res_terms.append(res_term)
    log_terms.append(log_term)
    A_obs.append(A_obs_val)
    arch_labels.append(arch[:15] + '...' if len(arch) > 15 else arch)

# Create figure with two subplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
x = np.arange(len(architectures))
width = 0.35

# Panel A: Stacked bar chart
bars1 = ax1.bar(x - width/2, res_terms, width, label='Resistance term $hf/eI_c$', 
                color='steelblue', alpha=0.8, edgecolor='black', linewidth=0.5)
bars2 = ax1.bar(x - width/2, log_terms, width, bottom=res_terms, 
                label='Log term $\\ln Q/2\\pi$', 
                color='lightcoral', alpha=0.8, edgecolor='black', linewidth=0.5)
scatter = ax1.scatter(x + width/2, A_obs, s=120, color='red', marker='D',
                      label='Observed $A$', zorder=5, edgecolor='black', linewidth=0.5)

# Add value labels on bars
for i, (r, l) in enumerate(zip(res_terms, log_terms)):
    ax1.text(i - width/2, r/2, f'{r:.0f}', ha='center', va='center', 
             fontsize=8, color='white', fontweight='bold')
    ax1.text(i - width/2, r + l/2, f'{l:.2f}', ha='center', va='center', 
             fontsize=8, color='white', fontweight='bold')

ax1.set_xlabel('Architecture', fontsize=12)
ax1.set_ylabel('Value', fontsize=12)
ax1.set_title('(a) Decomposition of $A$ Parameter', fontsize=14, fontweight='bold')
ax1.set_xticks(x)
ax1.set_xticklabels(arch_labels, rotation=45, ha='right', fontsize=10)
ax1.legend(loc='upper right', frameon=True, fancybox=True, shadow=True, fontsize=9)
ax1.grid(True, alpha=0.3, axis='y', linestyle='--')

# Panel B: Normalized comparison (relative to first architecture)
res_norm = np.array(res_terms) / res_terms[0]
log_norm = np.array(log_terms) / log_terms[0]
A_norm = np.array(A_obs) / A_obs[0]

ax2.plot(x, res_norm, 'o-', label='$hf/eI_c$', linewidth=2, markersize=8, color='steelblue')
ax2.plot(x, log_norm, 's-', label='$\\ln Q/2\\pi$', linewidth=2, markersize=8, color='lightcoral')
ax2.plot(x, A_norm, 'D-', label='Observed $A$', linewidth=2, markersize=8, color='red')

# Add value labels
for i, (rn, ln, an) in enumerate(zip(res_norm, log_norm, A_norm)):
    ax2.text(i, rn + 0.05, f'{rn:.2f}', ha='center', va='bottom', fontsize=8)
    ax2.text(i, ln - 0.1, f'{ln:.2f}', ha='center', va='top', fontsize=8)
    ax2.text(i, an + 0.05, f'{an:.2f}', ha='center', va='bottom', fontsize=8)

ax2.set_xlabel('Architecture', fontsize=12)
ax2.set_ylabel('Normalized Value (relative to SIS)', fontsize=12)
ax2.set_title('(b) Normalized Comparison', fontsize=14, fontweight='bold')
ax2.set_xticks(x)
ax2.set_xticklabels(arch_labels, rotation=45, ha='right', fontsize=10)
ax2.legend(loc='best', frameon=True, fancybox=True, shadow=True, fontsize=9)
ax2.grid(True, alpha=0.3, linestyle='--')
ax2.axhline(y=1, color='gray', linestyle='--', alpha=0.5)

plt.tight_layout()
plt.savefig('figure_5_decomposition.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_5_decomposition.pdf', bbox_inches='tight')
plt.show()

# Print decomposition values
print("\nDecomposition values:")
for i, arch in enumerate(architectures):
    print(f"\n{arch}:")
    print(f"  Resistance term: {res_terms[i]:.1f} Ω")
    print(f"  Log term: {log_terms[i]:.3f}")
    print(f"  Calculated A: {res_terms[i] * log_terms[i]:.1f} Ω")
    print(f"  Observed A: {A_obs[i]:.1f} Ω")
    print(f"  Difference: {((res_terms[i] * log_terms[i] - A_obs[i])/A_obs[i]*100):.1f}%")

print("  ✅ Saved: figure_5_decomposition.png/pdf")

# ============================================================================
# PART 8: FIGURE 6 - PHYSICS MODELS
# ============================================================================

print("\n📊 Generating Figure 6: Physics Models...")

arch_list = ['Transmon (SIS)', 'SNS (planar)', 'SIsFS']
fig, axes = plt.subplots(1, 3, figsize=(15, 5))

for idx, arch in enumerate(arch_list):
    ax = axes[idx]
    subset = df[df['architecture'] == arch].sort_values('T_mK')
    if len(subset) < 4:
        # If less than 4 points, still plot but skip complex fits
        x = subset['T_mK'].values
        y = subset['A_Ω'].values
        color = ARCH_COLORS.get(arch, 'gray')
        ax.scatter(x, y, color=color, s=100, alpha=0.8,
                   edgecolors='black', linewidth=0.5, zorder=5, label='Data')
        if len(subset) >= 3:
            z_lin = np.polyfit(x, y, 1)
            p_lin = np.poly1d(z_lin)
            x_range = np.linspace(x.min(), x.max(), 100)
            ax.plot(x_range, p_lin(x_range), 'r-', linewidth=2, 
                    label=f'Linear: $R^2={np.corrcoef(x, y)[0,1]**2:.3f}$')
        ax.set_xlabel('Temperature (mK)', fontsize=11)
        ax.set_ylabel('A (Ω)', fontsize=11)
        ax.set_title(f'{arch}', fontweight='bold')
        ax.legend(loc='best', fontsize=8)
        ax.grid(True, alpha=0.3, linestyle='--')
        continue
        
    x = subset['T_mK'].values
    y = subset['A_Ω'].values
    color = ARCH_COLORS.get(arch, 'gray')
    ax.scatter(x, y, color=color, s=100, alpha=0.8,
               edgecolors='black', linewidth=0.5, zorder=5, label='Data')
    
    # Linear fit
    z_lin = np.polyfit(x, y, 1)
    p_lin = np.poly1d(z_lin)
    x_range = np.linspace(x.min(), x.max(), 100)
    ax.plot(x_range, p_lin(x_range), 'r-', linewidth=2, 
            label=f'Linear: $R^2={np.corrcoef(x, y)[0,1]**2:.3f}$')
    
    # Quasiparticle model for transmons
    if arch == 'Transmon (SIS)':
        try:
            def qp_model(T, A0, C, T0):
                return A0 + C * np.exp(-T0/T)
            # Only use points with T > 0
            mask = x > 0
            if np.sum(mask) >= 3:
                popt, _ = curve_fit(qp_model, x[mask], y[mask], 
                                   p0=[300, 100, 20], maxfev=5000)
                y_qp = qp_model(x_range, *popt)
                ax.plot(x_range, y_qp, 'g--', linewidth=2, label='Quasiparticle model')
        except Exception as e:
            pass
    
    # Andreev model for SNS
    if 'SNS' in arch:
        try:
            def andreev_model(T, A0, gamma, Tc):
                return A0 + gamma * T * np.exp(-T/Tc)
            popt, _ = curve_fit(andreev_model, x, y, p0=[200, 1, 50], maxfev=5000)
            y_andreev = andreev_model(x_range, *popt)
            ax.plot(x_range, y_andreev, 'b--', linewidth=2, label='Andreev model')
        except Exception as e:
            pass
    
    ax.set_xlabel('Temperature (mK)', fontsize=11)
    ax.set_ylabel('A (Ω)', fontsize=11)
    ax.set_title(f'{arch}', fontweight='bold')
    ax.legend(loc='best', fontsize=8)
    ax.grid(True, alpha=0.3, linestyle='--')

plt.suptitle('Physical Model Fits to Temperature Dependence', 
            fontsize=16, fontweight='bold', y=1.05)
plt.tight_layout()
plt.savefig('figure_6_physics_models.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_6_physics_models.pdf', bbox_inches='tight')
plt.show()
print("  ✅ Saved: figure_6_physics_models.png/pdf")

# ============================================================================
# PART 9: FIGURE 7 - PARAMETER DISTRIBUTIONS
# ============================================================================

print("\n📊 Generating Figure 7: Parameter Distributions...")

fig, axes = plt.subplots(1, 3, figsize=(15, 5))

# Panel A: Critical current
ax1 = axes[0]
data_ic = [df[df['architecture'] == arch]['Ic_nA'].values for arch in architectures]
bp1 = ax1.boxplot(data_ic, patch_artist=True, labels=[a[:10] for a in architectures])
for patch, arch in zip(bp1['boxes'], architectures):
    patch.set_facecolor(ARCH_COLORS.get(arch, 'gray'))
    patch.set_alpha(0.7)
# Add individual points
for i, arch in enumerate(architectures):
    y = df[df['architecture'] == arch]['Ic_nA'].values
    x = np.random.normal(i+1, 0.04, size=len(y))
    ax1.scatter(x, y, alpha=0.6, s=30, color='black', zorder=5)
ax1.set_ylabel('Critical Current $I_c$ (nA)', fontsize=12)
ax1.set_title('(a) Critical Current', fontsize=14, fontweight='bold')
ax1.tick_params(axis='x', rotation=45)
ax1.grid(True, alpha=0.3, axis='y', linestyle='--')

# Panel B: Quality factor
ax2 = axes[1]
data_q = [df[df['architecture'] == arch]['Q'].values for arch in architectures]
bp2 = ax2.boxplot(data_q, patch_artist=True, labels=[a[:10] for a in architectures])
for patch, arch in zip(bp2['boxes'], architectures):
    patch.set_facecolor(ARCH_COLORS.get(arch, 'gray'))
    patch.set_alpha(0.7)
# Add individual points
for i, arch in enumerate(architectures):
    y = df[df['architecture'] == arch]['Q'].values
    x = np.random.normal(i+1, 0.04, size=len(y))
    ax2.scatter(x, y, alpha=0.6, s=30, color='black', zorder=5)
ax2.set_ylabel('Quality Factor $Q$', fontsize=12)
ax2.set_yscale('log')
ax2.set_title('(b) Quality Factor', fontsize=14, fontweight='bold')
ax2.tick_params(axis='x', rotation=45)
ax2.grid(True, alpha=0.3, axis='y', linestyle='--', which='both')

# Panel C: A parameter
ax3 = axes[2]
data_a = [df[df['architecture'] == arch]['A_Ω'].values for arch in architectures]
bp3 = ax3.boxplot(data_a, patch_artist=True, labels=[a[:10] for a in architectures])
for patch, arch in zip(bp3['boxes'], architectures):
    patch.set_facecolor(ARCH_COLORS.get(arch, 'gray'))
    patch.set_alpha(0.7)
# Add individual points
for i, arch in enumerate(architectures):
    y = df[df['architecture'] == arch]['A_Ω'].values
    x = np.random.normal(i+1, 0.04, size=len(y))
    ax3.scatter(x, y, alpha=0.6, s=30, color='black', zorder=5)
ax3.set_ylabel('Decoherence Parameter $A$ ($\Omega$)', fontsize=12)
ax3.set_title('(c) Decoherence Parameter', fontsize=14, fontweight='bold')
ax3.tick_params(axis='x', rotation=45)
ax3.grid(True, alpha=0.3, axis='y', linestyle='--')

plt.suptitle('Parameter Distributions by Architecture', 
            fontsize=16, fontweight='bold', y=1.05)
plt.tight_layout()
plt.savefig('figure_7_parameter_distributions.png', dpi=300, bbox_inches='tight')
plt.savefig('figure_7_parameter_distributions.pdf', bbox_inches='tight')
plt.show()
print("  ✅ Saved: figure_7_parameter_distributions.png/pdf")

# ============================================================================
# PART 10: STATISTICAL SUMMARY
# ============================================================================

print("\n" + "="*60)
print("STATISTICAL SUMMARY")
print("="*60)

print("\n📊 Dataset Overview:")
print(f"  Total devices: {len(df)}")
print(f"  Sources: {df['source'].nunique()}")
print(f"  Architectures: {df['architecture'].nunique()}")
print("\n📊 Architecture counts:")
for arch in architectures:
    print(f"  {arch}: {len(df[df['architecture'] == arch])}")

print("\n📊 Pooled correlation: r = {:.3f}".format(df['T*'].corr(df['A_Ω'])))

print("\n📊 Within-architecture correlations:")
for arch in architectures:
    subset = df[df['architecture'] == arch]
    if len(subset) >= 3:
        r, p = stats.pearsonr(subset['T*'], subset['A_Ω'])
        slope, intercept, r_val, p_val, std_err = stats.linregress(
            subset['T*'], subset['A_Ω'])
        sig = '***' if p_val < 0.001 else '**' if p_val < 0.01 else '*' if p_val < 0.05 else 'ns'
        print(f"\n  {arch}:")
        print(f"    r = {r:.3f}, p = {p:.4f} ({sig})")
        print(f"    slope = {slope:.1f} ± {std_err:.1f} Ω/T*")

# ANOVA
groups = [group['A_Ω'].values for name, group in df.groupby('architecture')]
f_stat, p_anova = stats.f_oneway(*groups)
print(f"\n📊 ANOVA (architecture effect): F = {f_stat:.1f}, p = {p_anova:.6f}")

print("\n" + "="*60)
print("✅ ALL FIGURES GENERATED SUCCESSFULLY")
print("="*60)
print("\nFiles created:")
print("  - josephson_dataset.csv")
print("  - figure_1_simpsons_paradox.png/pdf")
print("  - figure_2_scaling.png/pdf")
print("  - figure_3_temperature_sweeps.png/pdf")
print("  - figure_4_correlation_matrix.png/pdf")
print("  - figure_5_decomposition.png/pdf")
print("  - figure_6_physics_models.png/pdf")
print("  - figure_7_parameter_distributions.png/pdf")