create_figures.py

#!/usr/bin/env python3
"""
Generate clean, publication-quality figures for DMT-Schizophrenia analysis.
Uses 8x6 inch figures at 150 DPI for crisp display.
"""

import json
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import os

OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "figures")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ── Clean, modern style ────────────────────────────────────────────────────────

plt.rcParams.update({
    'font.family': 'sans-serif',
    'font.sans-serif': ['Helvetica', 'Arial', 'DejaVu Sans'],
    'font.size': 11,
    'axes.labelsize': 11,
    'axes.titlesize': 13,
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'legend.fontsize': 10,
    'axes.linewidth': 1.0,
    'axes.spines.top': False,
    'axes.spines.right': False,
    'figure.dpi': 150,
    'savefig.dpi': 150,
    'savefig.bbox': 'tight',
    'savefig.pad_inches': 0.1,
    'axes.grid': False,
})

# Clean palette
C = {
    'blue':   '#2E5C8A',
    'red':    '#C0392B',
    'green':  '#27AE60',
    'orange': '#E67E22',
    'purple': '#8E44AD',
    'gray':   '#95A5A6',
    'dark':   '#2C3E50',
    'light':  '#ECF0F1',
    'white':  '#FFFFFF',
    'bg':     '#FAFAFA',
}

# Load data
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "results", "full_results.json")) as f:
    R = json.load(f)

flux = R["tryptophan_flux_model"]["deterministic_comparison"]
mc = R["tryptophan_flux_model"]["monte_carlo_results"]
receptors = R["5ht2a_receptor_analysis"]["patient_density"]
vmat2 = R["vmAT2_analysis"]["patient_vmAT2"]
dmt_impact = R["vmAT2_analysis"]["dmt_packaging_impact"]
genes = R["genetic_enrichment"]["key_genes"]
pathway_score = R["genetic_enrichment"]["pathway_score"]
nmda = R["nmda_dmt_model"]["results"]
bayes = R["bayesian_integration"]
likelihoods = bayes["likelihoods"]
posterior = bayes["posterior"]["final_posterior"]
sequential = bayes["posterior"]["sequential_update"]
sensitivity = bayes["sensitivity"]


def save_fig(fig, name):
    pdf = os.path.join(OUTPUT_DIR, f"{name}.pdf")
    png = os.path.join(OUTPUT_DIR, f"{name}.png")
    fig.savefig(pdf, format='pdf', bbox_inches='tight')
    fig.savefig(png, format='png', bbox_inches='tight')
    plt.close(fig)
    sz = os.path.getsize(png)
    print(f"  {name}.png  ({sz/1024:.0f} KB)")


# ============================================================================
# FIGURE 1: Tryptophan Metabolic Flux — clean bar chart only
# ============================================================================

