import pandas as pd, numpy as np
from scipy import stats
from scipy.cluster.hierarchy import linkage, fcluster, dendrogram
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

df = pd.read_csv('table1_data.csv')
tasks = ['Classification','Clustering','PairClassification','Reranking','Retrieval','STS','Summarization']

# 1. Pearson correlation matrix across tasks (using the 30 models as samples)
corr = df[tasks].corr(method='pearson')
corr.to_csv('task_correlation_matrix.csv')
print("=== Pearson correlation matrix (tasks x tasks, n=30 models) ===")
print(corr.round(2))

# 2. Z-score each task column across models, compute per-model profile dispersion (std of z-scores)
z = (df[tasks] - df[tasks].mean()) / df[tasks].std(ddof=0)
df['profile_dispersion'] = z.std(axis=1)
df['z_mean'] = z.mean(axis=1)

# 3. Correlate dispersion with Average score and with log(params)
df['log_params'] = np.log10(df['Params_M'])
r_disp_avg, p_disp_avg = stats.pearsonr(df['profile_dispersion'], df['Average'])
r_disp_size, p_disp_size = stats.pearsonr(df['profile_dispersion'], df['log_params'])
r_avg_size, p_avg_size = stats.pearsonr(df['Average'], df['log_params'])
print(f"\nCorrelation dispersion vs Average score: r={r_disp_avg:.3f}, p={p_disp_avg:.4f}")
print(f"Correlation dispersion vs log(params): r={r_disp_size:.3f}, p={p_disp_size:.4f}")
print(f"Correlation Average vs log(params): r={r_avg_size:.3f}, p={p_avg_size:.4f}")

# 4. Hierarchical clustering of models by task profile (z-scores)
Z = linkage(z.values, method='ward')
clusters = fcluster(Z, t=4, criterion='maxclust')
df['cluster'] = clusters
print("\n=== Cluster assignment (Ward, k=4) ===")
for c in sorted(df['cluster'].unique()):
    members = df[df['cluster']==c]['Model'].tolist()
    print(f"Cluster {c} (n={len(members)}): {members}")
    print("  mean task z-profile:", z[df['cluster']==c].mean().round(2).to_dict())

# 5. Rank-based task similarity (Spearman) to compare with Pearson
spear = df[tasks].corr(method='spearman')
print("\n=== Spearman rank correlation matrix ===")
print(spear.round(2))

# Save full augmented table
df.to_csv('augmented_results.csv', index=False)

# --- Figure 1: heatmap of Pearson correlations ---
fig, ax = plt.subplots(figsize=(6,5))
im = ax.imshow(corr.values, cmap='RdBu_r', vmin=-1, vmax=1)
ax.set_xticks(range(len(tasks))); ax.set_xticklabels(tasks, rotation=45, ha='right', fontsize=8)
ax.set_yticks(range(len(tasks))); ax.set_yticklabels(tasks, fontsize=8)
for i in range(len(tasks)):
    for j in range(len(tasks)):
        ax.text(j, i, f"{corr.values[i,j]:.2f}", ha='center', va='center', fontsize=7,
                 color='white' if abs(corr.values[i,j])>0.6 else 'black')
plt.colorbar(im, ax=ax, label='Pearson r')
plt.title('Cross-task Pearson correlation (n=30 models)', fontsize=10)
plt.tight_layout()
plt.savefig('fig1_task_correlation_heatmap.png', dpi=200)
plt.close()

# --- Figure 2: dispersion vs average score scatter, colored by family ---
fig, ax = plt.subplots(figsize=(6,5))
families = df['Family'].unique()
colors = plt.cm.tab10(np.linspace(0,1,len(families)))
for fam, col in zip(families, colors):
    sub = df[df['Family']==fam]
    ax.scatter(sub['Average'], sub['profile_dispersion'], label=fam, color=col, s=50, edgecolor='k', linewidth=0.3)
ax.set_xlabel('Average MTEB score (7-task mean)')
ax.set_ylabel('Cross-task profile dispersion (SD of z-scores)')
ax.set_title('Overall performance vs. cross-task consistency')
ax.legend(fontsize=7, loc='best')
plt.tight_layout()
plt.savefig('fig2_dispersion_vs_average.png', dpi=200)
plt.close()

# --- Figure 3: dendrogram ---
fig, ax = plt.subplots(figsize=(10,5))
dendrogram(Z, labels=df['Model'].tolist(), leaf_rotation=90, leaf_font_size=7, ax=ax)
ax.set_title('Hierarchical clustering of models by 7-task performance profile (Ward linkage)')
ax.set_ylabel('Distance')
plt.tight_layout()
plt.savefig('fig3_dendrogram.png', dpi=200)
plt.close()

print("\nSaved: fig1_task_correlation_heatmap.png, fig2_dispersion_vs_average.png, fig3_dendrogram.png")