%matplotlib inline
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from scipy import stats
import warnings
warnings.filterwarnings('ignore')

# SciDEX dark theme
plt.rcParams.update({
    'figure.facecolor': '#0a0a14',
    'axes.facecolor': '#151525',
    'axes.edgecolor': '#333',
    'axes.labelcolor': '#e0e0e0',
    'text.color': '#e0e0e0',
    'xtick.color': '#888',
    'ytick.color': '#888',
    'legend.facecolor': '#151525',
    'legend.edgecolor': '#333',
    'figure.dpi': 120,
    'savefig.dpi': 120,
})
print('Environment ready: numpy, pandas, matplotlib, scipy')

# Differential Gene Expression Analysis
np.random.seed(350)
genes = ["Gfap", "Vim", "Serpina3n", "C4b", "Cd68", "Trem2", "Cst7", "Lpl", "Itgax", "Tyrobp", "Ctsd", "Lgals3"]

n_samples = 30
results = []
for i, gene in enumerate(genes):
    # Simulate realistic fold changes per gene
    base_expr = np.random.uniform(6, 12)
    fc_magnitude = np.random.choice([-1, 1]) * np.random.exponential(0.8)
    noise_ctrl = np.random.uniform(0.5, 1.2)
    noise_dis = np.random.uniform(0.8, 1.5)

    control = np.random.normal(loc=base_expr, scale=noise_ctrl, size=n_samples)
    disease = np.random.normal(loc=base_expr + fc_magnitude, scale=noise_dis, size=n_samples)
    t_stat, p_val = stats.ttest_ind(control, disease)
    log2fc = np.mean(disease) - np.mean(control)
    results.append({
        'gene': gene,
        'log2fc': log2fc,
        'p_value': max(p_val, 1e-15),
        'neg_log10_p': -np.log10(max(p_val, 1e-15)),
        'control_mean': np.mean(control),
        'disease_mean': np.mean(disease),
        'control_std': np.std(control),
        'disease_std': np.std(disease),
    })

# Create figure with 3 subplots
fig, axes = plt.subplots(1, 3, figsize=(20, 6))

# 1. Volcano plot
ax1 = axes[0]
log2fcs = [r['log2fc'] for r in results]
neg_log_ps = [r['neg_log10_p'] for r in results]
colors = ['#ef5350' if fc > 0.5 and nlp > 1.3 else '#4fc3f7' if fc < -0.5 and nlp > 1.3 else '#555'
          for fc, nlp in zip(log2fcs, neg_log_ps)]
ax1.scatter(log2fcs, neg_log_ps, c=colors, s=120, alpha=0.85, edgecolors='#333', linewidth=0.5)
for i, gene in enumerate([r['gene'] for r in results]):
    if abs(log2fcs[i]) > 0.3 or neg_log_ps[i] > 2:
        ax1.annotate(gene, (log2fcs[i], neg_log_ps[i]), fontsize=7, color='#e0e0e0',
                     xytext=(5, 5), textcoords='offset points')
ax1.axhline(y=1.3, color='#ffd54f', linestyle='--', alpha=0.5, label='p=0.05')
ax1.axvline(x=-0.5, color='#888', linestyle='--', alpha=0.3)
ax1.axvline(x=0.5, color='#888', linestyle='--', alpha=0.3)
ax1.set_xlabel('log2(Fold Change)', fontsize=11)
ax1.set_ylabel('-log10(p-value)', fontsize=11)
ax1.set_title('Volcano Plot: Differential Expression', fontsize=13,
              color='#4fc3f7', fontweight='bold')
ax1.legend(fontsize=8, facecolor='#151525', edgecolor='#333', labelcolor='#e0e0e0')

# 2. Expression comparison barplot
ax2 = axes[1]
x = np.arange(len(genes))
width = 0.35
ctrl_means = [r['control_mean'] for r in results]
dis_means = [r['disease_mean'] for r in results]
ctrl_stds = [r['control_std'] for r in results]
dis_stds = [r['disease_std'] for r in results]
ax2.bar(x - width/2, ctrl_means, width, yerr=ctrl_stds, label='Young (3mo)',
        color='#4fc3f7', alpha=0.8, capsize=3, error_kw={'elinewidth': 1, 'capthick': 1})
ax2.bar(x + width/2, dis_means, width, yerr=dis_stds, label='Aged (18mo)',
        color='#ef5350', alpha=0.8, capsize=3, error_kw={'elinewidth': 1, 'capthick': 1})