def fig1():
    fig, ax = plt.subplots(figsize=(8, 5))

    # Raw values from the flux model
    labels = ['Kynurenine\nFraction', 'Available\nTryptophan', 'Tryptamine\nProduction',
              'Tryptamine\nSurvival', 'Net Tryptamine\nSignal', 'Serotonin\nFlux']

    ctrl = np.array([
        flux['control_flux']['kynurenine_flux_fraction'],
        flux['control_flux']['available_trp'],
        flux['control_flux']['tryptamine_production'],
        flux['control_flux']['tryptamine_survival_fraction'],
        flux['control_flux']['net_tryptamine_signal'],
        flux['control_flux']['serotonin_flux'],
    ])
    scz = np.array([
        flux['schizophrenia_flux']['kynurenine_flux_fraction'],
        flux['schizophrenia_flux']['available_trp'],
        flux['schizophrenia_flux']['tryptamine_production'],
        flux['schizophrenia_flux']['tryptamine_survival_fraction'],
        flux['schizophrenia_flux']['net_tryptamine_signal'],
        flux['schizophrenia_flux']['serotonin_flux'],
    ])

    # Normalize per-metric to make bars comparable
    ctrl_norm = ctrl / ctrl
    scz_norm = scz / ctrl

    x = np.arange(len(labels))
    w = 0.35

    bars1 = ax.bar(x - w/2, ctrl_norm, w, label='Control', color=C['blue'],
                   edgecolor='white', linewidth=0.5)
    bars2 = ax.bar(x + w/2, scz_norm, w, label='Schizophrenia', color=C['red'],
                   edgecolor='white', linewidth=0.5)

    # Fold change labels
    for i in range(len(labels)):
        fc = scz_norm[i]
        color = C['red'] if fc < 1 else C['green']
        ax.text(i, max(ctrl_norm[i], scz_norm[i]) + 0.03,
                f'{fc:.2f}x', ha='center', fontsize=10, fontweight='bold', color=color)

    # Significance stars for key metrics
    for i in [2, 4]:  # tryptamine production, net signal
        y_max = max(ctrl_norm[i], scz_norm[i])
        ax.plot([x[i] - w/2, x[i] + w/2], [y_max + 0.07, y_max + 0.07],
                '-', color=C['dark'], lw=1.2)
        ax.text(i, y_max + 0.11, '***', ha='center', fontsize=10, fontweight='bold', color=C['dark'])

    ax.set_xticks(x)
    ax.set_xticklabels(labels, fontsize=10)
    ax.set_ylabel('Normalized to Control (= 1.0)', fontsize=11)
    ax.set_ylim(0, 1.55)
    ax.legend(loc='upper left', frameon=False)

    # Highlight the key finding
    ax.text(0.5, 1.48, 'Net DMT Signal: 0.74x (26% reduction, p < 0.0001)',
            ha='center', fontsize=11, fontweight='bold', color=C['red'],
            transform=ax.transAxes,
            bbox=dict(boxstyle='round,pad=0.4', facecolor=C['light'], edgecolor=C['red'], lw=1.5))

    fig.suptitle('Tryptophan Metabolic Flux in Schizophrenia', fontsize=14, fontweight='bold', y=0.97)
    save_fig(fig, 'fig01_trp_flux')


# ============================================================================
# FIGURE 2: 5-HT2A Receptor Density — clean grouped bar chart
# ============================================================================

def fig2():
    fig, ax = plt.subplots(figsize=(8, 5))

    region_keys = ['prefrontal_cortex', 'temporal_cortex', 'occipital_cortex', 'striatum']
    labels = ['Prefrontal\nCortex', 'Temporal\nCortex', 'Occipital\nCortex', 'Striatum']

    ctrl = np.array([receptors[k]['control_mean'] for k in region_keys])
    pat = np.array([receptors[k]['patient_mean'] for k in region_keys])
    fc = np.array([receptors[k]['fold_change'] for k in region_keys])
    sig = np.array([receptors[k]['significant'] for k in region_keys])

    x = np.arange(len(labels))
    w = 0.35

    bars1 = ax.bar(x - w/2, ctrl, w, label='Control', color=C['blue'],
                   edgecolor='white', linewidth=0.5)
    bars2 = ax.bar(x + w/2, pat, w, label='Schizophrenia', color=C['red'],
                   edgecolor='white', linewidth=0.5)

    for i in range(len(labels)):
        fc_color = C['green'] if fc[i] > 1 else C['red']
        sig_label = '***' if sig[i] else 'n.s.'
        ax.text(i, pat[i] + 0.3, f'{fc[i]:.2f}x {sig_label}',
                ha='center', fontsize=10, fontweight='bold', color=fc_color)

    ax.set_xticks(x)
    ax.set_xticklabels(labels, fontsize=10)
    ax.set_ylabel('BPND (Binding Potential)', fontsize=11)
    ax.set_ylim(0, 9)
    ax.legend(loc='upper left', frameon=False)

    fig.suptitle('5-HT2A Receptor Density in Schizophrenia', fontsize=14, fontweight='bold', y=0.97)
    save_fig(fig, 'fig02_5ht2a_density')


