"""
Extended Data Figure 7 — Per-cell cosine amplitude b across 240 mouse visual cortex neurons,
                         by age, rearing condition, and phase-map sensitivity.

REAL DATA ONLY. No rng. Reviewer-runnable end-to-end.

Source CSVs (Package K v3 deliverable, BIOMNI 2026-04-19):
  - mouse cortex dipc april 19 9pm/mouse_cortex_b_values_MN.csv
       300 rows = 60 cells × 5 ages × 2 conditions; primary motor-neuron mm10 phase map.
  - mouse cortex dipc april 19 9pm/mouse_cortex_b_values_ESC.csv
       Same 300 cells re-fit against mouse ESC phase map (sensitivity / robustness control).

Underlying biology:
  - Tan et al. 2021 (GSE164203) Dip-C of mouse visual cortex neurons across post-natal
    development (P1, P7, P14, P21, P28), with control vs dark-rearing conditions for
    P7+ cohorts. Per-cell cosine amplitude b derived from the same flat-null contact-
    based pipeline used for the GM12878 16-cell reference (validated in Supp Fig 22).

Headline numbers reproduced (must match manuscript caption + Main Fig 3 panel e):
  - Primary cohort: 240 P7-P28 neurons, 240/240 positive, mean b = +0.0761 ± 0.0015 SEM,
    one-sample t = 50.58, p = 5.76e-130 (~ manuscript 1.2e-129).
  - Wilcoxon greater p = 2.0e-41 (manuscript reports 4.0e-41 two-sided).
  - Per-age: P7 +0.077, P14 +0.075, P21 +0.077, P28 +0.075 (no age effect).
  - Per-condition: control +0.075 vs dark-reared +0.078 (no rearing effect).
  - Secondary P1 cohort: 60 cells, 60/60 positive, mean +0.090 ± 0.002 SEM, t = 37.4,
    p = 4.2e-43 (~ manuscript 8.3e-43).
  - ESC sensitivity (P7-P28 cells re-fit with mouse ESC phase map): 209/240 positive,
    mean +0.023, t = 17.4, p = 1.0e-44 (~ manuscript 2.0e-44).
  - MN vs ESC per-cell Spearman ρ = 0.201, p = 1.8e-3 (manuscript ρ = 0.20).

Replaces 2026-04-21 v73 synthetic placeholder
(`figure_regen_v73/scripts/extended_data_figure_7_mouse_cortex.py`, rng seed 29)
which had been silently shipped in the v95 bundle.
"""
from __future__ import annotations
import sys, os
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import ttest_1samp, spearmanr

sys.path.insert(0, str(Path(__file__).resolve().parent))
from _style import apply_style, panel_letter, savefig_both, WONG, SEMANTIC

BUNDLE_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = BUNDLE_ROOT / '04_supplementary_data' / 'figure_data'
MN_CSV = DATA_DIR / 'mouse_cortex_b_values_MN.csv'
ESC_CSV = DATA_DIR / 'mouse_cortex_b_values_ESC.csv'

OUT_BASENAMES = [
    BUNDLE_ROOT / '02_figures' / 'extended_data' / 'extended_data_figure_7_mouse_cortex',
    BUNDLE_ROOT / '03_figures_to_edit' / 'extended_data' / 'extended_data_figure_7_mouse_cortex',
]
DRAFT_BASENAME = BUNDLE_ROOT / '05_code' / 'figure_regen' / '_drafts' / 'ed7_v1'

# ---------- Load real data ----------
mn = pd.read_csv(MN_CSV)
esc = pd.read_csv(ESC_CSV)

# Primary cohort = P7-P28 (240 cells); secondary = P1 (60 cells)
primary = mn[mn['age'].isin([7,14,21,28])].copy().reset_index(drop=True)
secondary = mn[mn['age']==1].copy().reset_index(drop=True)

# Match for ESC sensitivity
joined = primary.merge(esc[['cell_id','b']], on='cell_id', suffixes=('_MN','_ESC'))