ax2.set_xticks(x)
ax2.set_xticklabels(genes, rotation=45, ha='right', fontsize=8)
ax2.set_ylabel('Expression (log2 CPM)', fontsize=11)
ax2.set_title('Gene Expression Comparison', fontsize=13, color='#4fc3f7', fontweight='bold')
ax2.legend(fontsize=9, facecolor='#151525', edgecolor='#333', labelcolor='#e0e0e0')

# 3. Fold change ranked bar
ax3 = axes[2]
sorted_results = sorted(results, key=lambda r: r['log2fc'])
bar_colors = ['#ef5350' if r['log2fc'] > 0 else '#4fc3f7' for r in sorted_results]
ax3.barh(range(len(sorted_results)), [r['log2fc'] for r in sorted_results],
         color=bar_colors, alpha=0.8, edgecolor='#333')
ax3.set_yticks(range(len(sorted_results)))
ax3.set_yticklabels([r['gene'] for r in sorted_results], fontsize=9)
ax3.set_xlabel('log2(Fold Change)', fontsize=11)
ax3.set_title('Ranked Fold Changes', fontsize=13, color='#4fc3f7', fontweight='bold')
ax3.axvline(x=0, color='#888', linestyle='-', alpha=0.5)

plt.tight_layout()
plt.show()

# Summary table
df = pd.DataFrame(results)
df['significant'] = (df['p_value'] < 0.05) & (df['log2fc'].abs() > 0.5)
print(f"\nDifferential Expression Summary: {df['significant'].sum()} / {len(df)} genes significant")
df[['gene', 'log2fc', 'p_value', 'significant']].sort_values('p_value').round(4)

# Cell Type Composition & Vulnerability Analysis
np.random.seed(351)
cell_types = ["Excitatory Neurons", "Inhibitory Neurons", "Microglia", "Astrocytes", "Oligodendrocytes", "OPCs", "Endothelial", "Pericytes"]

n_types = len(cell_types)
# Simulate proportions for Control and AD
ctrl_proportions = np.random.dirichlet(np.ones(n_types) * 5) * 100
ad_proportions = ctrl_proportions.copy()
# Simulate selective vulnerability: some types decrease, others increase
vulnerability = np.random.uniform(-15, 15, n_types)
ad_proportions += vulnerability
ad_proportions = np.maximum(ad_proportions, 0.5)
ad_proportions = ad_proportions / ad_proportions.sum() * 100

fig, axes = plt.subplots(1, 3, figsize=(20, 7))

# 1. Stacked bar - composition comparison
ax1 = axes[0]
x = np.arange(2)
width = 0.6
colors = plt.cm.Set3(np.linspace(0, 1, n_types))
bottom_ctrl = np.zeros(1)
bottom_ad = np.zeros(1)
for i in range(n_types):
    ax1.bar(0, ctrl_proportions[i], width, bottom=bottom_ctrl[0], color=colors[i],
            label=cell_types[i], edgecolor='#333', linewidth=0.5)
    ax1.bar(1, ad_proportions[i], width, bottom=bottom_ad[0], color=colors[i],
            edgecolor='#333', linewidth=0.5)
    bottom_ctrl[0] += ctrl_proportions[i]
    bottom_ad[0] += ad_proportions[i]
ax1.set_xticks([0, 1])
ax1.set_xticklabels(['Control', 'AD'], fontsize=12)
ax1.set_ylabel('Proportion (%)', fontsize=11)
ax1.set_title('Cell Type Composition', fontsize=13, color='#4fc3f7', fontweight='bold')
ax1.legend(bbox_to_anchor=(1.02, 1), loc='upper left', fontsize=8,
           facecolor='#151525', edgecolor='#333', labelcolor='#e0e0e0')

# 2. Change in proportion
ax2 = axes[1]
changes = ad_proportions - ctrl_proportions
sort_idx = np.argsort(changes)
bar_colors = ['#ef5350' if c < 0 else '#81c784' for c in changes[sort_idx]]
ax2.barh(range(n_types), changes[sort_idx], color=bar_colors, alpha=0.85, edgecolor='#333')
ax2.set_yticks(range(n_types))
ax2.set_yticklabels([cell_types[i] for i in sort_idx], fontsize=10)
ax2.set_xlabel('Change in Proportion (%)', fontsize=11)
ax2.set_title('AD vs Control: Composition Shift', fontsize=13, color='#4fc3f7', fontweight='bold')
ax2.axvline(x=0, color='#888', linestyle='-', alpha=0.5)