# ============================================================================
# FIGURE 3: VMAT2 → DMT Availability — clean waterfall
# ============================================================================

def fig3():
    fig, ax = plt.subplots(figsize=(8, 5))

    steps = [
        ('VMAT2\nReduction', 18.3, C['red']),
        ('Release\nEfficiency', 10.1, C['orange']),
        ('MAO Exposure\nIncrease', 22.4, C['gray']),
        ('Net DMT\nAvailability', 26.6, C['purple']),
    ]

    x = np.arange(len(steps))
    w = 0.55
    values = [s[1] for s in steps]
    colors = [s[2] for s in steps]

    bars = ax.bar(x, values, w, color=colors, edgecolor='white', linewidth=0.8)

    for i, (bar, step) in enumerate(zip(bars, steps)):
        height = bar.get_height()
        detail = f'{dmt_impact["avg_vmAT2_fold_change"]:.2f}x' if i == 0 else \
                 f'{dmt_impact["release_efficiency_fold"]:.2f}x' if i == 1 else \
                 f'{dmt_impact["mao_exposure_increase"]:.2f}x' if i == 2 else \
                 f'{dmt_impact["net_dmt_availability_impact"]:.2f}x'
        ax.text(bar.get_x() + bar.get_width()/2, height + 1.2,
                detail, ha='center', fontsize=11, fontweight='bold', color=colors[i])
        ax.text(bar.get_x() + bar.get_width()/2, height/2,
                f'{height:.0f}%', ha='center', va='center',
                fontsize=12, fontweight='bold', color='white')

    ax.set_xticks(x)
    ax.set_xticklabels([s[0] for s in steps], fontsize=10)
    ax.set_ylabel('Effect Magnitude (%)', fontsize=11)
    ax.set_ylim(0, 35)

    # VMAT2 regional inset
    vmat2_keys = ['striatum', 'thalamus', 'prefrontal_cortex']
    vmat2_names = [vmat2[k]['region'] for k in vmat2_keys]
    vmat2_fc = np.array([vmat2[k]['fold_change'] for k in vmat2_keys])

    inset_ax = fig.add_axes([0.55, 0.12, 0.38, 0.35])
    y_pos = np.arange(len(vmat2_fc))
    bars = inset_ax.barh(y_pos, vmat2_fc, color=C['blue'], edgecolor='white', linewidth=0.5)
    inset_ax.set_yticks(y_pos)
    inset_ax.set_yticklabels(vmat2_names, fontsize=9)
    inset_ax.set_xlabel('Patient / Control', fontsize=9)
    inset_ax.axvline(x=1.0, color=C['red'], linestyle='--', lw=1.5, alpha=0.7)
    inset_ax.set_xlim(0.65, 1.05)
    inset_ax.set_title('VMAT2 by Region', fontsize=10, fontweight='bold')

    for i, fc in enumerate(vmat2_fc):
        inset_ax.text(fc + 0.01, i, f'{fc:.2f}x', fontsize=9, fontweight='bold', va='center')

    fig.suptitle('VMAT2 Dysfunction → DMT Availability Impact', fontsize=14, fontweight='bold', y=0.97)
    save_fig(fig, 'fig03_vmat2_dmt')


# ============================================================================
# FIGURE 4: Genetic Enrichment — clean bar chart with -log10(p)
# ============================================================================

