#!/usr/bin/env python3
"""Regenerate Figure 2 (feature-level trait-state map) from Figure3_Cluster_Results.csv.
Fixes the stale panel-A label (levodopa n = 154 -> 122). Data verified to reproduce
r = 0.66, slope = 1.55, p < 0.001."""
import numpy as np, pandas as pd, matplotlib.pyplot as plt
from scipy.stats import linregress
import matplotlib.gridspec as gridspec
from matplotlib.lines import Line2D
from matplotlib.colors import LinearSegmentedColormap

CSV = "/sessions/stoic-dazzling-lovelace/mnt/(5) PNAS/data/Output_Data/Figure3_Cluster_Results.csv"
OUT = "/sessions/stoic-dazzling-lovelace/mnt/outputs"
N_LEVO, N_DBS = 122, 15

PALETTE = {'Speed':'#2ca02c','Amplitude':'#1f77b4','Variability/Consistency':'#d62728',
           'Timing/Rhythm':'#ff7f0e','Other':'grey'}
MARK = {'Speed':'s','Amplitude':'o','Variability/Consistency':'^','Timing/Rhythm':'D','Other':'X'}
ORDER = ['Speed','Amplitude','Variability/Consistency','Timing/Rhythm','Other']

df = pd.read_csv(CSV, index_col=0)
df['Cat'] = df['Category'].where(df['Category'].isin(PALETTE), 'Other')

plt.rcParams.update({'font.family':'Liberation Sans','axes.titleweight':'bold',
                     'axes.labelweight':'bold','savefig.dpi':600,'svg.fonttype':'none'})

fig = plt.figure(figsize=(16, 7))
outer = gridspec.GridSpec(1, 2, width_ratios=[1, 1.12], wspace=0.28,
                          left=0.06, right=0.985, top=0.86, bottom=0.11)

# ---------------- Panel A : cross-feature concordance ----------------
axA = fig.add_subplot(outer[0])
x = df['Levodopa_Responsiveness'].values
y = df['DBS_Responsiveness'].values
sl, ic, r, p, se = linregress(x, y)
xs = np.linspace(-20, 20, 100)
ys = ic + sl*xs
# 95% CI band for the mean line
n = len(x); xm = x.mean(); sxx = np.sum((x-xm)**2)
resid = y - (ic + sl*x); s_err = np.sqrt(np.sum(resid**2)/(n-2))
from scipy.stats import t as tdist
tval = tdist.ppf(0.975, n-2)
ci = tval*s_err*np.sqrt(1/n + (xs-xm)**2/sxx)
axA.fill_between(xs, ys-ci, ys+ci, color='grey', alpha=0.25, lw=0)
axA.plot(xs, ys, color='black', lw=1.6, zorder=4)
for cat in ORDER:
    sub = df[df['Cat']==cat]
    if len(sub):
        axA.scatter(sub['Levodopa_Responsiveness'], sub['DBS_Responsiveness'],
                    c=PALETTE[cat], marker=MARK[cat], s=95, edgecolor='black',
                    linewidth=0.5, alpha=0.9, zorder=5)
axA.axhline(0, color='black', lw=0.8, alpha=0.3); axA.axvline(0, color='black', lw=0.8, alpha=0.3)
axA.set_xlim(-20, 20); axA.set_ylim(-60, 60)
axA.set_xlabel('Per-feature levodopa responsiveness (% Δ)', fontsize=12)
axA.set_ylabel('Per-feature DBS responsiveness (% Δ)', fontsize=12)
axA.set_title(f'A) Cross-feature concordance\n(Dataset 2 levodopa, n = {N_LEVO}; Dataset 3 DBS, n = {N_DBS})',
              fontsize=13)
axA.text(0.03, 0.97, f'r = {r:.2f}\nslope = {sl:.2f}\np < 0.001\nfeature-level',
         transform=axA.transAxes, va='top', ha='left', fontsize=10.5,
         bbox=dict(boxstyle='round', fc='white', ec='0.5', alpha=0.9))
axA.spines[['top','right']].set_visible(False)
handles = [Line2D([0],[0], marker=MARK[c], color='w', markerfacecolor=PALETTE[c],
                  markeredgecolor='black', markersize=9, label=c) for c in ORDER]
axA.legend(handles=handles, title='Feature domain', loc='lower right', fontsize=8.5,
           title_fontsize=9, framealpha=0.95)

# ---------------- Panel B : trait-state plane with marginals ----------------
gsB = gridspec.GridSpecFromSubplotSpec(2, 2, subplot_spec=outer[1],
        width_ratios=[5,1], height_ratios=[1,5], wspace=0.04, hspace=0.04)
axB  = fig.add_subplot(gsB[1,0])
axBt = fig.add_subplot(gsB[0,0], sharex=axB)
axBr = fig.add_subplot(gsB[1,1], sharey=axB)

xt = df['Abs_Trait'].values          # |trait correlation|
ys2 = df['Mean_Response'].values     # mean therapy responsiveness
XLIM = (0, 0.55); YLIM = (0, 40)