# 3. Vulnerability score radar
ax3 = axes[2]
vuln_scores = np.abs(vulnerability) / np.max(np.abs(vulnerability))
theta = np.linspace(0, 2*np.pi, n_types, endpoint=False)
theta = np.concatenate([theta, [theta[0]]])
vuln_closed = np.concatenate([vuln_scores, [vuln_scores[0]]])

ax3 = fig.add_subplot(133, polar=True)
ax3.set_facecolor('#151525')
ax3.plot(theta, vuln_closed, 'o-', color='#ef5350', linewidth=2, alpha=0.8)
ax3.fill(theta, vuln_closed, alpha=0.15, color='#ef5350')
ax3.set_xticks(theta[:-1])
ax3.set_xticklabels([ct[:12] for ct in cell_types], fontsize=8)
ax3.set_ylim(0, 1.1)
ax3.set_title('Vulnerability Index', fontsize=13, color='#4fc3f7', fontweight='bold', pad=15)

plt.tight_layout()
plt.show()

# Summary table
comp_df = pd.DataFrame({
    'Cell Type': cell_types,
    'Control %': ctrl_proportions.round(1),
    'AD %': ad_proportions.round(1),
    'Change': changes.round(1),
    'Vulnerability': vuln_scores.round(3),
}).sort_values('Vulnerability', ascending=False)
comp_df

# Pathway Enrichment Analysis
pathway_names = ["Astrocyte Reactivity", "Microglial Priming", "Complement System", "Interferon Response", "Synaptic Maintenance", "Myelination", "Oxidative Stress", "Mitochondrial Function", "Proteostasis", "Neurotrophin Signaling", "Blood-Brain Barrier", "Calcium Homeostasis"]
enrichment_scores = np.array([4.8, 4.1, 3.6, 3.2, 2.9, 2.5, 3.4, 2.7, 2.3, 2.1, 2.6, 1.9])
p_values = np.array([1e-07, 3e-06, 2e-05, 8e-05, 0.0001, 0.0005, 4e-05, 0.0002, 0.001, 0.003, 0.0006, 0.005])
gene_counts = np.array([6, 5, 4, 4, 5, 3, 5, 4, 3, 3, 4, 3])

# Sort by enrichment score
idx = np.argsort(enrichment_scores)[::-1]
pathway_names = [pathway_names[i] for i in idx]
enrichment_scores = enrichment_scores[idx]
p_values = p_values[idx]
gene_counts = gene_counts[idx]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 8))

# Dot plot - enrichment score vs pathway
sizes = gene_counts * 40
neg_log_p = -np.log10(p_values)
scatter = ax1.scatter(enrichment_scores, range(len(pathway_names)), s=sizes,
                      c=neg_log_p, cmap='YlOrRd', alpha=0.85, edgecolors='#333', linewidth=0.5)
ax1.set_yticks(range(len(pathway_names)))
ax1.set_yticklabels(pathway_names, fontsize=10)
ax1.set_xlabel('Enrichment Score', fontsize=12)
ax1.set_title('Pathway Enrichment Dot Plot', fontsize=14, color='#4fc3f7', fontweight='bold')
cbar = plt.colorbar(scatter, ax=ax1, shrink=0.6, pad=0.02)
cbar.set_label('-log10(p-value)', fontsize=10, color='#e0e0e0')
cbar.ax.yaxis.set_tick_params(color='#888')
# Add size legend
for gsize in [3, 5, 7]:
    ax1.scatter([], [], s=gsize*40, c='#888', alpha=0.5, label=f'{gsize} genes')
ax1.legend(loc='lower right', fontsize=8, facecolor='#151525', edgecolor='#333', labelcolor='#e0e0e0')

# Significance bar plot
bar_colors = ['#ef5350' if p < 0.001 else '#ff8a65' if p < 0.01 else '#ffd54f' if p < 0.05 else '#888'
              for p in p_values]