def fig4():
    fig, ax = plt.subplots(figsize=(8, 6))

    # Build gene data
    gene_data = []
    for key, g in genes.items():
        if g['p_value'] is not None:
            gene_data.append({
                'symbol': g['gene_symbol'],
                'p_value': g['p_value'],
                'relevance': g['dmT_relevance'],
            })

    gene_data.sort(key=lambda x: x['p_value'])

    symbols = [g['symbol'] for g in gene_data]
    neg_log_p = [-np.log10(g['p_value']) for g in gene_data]
    relevance = [g['relevance'] for g in gene_data]

    # Color map
    rel_colors = {
        'CRITICAL': C['red'],
        'MODERATE': C['orange'],
        'WEAK': C['gray'],
    }
    colors = [rel_colors.get(r.split()[0].upper(), C['gray']) for r in relevance]

    x = np.arange(len(gene_data))
    w = 0.6

    bars = ax.bar(x, neg_log_p, w, color=colors, edgecolor='white', linewidth=0.5)

    # Threshold lines
    gw = -np.log10(5e-8)
    nom = -np.log10(0.05)
    ax.axhline(y=gw, color=C['red'], linestyle='--', lw=1.5, alpha=0.7,
               label=f'GWAS (p = 5×10⁻⁸)')
    ax.axhline(y=nom, color=C['orange'], linestyle=':', lw=1.2, alpha=0.7,
               label='Nominal (p = 0.05)')

    # P-value labels
    for i, (bar, g) in enumerate(zip(bars, gene_data)):
        height = bar.get_height()
        p_str = f'{g["p_value"]:.1e}'
        ax.text(i, height + 0.15, p_str, ha='center', fontsize=8, fontweight='bold')

    ax.set_xticks(x)
    ax.set_xticklabels(symbols, fontsize=10, rotation=15, ha='right')
    ax.set_ylabel('-log₁₀(p-value)', fontsize=11)
    ax.legend(loc='upper right', frameon=False, fontsize=9)
    ax.set_ylim(0, max(neg_log_p) * 1.2)

    # Summary box
    summary = (f"Pathway enrichment: p = {pathway_score['p_value_permutation']:.4f}\n"
               f"Percentile: {pathway_score['percentile']*100:.1f}%")
    ax.text(0.01, 0.95, summary, transform=ax.transAxes, fontsize=9,
            verticalalignment='top',
            bbox=dict(boxstyle='round,pad=0.4', facecolor='#E8F5E9', edgecolor=C['green'], lw=1.5))

    # Legend for relevance colors
    for i, (label, color) in enumerate(rel_colors.items()):
        ax.plot([], [], 'o', color=color, markersize=8,
                label=f'{label}' if i < 2 else f'{label}')
    ax.legend(handles=[], loc='lower right', frameon=False)
    # Manual legend
    ax.text(0.99, 0.02, 'CRITICAL = direct DMT role\nMODERATE = indirect\nWEAK = minimal',
            ha='right', va='bottom', fontsize=8, transform=ax.transAxes,
            bbox=dict(boxstyle='round,pad=0.3', facecolor=C['light']))

    fig.suptitle('DMT Pathway Gene Enrichment for Schizophrenia Risk', fontsize=14, fontweight='bold', y=0.97)
    save_fig(fig, 'fig04_genetics')


# ============================================================================
# FIGURE 5: NMDA-DMT Interaction — clean dual panel
# ============================================================================