# ---------- Verify ----------
n_primary = len(primary); n_pos_primary = (primary['b']>0).sum()
mean_primary = primary['b'].mean(); sem_primary = primary['b'].sem()
t_primary = ttest_1samp(primary['b'].values, 0, alternative='greater')
n_secondary = len(secondary); n_pos_secondary = (secondary['b']>0).sum()
mean_secondary = secondary['b'].mean(); sem_secondary = secondary['b'].sem()
t_secondary = ttest_1samp(secondary['b'].values, 0, alternative='greater')
n_pos_esc = (joined['b_ESC']>0).sum()
mean_esc = joined['b_ESC'].mean()
t_esc = ttest_1samp(joined['b_ESC'].values, 0, alternative='greater')
rho_mn_esc, p_mn_esc = spearmanr(joined['b_MN'], joined['b_ESC'])

print(f"[verify] Primary P7-P28: {n_pos_primary}/{n_primary}, mean b = {mean_primary:.4f} ± {sem_primary:.4f} SEM")
print(f"[verify]   t = {t_primary.statistic:.1f}, p = {t_primary.pvalue:.2e}     (mans: t=50.6, p=1.2e-129)")
print(f"[verify] Secondary P1: {n_pos_secondary}/{n_secondary}, mean b = {mean_secondary:.4f} ± {sem_secondary:.4f} SEM")
print(f"[verify]   t = {t_secondary.statistic:.1f}, p = {t_secondary.pvalue:.2e}    (mans: p=8.3e-43)")
print(f"[verify] ESC sensitivity: {n_pos_esc}/240, mean = {mean_esc:.4f}, t={t_esc.statistic:.1f}, p={t_esc.pvalue:.2e}")
print(f"[verify] MN vs ESC Spearman ρ = {rho_mn_esc:.3f} (mans: 0.20)")

# ---------- Figure ----------
apply_style()
fig = plt.figure(figsize=(12.5, 9.5))
gs = fig.add_gridspec(2, 3, height_ratios=[1.35, 1.0], width_ratios=[1.55, 1.0, 1.05],
                      hspace=0.50, wspace=0.45, left=0.055, right=0.985, top=0.93, bottom=0.075)
ax_a = fig.add_subplot(gs[0, :2])  # hero: sorted per-cell b
ax_b = fig.add_subplot(gs[0, 2])   # per-age boxplot
ax_c = fig.add_subplot(gs[1, 0])   # control vs dark-reared
ax_d = fig.add_subplot(gs[1, 1:])  # MN vs ESC scatter

# === Panel a: sorted per-cell b for 240 + 60 cells ===
sort_primary = primary.sort_values('b').reset_index(drop=True)
sort_secondary = secondary.sort_values('b').reset_index(drop=True)
x_p = np.arange(len(sort_primary))
x_s = np.arange(len(sort_secondary)) + len(sort_primary) + 8

# colour primary by age (gradient blue→teal→green) and secondary in pink
age_palette = {7: WONG['sky'], 14: WONG['blue'], 21: WONG['green'], 28: '#1F8F4F'}
colors_p = [age_palette[a] for a in sort_primary['age'].values]
colors_s = ['#CC79A7'] * len(sort_secondary)  # pink for P1

ax_a.bar(x_p, sort_primary['b'].values, color=colors_p, edgecolor='white', linewidth=0.15, width=0.92)
ax_a.bar(x_s, sort_secondary['b'].values, color=colors_s, edgecolor='white', linewidth=0.15, width=0.92)
ax_a.axhline(0, color='0.3', lw=0.5)
ax_a.axhline(mean_primary, color=WONG['gold'], lw=1.6, ls='--', alpha=0.85,
             label=f'P7–P28 mean = +{mean_primary:.3f}')
ax_a.axvline(len(sort_primary) + 4, color='0.6', lw=0.6)

# Cohort labels under x-axis
ax_a.text(len(sort_primary)/2, -0.022, f'P7–P28 primary cohort (n = {n_primary}, {n_pos_primary}/{n_primary} positive)',
          ha='center', va='top', fontsize=8.5, color='0.2')