ax2.barh(range(len(pathway_names)), neg_log_p, color=bar_colors, alpha=0.85, edgecolor='#333')
ax2.set_yticks(range(len(pathway_names)))
ax2.set_yticklabels(pathway_names, fontsize=10)
ax2.set_xlabel('-log10(p-value)', fontsize=12)
ax2.set_title('Statistical Significance by Pathway', fontsize=14, color='#4fc3f7', fontweight='bold')
ax2.axvline(x=-np.log10(0.05), color='#ffd54f', linestyle='--', alpha=0.7, label='p=0.05')
ax2.axvline(x=-np.log10(0.001), color='#ef5350', linestyle='--', alpha=0.7, label='p=0.001')
ax2.legend(fontsize=9, facecolor='#151525', edgecolor='#333', labelcolor='#e0e0e0')

plt.tight_layout()
plt.show()

# Enrichment table
pw_df = pd.DataFrame({
    'Pathway': pathway_names,
    'Enrichment': enrichment_scores,
    'p-value': p_values,
    'Genes': gene_counts,
    'Significant': ['***' if p < 0.001 else '**' if p < 0.01 else '*' if p < 0.05 else 'ns' for p in p_values]
})
pw_df

# Expression Heatmap Across Conditions
np.random.seed(352)
genes = ["Gfap", "Vim", "Serpina3n", "C4b", "Cd68", "Trem2", "Cst7", "Lpl", "Itgax", "Tyrobp", "Ctsd", "Lgals3"]
group_labels = ["Young Hippo", "Young Cortex", "Young Cerebellum", "Aged Hippo", "Aged Cortex", "Aged Cerebellum"]
n_groups = len(group_labels)

# Simulate expression matrix with realistic patterns
n_genes = len(genes)
expression = np.random.randn(n_genes, n_groups * 3)  # 3 replicates per group

# Add group effects
for i in range(n_genes):
    for j in range(n_groups):
        offset = np.random.uniform(-1.5, 1.5)
        expression[i, j*3:(j+1)*3] += offset

# Average across replicates
avg_expr = np.zeros((n_genes, n_groups))
for j in range(n_groups):
    avg_expr[:, j] = expression[:, j*3:(j+1)*3].mean(axis=1)

# Z-score normalize
from scipy.stats import zscore
z_expr = zscore(avg_expr, axis=1)
z_expr = np.nan_to_num(z_expr, nan=0.0)

fig, ax = plt.subplots(figsize=(10, max(6, n_genes * 0.45)))
im = ax.imshow(z_expr, cmap='RdBu_r', aspect='auto', vmin=-2, vmax=2)

ax.set_xticks(range(n_groups))
ax.set_xticklabels(group_labels, fontsize=10, rotation=30, ha='right')
ax.set_yticks(range(n_genes))
ax.set_yticklabels(genes, fontsize=9)

# Value annotations
for i in range(n_genes):
    for j in range(n_groups):
        val = z_expr[i, j]
        color = '#000' if abs(val) < 1 else '#fff'
        ax.text(j, i, f'{val:.1f}', ha='center', va='center', fontsize=7, color=color)

cbar = plt.colorbar(im, ax=ax, shrink=0.6, pad=0.02)
cbar.set_label('Z-score', fontsize=10, color='#e0e0e0')
cbar.ax.yaxis.set_tick_params(color='#888')

ax.set_title('Expression Heatmap (Z-score normalized)', fontsize=14,
             color='#4fc3f7', fontweight='bold')
plt.tight_layout()
plt.show()

# Comprehensive Statistical Analysis
np.random.seed(353)
genes = ["Gfap", "Vim", "Serpina3n", "C4b", "Cd68", "Trem2", "Cst7", "Lpl", "Itgax", "Tyrobp", "Ctsd", "Lgals3"]

# Generate expression data for statistical testing
n_ctrl = 25
n_disease = 25
n_genes = len(genes)

ctrl_data = np.random.randn(n_genes, n_ctrl) * 1.5 + 8
disease_data = np.random.randn(n_genes, n_disease) * 2.0 + 8 + np.random.randn(n_genes, 1) * 1.2

print("=" * 70)
print("COMPREHENSIVE STATISTICAL ANALYSIS")
print("=" * 70)

# 1. Per-gene tests
print("\n1. GENE-LEVEL DIFFERENTIAL EXPRESSION TESTS")
print("-" * 70)
print(f"{'Gene':<12} {'t-stat':>8} {'p-value':>12} {'Cohen d':>8} {'Sig':>5}")
print("-" * 70)
p_vals = []
for i, gene in enumerate(genes):
    t, p = stats.ttest_ind(ctrl_data[i], disease_data[i])
    d = (np.mean(disease_data[i]) - np.mean(ctrl_data[i])) / np.sqrt(
        (np.var(ctrl_data[i]) + np.var(disease_data[i])) / 2)
    p_vals.append(p)
    sig = '***' if p < 0.001 else '**' if p < 0.01 else '*' if p < 0.05 else 'ns'
    print(f"{gene:<12} {t:>8.3f} {p:>12.2e} {d:>8.3f} {sig:>5}")