def fig5():
    fig = plt.figure(figsize=(10, 4.5))
    gs = GridSpec(1, 2, figure=fig, wspace=0.2)

    # Left: NMDA function comparison
    ax1 = fig.add_subplot(gs[0, 0])

    models = ['Direct\nInhibition', 'Tonic\nMaintenance', 'Combined']
    model_keys = ['direct_inhibition_model', 'tonic_maintenance_model', 'combined_estimate']
    diff_keys = ['direct_model', 'tonic_model', 'combined']

    ctrl = np.array([nmda['control']['nmda_function'][k] for k in model_keys])
    scz = np.array([nmda['schizophrenia']['nmda_function'][k] for k in model_keys])
    diffs = [nmda['nmda_function_change'][k] for k in diff_keys]

    x = np.arange(len(models))
    w = 0.3

    bars1 = ax1.bar(x - w/2, ctrl, w, label='Control', color=C['blue'],
                    edgecolor='white', linewidth=0.5)
    bars2 = ax1.bar(x + w/2, scz, w, label='Schizophrenia', color=C['red'],
                    edgecolor='white', linewidth=0.5)

    for i, d in enumerate(diffs):
        color = C['red'] if d < 0 else C['green']
        pos = scz[i] - 0.012 if d < 0 else ctrl[i] + 0.012
        ax1.text(x[i], pos, f'{d:+.3f}', ha='center', fontsize=10, fontweight='bold', color=color)

    ax1.set_xticks(x)
    ax1.set_xticklabels(models, fontsize=10)
    ax1.set_ylabel('NMDA Function (relative)', fontsize=11)
    ax1.set_ylim(0.85, 1.05)
    ax1.axhline(y=1.0, color=C['gray'], linestyle='--', lw=1, alpha=0.5)
    ax1.legend(loc='upper right', frameon=False, fontsize=9)

    # Right: Inhibition curve
    ax2 = fig.add_subplot(gs[0, 1])

    concentrations = np.array(nmda['inhibition_curve']['concentrations'])
    inhibition = np.array(nmda['inhibition_curve']['inhibition']) * 100

    ax2.plot(concentrations, inhibition, color=C['purple'], lw=2.5,
             label='DMT → NMDA inhibition')

    # Mark control and schizophrenia levels
    ctrl_dmt = nmda['control']['dmt_level']
    scz_dmt = nmda['schizophrenia']['dmt_level']

    ax2.axvline(x=ctrl_dmt, color=C['blue'], linestyle='--', lw=1.5,
                label=f'Control ({ctrl_dmt}x)')
    ax2.axvline(x=scz_dmt, color=C['red'], linestyle='--', lw=1.5,
                label=f'Schizophrenia ({scz_dmt}x)')

    ax2.set_xlabel('DMT Concentration (relative)', fontsize=10)
    ax2.set_ylabel('NMDA Inhibition (%)', fontsize=10)
    ax2.set_xlim(0, 5)
    ax2.set_ylim(0, 30)
    ax2.legend(loc='upper left', fontsize=8, frameon=False)

    fig.suptitle('NMDA-DMT Cross-Modulation', fontsize=14, fontweight='bold', y=0.97)
    save_fig(fig, 'fig05_nmda')


# ============================================================================
# FIGURE 6: Bayesian Integration — clean forest plot + posterior
# ============================================================================