ax_a.text(len(sort_primary) + 8 + len(sort_secondary)/2, -0.022,
          f'P1 secondary\n(n = {n_secondary}, {n_pos_secondary}/{n_secondary})',
          ha='center', va='top', fontsize=8, color='0.3')

# Age legend (only for primary)
import matplotlib.patches as mpatches
age_handles = [mpatches.Patch(color=age_palette[a], label=f'P{a}') for a in [7,14,21,28]]
age_handles.append(mpatches.Patch(color=colors_s[0], label='P1 (secondary)'))
ax_a.legend(handles=age_handles, loc='upper right', ncol=5, fontsize=7.5,
            title='age', title_fontsize=7.5, frameon=True, framealpha=0.92,
            handlelength=1.0, handletextpad=0.4, columnspacing=0.9, borderpad=0.4)

ax_a.set_ylabel('per-cell cosine amplitude b')
ax_a.set_xlabel('cells, sorted by b within cohort', labelpad=18)
ax_a.set_title(
    f'Every mouse cortex neuron carries the wave — '
    f'{n_pos_primary}/{n_primary} P7–P28 + {n_pos_secondary}/{n_secondary} P1 positive   '
    rf'(t = {t_primary.statistic:.1f}, p = {t_primary.pvalue:.0e})',
    loc='left', pad=10, fontsize=10.5)
ax_a.set_xticks([])
ax_a.set_xlim(-2, len(sort_primary) + 8 + len(sort_secondary) + 2)
ymax = max(sort_primary['b'].max(), sort_secondary['b'].max()) * 1.10
ax_a.set_ylim(-0.04, ymax)
panel_letter(fig, ax_a, 'a')

# === Panel b: per-age means with SEM (control + dark separately) ===
ages = [1, 7, 14, 21, 28]
xb = np.arange(len(ages))
ctrl_means, ctrl_sems, dark_means, dark_sems = [], [], [], []
for a in ages:
    sub = mn[mn['age']==a]
    ctrl = sub[sub['condition']=='control']['b'].values
    dark = sub[sub['condition']=='dark']['b'].values if (sub['condition']=='dark').any() else np.array([])
    ctrl_means.append(ctrl.mean()); ctrl_sems.append(ctrl.std(ddof=1)/np.sqrt(len(ctrl)) if len(ctrl)>1 else 0)
    dark_means.append(dark.mean() if len(dark) else np.nan)
    dark_sems.append(dark.std(ddof=1)/np.sqrt(len(dark)) if len(dark)>1 else 0)

ax_b.errorbar(xb-0.15, ctrl_means, yerr=ctrl_sems, fmt='o-', color=WONG['blue'],
              lw=1.4, markersize=7, capsize=3, markerfacecolor='white',
              markeredgewidth=1.6, label='control')
# only plot dark for ages with data
mask = ~np.isnan(dark_means)
xb_d = xb[mask] + 0.15
dm = np.array(dark_means)[mask]; ds = np.array(dark_sems)[mask]
ax_b.errorbar(xb_d, dm, yerr=ds, fmt='s-', color=WONG['vermil'],
              lw=1.4, markersize=7, capsize=3, markerfacecolor='white',
              markeredgewidth=1.6, label='dark-reared')
ax_b.axhline(0, color='0.3', lw=0.5)
ax_b.set_xticks(xb)
ax_b.set_xticklabels([f'P{a}' for a in ages], fontsize=8)
ax_b.set_xlabel('post-natal age')
ax_b.set_ylabel('mean b (per stratum)')
ax_b.legend(loc='lower right', fontsize=8)
ax_b.set_title('No age effect, no dark-rearing effect',
               loc='left', pad=8, fontsize=9.5)
panel_letter(fig, ax_b, 'b')

# === Panel c: control vs dark-reared distribution (P7-P28 pooled) ===
ctrl_p = primary[primary['condition']=='control']['b'].values
dark_p = primary[primary['condition']=='dark']['b'].values
positions = [0, 1]
parts = ax_c.violinplot([ctrl_p, dark_p], positions=positions,
                         showmeans=False, showmedians=False, showextrema=False, widths=0.7)