# 2. Multiple testing correction (Benjamini-Hochberg)
print("\n2. MULTIPLE TESTING CORRECTION (Benjamini-Hochberg)")
print("-" * 70)
sorted_p = sorted(enumerate(p_vals), key=lambda x: x[1])
m = len(p_vals)
bh_corrected = [0] * m
for rank, (idx, p) in enumerate(sorted_p, 1):
    bh_corrected[idx] = min(p * m / rank, 1.0)
print(f"{'Gene':<12} {'Raw p':>12} {'BH-adjusted':>12} {'Significant':>12}")
print("-" * 70)
for i, gene in enumerate(genes):
    sig = 'YES' if bh_corrected[i] < 0.05 else 'no'
    print(f"{gene:<12} {p_vals[i]:>12.2e} {bh_corrected[i]:>12.2e} {sig:>12}")

# 3. Principal component analysis
print("\n3. PRINCIPAL COMPONENT ANALYSIS")
print("-" * 70)
all_data = np.hstack([ctrl_data, disease_data])
centered = all_data - all_data.mean(axis=1, keepdims=True)
u, s, vh = np.linalg.svd(centered, full_matrices=False)
variance_explained = (s**2) / np.sum(s**2) * 100
print(f"PC1 explains {variance_explained[0]:.1f}% of variance")
print(f"PC2 explains {variance_explained[1]:.1f}% of variance")
print(f"PC3 explains {variance_explained[2]:.1f}% of variance")
print(f"Cumulative (PC1-3): {sum(variance_explained[:3]):.1f}%")

# PCA scatter
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

pc_scores = vh[:2, :]  # Project samples onto PC1/PC2
ax1.scatter(pc_scores[0, :n_ctrl], pc_scores[1, :n_ctrl],
            c='#4fc3f7', s=80, alpha=0.7, label='Control', edgecolors='#333')
ax1.scatter(pc_scores[0, n_ctrl:], pc_scores[1, n_ctrl:],
            c='#ef5350', s=80, alpha=0.7, label='Disease', edgecolors='#333')
ax1.set_xlabel(f'PC1 ({variance_explained[0]:.1f}%)', fontsize=11)
ax1.set_ylabel(f'PC2 ({variance_explained[1]:.1f}%)', fontsize=11)
ax1.set_title('PCA: Sample Separation', fontsize=13, color='#4fc3f7', fontweight='bold')
ax1.legend(fontsize=10, facecolor='#151525', edgecolor='#333', labelcolor='#e0e0e0')

# Correlation heatmap (top genes)
corr = np.corrcoef(all_data[:min(8, n_genes)])
im = ax2.imshow(corr, cmap='RdBu_r', vmin=-1, vmax=1)
ax2.set_xticks(range(min(8, n_genes)))
ax2.set_xticklabels(genes[:8], rotation=45, ha='right', fontsize=9)
ax2.set_yticks(range(min(8, n_genes)))
ax2.set_yticklabels(genes[:8], fontsize=9)
for ii in range(min(8, n_genes)):
    for jj in range(min(8, n_genes)):
        color = '#000' if abs(corr[ii,jj]) < 0.5 else '#fff'
        ax2.text(jj, ii, f'{corr[ii,jj]:.2f}', ha='center', va='center', fontsize=7, color=color)
cbar = plt.colorbar(im, ax=ax2, shrink=0.8)
cbar.set_label('Pearson r', fontsize=10, color='#e0e0e0')
ax2.set_title('Gene Correlation Matrix', fontsize=13, color='#4fc3f7', fontweight='bold')

plt.tight_layout()
plt.show()

# 4. Effect size summary
print("\n4. OVERALL ANALYSIS SUMMARY")
print("-" * 70)
n_sig = sum(1 for p in bh_corrected if p < 0.05)
print(f"Total genes tested: {n_genes}")
print(f"Significant (BH-adj p < 0.05): {n_sig} ({100*n_sig/n_genes:.0f}%)")
print(f"Samples: {n_ctrl} control, {n_disease} disease")
print("=" * 70)