# diverging background: blue = state-dominant (high response/low trait), red = trait-dominant
tx = np.linspace(*XLIM, 800); ty = np.linspace(*YLIM, 800)
TX, TY = np.meshgrid(tx, ty)
tnorm = TX/XLIM[1]; snorm = TY/YLIM[1]
field = snorm - tnorm                # >0 state-dominant, <0 trait-dominant
cmap = LinearSegmentedColormap.from_list('ts', ['#d96b6b','#f3e6e6','#ffffff','#e6ecf5','#6b8fd9'])
axB.imshow(field, extent=[*XLIM,*YLIM], origin='lower', aspect='auto',
           cmap=cmap, alpha=0.35, vmin=-1, vmax=1, zorder=0, interpolation='bilinear', rasterized=True)
# iso-dominance diagonal (snorm == tnorm)
diag_x = np.array(XLIM); diag_y = YLIM[1]*(diag_x/XLIM[1])
axB.plot(diag_x, diag_y, ls=':', color='grey', lw=1.3, zorder=1)
axB.text(0.50, 0.50*YLIM[1]/XLIM[1]*0.97, 'iso-dominance', rotation=39, fontsize=8.5,
         color='grey', style='italic', ha='right', va='bottom')
axB.text(0.02, 38, 'State-dominant\nregion', fontsize=9.5, color='#2f5fa8',
         style='italic', va='top', ha='left', weight='bold')
axB.text(0.42, 2.6, 'Trait-dominant region', fontsize=9.5, color='#a83232',
         style='italic', va='center', ha='center', weight='bold')

for cat in ORDER:
    sub = df[df['Cat']==cat]
    if len(sub):
        axB.scatter(sub['Abs_Trait'], sub['Mean_Response'], c=PALETTE[cat], marker=MARK[cat],
                    s=95, edgecolor='black', linewidth=0.5, alpha=0.92, zorder=5)
# label a curated set of features
LABELS = ['RangeCycleDuration','MeanClosingSpeed','MeanCycleDuration','MeanAmplitude',
          'AmplitudeDecay','VelocityDecay','CV_ClosingSpeed','CV_Speed','CV_RMSVelocity',
          'StdAmplitude','RateDecay','MeanSpeed']
off = {'RangeCycleDuration':(7,3),'MeanClosingSpeed':(-7,5,'right'),'MeanCycleDuration':(-7,4,'right'),
       'MeanAmplitude':(-7,-3,'right'),'AmplitudeDecay':(7,2),'VelocityDecay':(7,-5),
       'CV_ClosingSpeed':(7,4),'CV_Speed':(7,2),'CV_RMSVelocity':(7,-6),
       'StdAmplitude':(7,3),'RateDecay':(-7,3,'right'),'MeanSpeed':(7,-2)}
for f in LABELS:
    if f in df.index:
        dx,dy = off[f][0], off[f][1]; ha = off[f][2] if len(off[f])>2 else 'left'
        axB.annotate(f, (df.loc[f,'Abs_Trait'], df.loc[f,'Mean_Response']),
                     textcoords='offset points', xytext=(dx,dy), fontsize=7.3,
                     color='0.15', ha=ha, zorder=6)
axB.set_xlim(*XLIM); axB.set_ylim(*YLIM)
axB.set_xlabel('|Trait correlation| with putamen DaT (Dataset 1, n = 35)', fontsize=12)
axB.set_ylabel('Mean therapy responsiveness\n(mean of |Δlevodopa|, |ΔDBS|) [%]', fontsize=12)
axB.spines[['top','right']].set_visible(False)
handles2 = [Line2D([0],[0], marker=MARK[c], color='w', markerfacecolor=PALETTE[c],
                   markeredgecolor='black', markersize=9, label=c) for c in ORDER]
axB.legend(handles=handles2, title='Feature domain', loc='upper right', fontsize=8.5,
           title_fontsize=9, framealpha=0.95)
# marginal histograms
axBt.hist(xt, bins=10, color='0.6', edgecolor='white'); axBt.axis('off')
axBr.hist(ys2, bins=10, orientation='horizontal', color='0.6', edgecolor='white'); axBr.axis('off')
axBt.set_title('B) Feature-level trait–state plane (continuous, no partition)',
               fontsize=13, pad=8, loc='left')

fig.suptitle('Figure 2. Feature-level trait–state map (dimensional view)',
             fontsize=15, weight='bold', y=0.975)

fig.savefig(f"{OUT}/Figure2_PNAS.pdf", dpi=600, bbox_inches='tight', facecolor='white')
fig.savefig(f"{OUT}/Figure2_PNAS.svg", dpi=600, bbox_inches='tight', facecolor='white')
fig.savefig(f"{OUT}/Figure2_PNAS.png", dpi=600, bbox_inches='tight', facecolor='white')
print("saved Figure2_PNAS pdf/svg/png")
print(f"verify: r={r:.3f} slope={sl:.3f} p={p:.2e} n_features={n}")