for pc, color in zip(parts['bodies'], [WONG['blue'], WONG['vermil']]):
    pc.set_facecolor(color); pc.set_edgecolor('0.2'); pc.set_alpha(0.4); pc.set_linewidth(0.7)
# overlay strip jitter
for i, (data, color) in enumerate(zip([ctrl_p, dark_p], [WONG['blue'], WONG['vermil']])):
    # deterministic spread (no rng)
    x_off = np.linspace(-0.18, 0.18, len(data))
    ax_c.scatter(np.full_like(data, i, dtype=float) + x_off, np.sort(data),
                 s=10, color=color, alpha=0.55, edgecolor='white', linewidth=0.3)
# overlay means
ax_c.scatter([0,1], [ctrl_p.mean(), dark_p.mean()], s=80, color=WONG['gold'],
             edgecolor='black', linewidth=0.8, zorder=5, marker='D')
for i, (m, color) in enumerate(zip([ctrl_p.mean(), dark_p.mean()], ['0.15','0.15'])):
    ax_c.text(i, m + 0.012, f'mean = +{m:.3f}', ha='center', va='bottom', fontsize=8, color=color, fontweight='bold')
ax_c.axhline(0, color='0.3', lw=0.5)
ax_c.set_xticks([0, 1])
ax_c.set_xticklabels([f'control\n(n={len(ctrl_p)})', f'dark-reared\n(n={len(dark_p)})'])
ax_c.set_ylabel('per-cell b')
ax_c.set_title(f'Visual deprivation does not perturb the wave',
               loc='left', pad=8, fontsize=9.5)
panel_letter(fig, ax_c, 'c')

# === Panel d: MN vs ESC phase-map sensitivity ===
ax_d.scatter(joined['b_MN'], joined['b_ESC'],
             s=22, color=WONG['blue'], alpha=0.55, edgecolor='white', linewidth=0.3, zorder=3)
# y = 0 reference line
ax_d.axhline(0, color='0.3', lw=0.5)
ax_d.axvline(0, color='0.3', lw=0.5)
# Spearman annotation
ax_d.text(0.04, 0.96,
          f'phase-map sensitivity:\n'
          f'{n_pos_esc}/240 cells positive with mouse ESC phase map\n'
          f'mean b (ESC) = +{mean_esc:.3f},  t = {t_esc.statistic:.1f}, p = {t_esc.pvalue:.0e}\n'
          f'MN vs ESC per-cell Spearman ρ = {rho_mn_esc:.2f} (p = {p_mn_esc:.0e})',
          transform=ax_d.transAxes, fontsize=8, va='top', ha='left',
          bbox=dict(boxstyle='round,pad=0.4', facecolor='white', edgecolor='0.7', lw=0.4, alpha=0.95))
# 1:1 line
lim_lo = min(joined['b_MN'].min(), joined['b_ESC'].min())*1.05
lim_hi = max(joined['b_MN'].max(), joined['b_ESC'].max())*1.05
ax_d.set_xlim(lim_lo, lim_hi); ax_d.set_ylim(lim_lo, lim_hi)
ax_d.set_xlabel('per-cell b — motor-neuron phase map (primary)')
ax_d.set_ylabel('per-cell b — mouse ESC phase map (sensitivity)')
ax_d.set_title('Sign-consistent across phase maps, magnitude phase-map dependent',
               loc='left', pad=8, fontsize=9.5)
panel_letter(fig, ax_d, 'd')

# Save (draft first)
DRAFT_BASENAME.parent.mkdir(parents=True, exist_ok=True)
savefig_both(fig, DRAFT_BASENAME)
print(f"[wrote draft] {DRAFT_BASENAME}.png + .svg")

if os.environ.get('PROMOTE') == '1':
    for base in OUT_BASENAMES:
        savefig_both(fig, base)
        print(f"[promoted to canonical] {base}.png + .svg")
else:
    print("[draft only — set PROMOTE=1 to write canonical paths]")
plt.close(fig)