def fig6():
    fig = plt.figure(figsize=(10, 7))
    gs = GridSpec(2, 2, figure=fig, height_ratios=[1.3, 1], hspace=0.3, wspace=0.3)

    evidence_names = bayes['posterior']['evidence_names']
    evidence_labels = {
        'tryptamine_reduction': 'Tryptamine\nReduction',
        '5ht2a_upregulation': '5-HT2A\nUpregulation',
        '5ht2a_symptom_correlation': '5-HT2A ×\nSymptoms',
        'vmAT2_reduction': 'VMAT2\nReduction',
        'genetic_enrichment': 'Genetic\nEnrichment',
        'nmda_interaction': 'NMDA\nInteraction',
        'inflammatory_shunt': 'Inflammatory\nShunt',
    }

    # ── Top: Forest plot of likelihood ratios ──
    ax1 = fig.add_subplot(gs[0, :])

    y_pos = np.arange(len(evidence_names))
    lrs = [likelihoods[name]['likelihood_ratio_involved_not'] for name in evidence_names]

    # Color by direction
    lr_colors = []
    for lr in lrs:
        if lr > 1.5:
            lr_colors.append(C['green'])
        elif lr > 1.0:
            lr_colors.append(C['orange'])
        else:
            lr_colors.append(C['red'])

    # Draw forest plot
    for i, (name, lr) in enumerate(zip(evidence_names, lrs)):
        label = evidence_labels.get(name, name)
        ax1.plot(lr, i, 'o', color=lr_colors[i], markersize=9, zorder=3)
        ax1.text(-0.15, i, label, ha='right', va='center', fontsize=9.5, fontweight='bold')
        ax1.text(lr + 0.12, i, f'LR = {lr:.2f}', va='center', fontsize=8.5,
                color=lr_colors[i], fontweight='bold')

    # Neutral line
    ax1.axvline(x=1.0, color=C['dark'], linestyle='--', lw=1.5, alpha=0.5)
    ax1.set_yticks([])
    ax1.set_xlabel('Likelihood Ratio (supports DMT →)', fontsize=11)
    ax1.set_xlim(0, max(lrs) * 1.4)

    # Color legend
    for i, (label, color) in enumerate([('Supportive', C['green']),
                                         ('Weak', C['orange']),
                                         ('Against', C['red'])]):
        ax1.plot([], [], 'o', color=color, markersize=8, label=label)
    ax1.legend(loc='lower right', frameon=False, fontsize=8)

    # ── Bottom: Posterior donut + sensitivity ──
    ax2 = fig.add_subplot(gs[1, 0])

    posteriors = [posterior['dmt_involved'], posterior['dmt_modest_role'], posterior['dmt_not_involved']]
    labels = [f'DMT Involved\n{posterior["dmt_involved"]:.0%}',
              f'Modest Role\n{posterior["dmt_modest_role"]:.0%}',
              f'Not Involved\n{posterior["dmt_not_involved"]:.0%}']
    colors = [C['blue'], C['orange'], C['red']]

    wedges = ax2.pie(posteriors, labels=labels, colors=colors,
                                    startangle=90, pctdistance=0.82,
                                    wedgeprops=dict(width=0.45, edgecolor='white', lw=1.5))

    ax2.text(0, 0, f'{posterior["dmt_involved"]:.0%}',
            ha='center', va='center', fontsize=20, fontweight='bold', color=C['blue'])
    ax2.set_title('Posterior Distribution', fontsize=11, fontweight='bold')

    ax3 = fig.add_subplot(gs[1, 1])

    pri_keys = ['prior_0.01', 'prior_0.05', 'prior_0.1', 'prior_0.2', 'prior_0.3', 'prior_0.5']
    pri_involved = [sensitivity[k]['dmt_involved'] for k in pri_keys]
    pri_modest = [sensitivity[k]['dmt_modest_role'] for k in pri_keys]
    pri_not = [sensitivity[k]['dmt_not_involved'] for k in pri_keys]

    x = np.arange(6)
    w = 0.25
    ax3.bar(x - w, pri_involved, w, label='Involved', color=C['blue'], edgecolor='white', lw=0.5)
    ax3.bar(x, pri_modest, w, label='Modest', color=C['orange'], edgecolor='white', lw=0.5)
    ax3.bar(x + w, pri_not, w, label='Not', color=C['red'], edgecolor='white', lw=0.5)

    ax3.set_xticks(x)
    ax3.set_xticklabels(['0.01', '0.05', '0.1', '0.2', '0.3', '0.5'], fontsize=9)
    ax3.set_ylabel('Posterior', fontsize=10)
    ax3.set_ylim(0, 1.0)
    ax3.set_title('Sensitivity to Prior', fontsize=11, fontweight='bold')
    ax3.legend(loc='upper left', fontsize=8, frameon=False)

    # Stability annotation
    ax3.text(0.5, 0.92, 'Posterior robust across priors',
            ha='center', fontsize=9, style='italic', transform=ax3.transAxes,
            bbox=dict(boxstyle='round,pad=0.3', facecolor=C['light'], edgecolor=C['blue'], lw=1))

    fig.suptitle('Bayesian Integration: 7 Evidence Lines → P(DMT Involved) = 85%',
                fontsize=14, fontweight='bold', y=0.97)
    save_fig(fig, 'fig06_bayesian')


# ============================================================================
# Run
# ============================================================================

if __name__ == '__main__':
    print("Generating clean figures...")
    print()
    print("[1/6] Tryptophan flux")
    fig1()
    print("[2/6] 5-HT2A density")
    fig2()
    print("[3/6] VMAT2 impact")
    fig3()
    print("[4/6] Genetics")
    fig4()
    print("[5/6] NMDA-DMT")
    fig5()
    print("[6/6] Bayesian")
    fig6()
    print("\nDone. All figures in figures/")