import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import seaborn as sns
import networkx as nx
from networkx.algorithms import community as nxcomm

PA = ['Interested', 'Excited', 'Strong', 'Enthusiastic', 'Proud',
      'Alert', 'Inspired', 'Determined', 'Attentive', 'Active']
NA = ['Distressed', 'Upset', 'Guilty', 'Scared', 'Hostile',
      'Irritable', 'Ashamed', 'Nervous', 'Jittery', 'Afraid']
ALL_ITEMS = PA + NA
VALENCE = {it: 'PA' for it in PA}
VALENCE.update({it: 'NA' for it in NA})

plt.rcParams.update({
    'font.family': 'serif', 'font.size': 11,
    'axes.titlesize': 12, 'axes.labelsize': 11,
    'figure.dpi': 150,
})

# Load similarity matrix
sim = pd.read_csv('out_sim_matrix_tfidf.csv', index_col=0)

# --------------------------------------------------------------------------- #
# Figure 1: Semantic similarity matrix (all 20 items, valence-ordered)
# --------------------------------------------------------------------------- #
order = PA + NA
sim_ord = sim.loc[order, order]

fig, ax = plt.subplots(figsize=(9, 8))
# Show off-diagonal only; mask the diagonal
mask = np.eye(len(order), dtype=bool)
sns.heatmap(sim_ord, mask=mask, cmap='YlGnBu', vmin=0, vmax=0.7,
            cbar_kws={'label': 'Cosine similarity'},
            xticklabels=order, yticklabels=order,
            linewidths=0.4, linecolor='white', ax=ax)
# Block lines to mark PA/NA boundary
ax.axhline(10, color='black', lw=1.0)
ax.axvline(10, color='black', lw=1.0)
ax.set_title('Semantic Similarity Among PANAS Items', fontsize=13)
ax.tick_params(axis='x', rotation=45)
ax.set_xticklabels(ax.get_xticklabels(), ha='right')
plt.tight_layout()
plt.savefig('fig1_similarity_matrix.png', dpi=300, bbox_inches='tight')
plt.close()

# --------------------------------------------------------------------------- #
# Figure 2: Within-valence vs cross-valence boxplot
# --------------------------------------------------------------------------- #
within_pa, within_na, cross = [], [], []
n = len(order)
for i in range(n):
    for j in range(i+1, n):
        v = sim_ord.iloc[i, j]
        li, lj = VALENCE[order[i]], VALENCE[order[j]]
        if li == 'PA' and lj == 'PA':
            within_pa.append(v)
        elif li == 'NA' and lj == 'NA':
            within_na.append(v)
        else:
            cross.append(v)

fig, ax = plt.subplots(figsize=(7, 5.5))
data = [within_pa, within_na, cross]
labels = ['Within positive\n(n = 45)', 'Within negative\n(n = 45)', 'Cross-valence\n(n = 100)']
positions = [1, 2, 3]
bp = ax.boxplot(data, positions=positions, widths=0.55, patch_artist=True,
                medianprops=dict(color='black', lw=1.5))
colors = ['#4daf4a', '#e41a1c', '#999999']
for patch, c in zip(bp['boxes'], colors):
    patch.set_facecolor(c)
    patch.set_alpha(0.6)
# Overlay individual points
for i, d in enumerate(data):
    x = np.random.default_rng(7).normal(positions[i], 0.06, size=len(d))
    ax.scatter(x, d, color=colors[i], alpha=0.7, s=14, edgecolors='black', linewidth=0.3)

ax.set_xticks(positions)
ax.set_xticklabels(labels)
ax.set_ylabel('Cosine similarity')
ax.set_title('Within-Valence vs. Cross-Valence Semantic Similarity', fontsize=12)
ax.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.savefig('fig2_valence_boxplot.png', dpi=300, bbox_inches='tight')
plt.close()

# --------------------------------------------------------------------------- #
# Figure 3: Semantic network (threshold = 0.15) with communities
# --------------------------------------------------------------------------- #
G = nx.Graph()
for it in order:
    G.add_node(it, valence=VALENCE[it])
for i in range(n):
    for j in range(i+1, n):
        w = sim_ord.iloc[i, j]
        if w >= 0.15:
            G.add_edge(order[i], order[j], weight=float(w))

communities = list(nxcomm.greedy_modularity_communities(G, weight='weight'))
# assign community colors
palette = ['#4daf4a', '#e41a1c', '#377eb8', '#984ea3', '#ff7f00', '#a65628']
node_color = {}
for c_idx, comm in enumerate(communities):
    for node in comm:
        node_color[node] = palette[c_idx % len(palette)]

# layout: spring with seed for reproducibility
pos = nx.spring_layout(G, k=1.4, seed=2026, weight='weight', iterations=200)

fig, ax = plt.subplots(figsize=(10, 8))
edge_weights = [G[u][v]['weight'] for u, v in G.edges()]
edge_widths = [0.5 + 4*((w-0.15)/(max(edge_weights)-0.15+1e-9)) for w in edge_weights]
edge_alphas = [0.25 + 0.65*((w-0.15)/(max(edge_weights)-0.15+1e-9)) for w in edge_weights]
for (u, v), w, ew, ea in zip(G.edges(), edge_weights, edge_widths, edge_alphas):
    ax.plot([pos[u][0], pos[v][0]], [pos[u][1], pos[v][1]],
            color='gray', linewidth=ew, alpha=ea, zorder=1)

# nodes
for node, (x, y) in pos.items():
    ax.scatter(x, y, s=900, c=node_color[node], edgecolors='black',
               linewidths=1.5, zorder=2)
    ax.text(x, y, node, ha='center', va='center', fontsize=9, fontweight='bold',
            zorder=3)

ax.set_axis_off()
ax.set_title('Semantic Network of PANAS Items (Edge Threshold = 0.15)', fontsize=13)

# Community legend
community_names = []
for i, comm in enumerate(communities):
    nodes_list = sorted(comm)
    label = f"Community {i+1} (n={len(comm)})"
    ax.scatter([], [], s=160, c=palette[i % len(palette)],
               edgecolors='black', linewidths=1.2, label=label)
ax.legend(loc='upper left', bbox_to_anchor=(1.0, 1.0), fontsize=9, frameon=False)

plt.tight_layout()
plt.savefig('fig3_network.png', dpi=300, bbox_inches='tight')
plt.close()

# --------------------------------------------------------------------------- #
# Figure 4: Mean similarity bar chart (valence structure)
# --------------------------------------------------------------------------- #
means = [np.mean(within_pa), np.mean(within_na), np.mean(cross)]
sds = [np.std(within_pa, ddof=1), np.std(within_na, ddof=1), np.std(cross, ddof=1)]
labels = ['Within positive', 'Within negative', 'Cross-valence']

fig, ax = plt.subplots(figsize=(7, 5))
bars = ax.bar(labels, means, yerr=sds, capsize=8,
              color=['#4daf4a', '#e41a1c', '#999999'], alpha=0.7,
              edgecolor='black', linewidth=1.2)
for b, m in zip(bars, means):
    ax.text(b.get_x() + b.get_width()/2, m + 0.012, f'M = {m:.3f}',
            ha='center', fontsize=11, fontweight='bold')
ax.set_ylabel('Mean cosine similarity')
ax.set_title('Mean Semantic Similarity by Pair Type', fontsize=12)
ax.set_ylim(0, 0.34)
ax.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.savefig('fig2_mean_similarity.png', dpi=300, bbox_inches='tight')
plt.close()

# --------------------------------------------------------------------------- #
# Supplementary Figure S2: TF-IDF vs SBERT validation scatter
# --------------------------------------------------------------------------- #
sim_sbert = pd.read_csv('out_sim_matrix_sbert.csv', index_col=0).loc[order, order]
iu = np.triu_indices(20, k=1)
tf_off = sim_ord.values[iu]
sb_off = sim_sbert.values[iu]
r_val = np.corrcoef(tf_off, sb_off)[0, 1]

# Also encode pair type
pair_types = []
for i, j in zip(*iu):
    li, lj = VALENCE[order[i]], VALENCE[order[j]]
    if li == 'PA' and lj == 'PA':
        pair_types.append('Within PA')
    elif li == 'NA' and lj == 'NA':
        pair_types.append('Within NA')
    else:
        pair_types.append('Cross')

fig, ax = plt.subplots(figsize=(7, 6))
colors_map = {'Within PA': '#4daf4a', 'Within NA': '#e41a1c', 'Cross': '#999999'}
for pt in ['Within PA', 'Within NA', 'Cross']:
    mask = [p == pt for p in pair_types]
    ax.scatter(np.array(tf_off)[mask], np.array(sb_off)[mask],
               c=colors_map[pt], label=pt, alpha=0.7, s=45, edgecolors='black',
               linewidths=0.5)
ax.set_xlabel('TF-IDF cosine similarity')
ax.set_ylabel('Sentence-BERT (MiniLM) cosine similarity')
ax.set_title(f'Convergence of TF-IDF and Transformer-Based Similarity\n(r = {r_val:.3f})',
             fontsize=11)
ax.legend(loc='upper left')
ax.grid(alpha=0.3)
plt.tight_layout()
plt.savefig('figS2_tfidf_vs_sbert.png', dpi=300, bbox_inches='tight')
plt.close()

# --------------------------------------------------------------------------- #
# Supplementary Figure S3: Coherence by item with 95% CIs
# --------------------------------------------------------------------------- #
coh = pd.read_csv('out_coherence_full.csv').sort_values('cos_mean')
fig, ax = plt.subplots(figsize=(8, 6.5))
ypos = np.arange(len(coh))
colors_v = ['#4daf4a' if v == 'PA' else '#e41a1c' for v in coh['Valence']]
ax.barh(ypos, coh['cos_mean'], color=colors_v, alpha=0.75,
        edgecolor='black', linewidth=0.6,
        xerr=[coh['cos_mean'] - coh['cos_95CI_low'],
              coh['cos_95CI_high'] - coh['cos_mean']], capsize=3)
ax.set_yticks(ypos)
ax.set_yticklabels(coh['Item'])
ax.set_xlabel('Mean within-item cosine similarity (95% CI)')
ax.set_title('Per-Item Semantic Coherence', fontsize=12)
ax.grid(axis='x', alpha=0.3)
# Legend
from matplotlib.patches import Patch
legend_elements = [Patch(facecolor='#4daf4a', alpha=0.75, edgecolor='black', label='Positive affect'),
                   Patch(facecolor='#e41a1c', alpha=0.75, edgecolor='black', label='Negative affect')]
ax.legend(handles=legend_elements, loc='lower right')
plt.tight_layout()
plt.savefig('figS1_coherence_ci.png', dpi=300, bbox_inches='tight')
plt.close()

print('All figures generated.')