#!/usr/bin/env python3
"""
Step 3: Statistical Significance Testing for GLP-1 Adverse Event Analysis

This script performs pairwise statistical comparisons between drugs for each
adverse event category, calculating Odds Ratios, 95% CIs, and p-values using
Fisher's Exact Test or Chi-Square test.

Outputs:
- results/03_statistical_results.json: Full statistical results
- results/03_stats_summary.txt: Human-readable interpretation
- figures/or_forest_plot.png: Forest plot of Odds Ratios
- figures/significance_heatmap.png: Heatmap of -log10(p-values)
"""

import json
import numpy as np
import pandas as pd
from scipy import stats
from itertools import combinations
from pathlib import Path
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# Configure paths
BASE_DIR = Path("/app/sandbox/session_20260203_091322_5981a70f834a")
RESULTS_DIR = BASE_DIR / "results"
FIGURES_DIR = BASE_DIR / "figures"
INPUT_FILE = RESULTS_DIR / "02_analysis_results.json"

print(f"=" * 60)
print("Step 3: Statistical Significance Testing")
print(f"=" * 60)
print(f"Start time: {datetime.now().isoformat()}")
print()

# Load data from Step 2
print("Loading data from Step 2...")
with open(INPUT_FILE, 'r') as f:
    data = json.load(f)

cross_tab = data['cross_tabulation']
summary_by_drug = data['summary_by_drug']
categories = data['metadata']['categories_analyzed']
drugs = data['metadata']['drugs_analyzed']

print(f"  Drugs: {drugs}")
print(f"  Categories: {categories}")
print(f"  Total records: {data['metadata']['total_records_processed']}")
print()


def calculate_odds_ratio_fisher(a, b, c, d):
    """
    Calculate Odds Ratio and 95% CI from 2x2 contingency table.

    Contingency table:
                Event+    Event-
    Drug A      a         b
    Drug B      c         d

    OR = (a*d) / (b*c)

    Uses log transformation for CI calculation.
    """
    # Add small constant to avoid division by zero
    epsilon = 0.5
    a_adj = a + epsilon if (a == 0 or b == 0 or c == 0 or d == 0) else a
    b_adj = b + epsilon if (a == 0 or b == 0 or c == 0 or d == 0) else b
    c_adj = c + epsilon if (a == 0 or b == 0 or c == 0 or d == 0) else c
    d_adj = d + epsilon if (a == 0 or b == 0 or c == 0 or d == 0) else d

    # Calculate OR
    or_value = (a_adj * d_adj) / (b_adj * c_adj)

    # Calculate 95% CI using log transformation
    log_or = np.log(or_value)
    se_log_or = np.sqrt(1/a_adj + 1/b_adj + 1/c_adj + 1/d_adj)

    ci_lower = np.exp(log_or - 1.96 * se_log_or)
    ci_upper = np.exp(log_or + 1.96 * se_log_or)

    return or_value, ci_lower, ci_upper


def perform_fisher_or_chi2(table):
    """
    Perform Fisher's Exact Test if any cell < 5, otherwise Chi-Square.
    Returns p-value and test name used.
    """
    min_expected = np.min(stats.chi2_contingency(table, correction=False)[3])

    if min_expected < 5:
        # Use Fisher's Exact Test
        _, p_value = stats.fisher_exact(table)
        test_used = "Fisher's Exact"
    else:
        # Use Chi-Square
        chi2, p_value, _, _ = stats.chi2_contingency(table, correction=True)
        test_used = "Chi-Square (Yates)"

    return p_value, test_used


# Perform pairwise comparisons
print("Performing pairwise comparisons...")
print("-" * 60)

drug_pairs = list(combinations(drugs, 2))
results = {
    "metadata": {
        "analysis_date": datetime.now().isoformat(),
        "comparisons_performed": len(drug_pairs) * len(categories),
        "alpha_level": 0.05,
        "multiple_testing_correction": "Bonferroni"
    },
    "pairwise_results": [],
    "summary_matrix": {}
}

# Total number of tests for Bonferroni correction
n_tests = len(drug_pairs) * len(categories)
bonferroni_alpha = 0.05 / n_tests

print(f"Number of drug pairs: {len(drug_pairs)}")
print(f"Number of event categories: {len(categories)}")
print(f"Total tests: {n_tests}")
print(f"Bonferroni-corrected alpha: {bonferroni_alpha:.6f}")
print()

all_results = []
comparison_count = 0

for drug_a, drug_b in drug_pairs:
    print(f"Comparing: {drug_a} vs {drug_b}")

    for category in categories:
        comparison_count += 1

        # Get counts for each drug in this category
        count_a_event = cross_tab[drug_a].get(category, 0)
        count_b_event = cross_tab[drug_b].get(category, 0)

        # Total reports for each drug
        total_a = summary_by_drug[drug_a]['total_reports']
        total_b = summary_by_drug[drug_b]['total_reports']

        # Not-event counts (other events for this drug)
        count_a_no_event = total_a - count_a_event
        count_b_no_event = total_b - count_b_event

        # Construct 2x2 contingency table
        # Rows: Drug A, Drug B
        # Cols: Event, No-Event
        table = np.array([
            [count_a_event, count_a_no_event],
            [count_b_event, count_b_no_event]
        ])

        # Calculate Odds Ratio and CI
        or_value, ci_lower, ci_upper = calculate_odds_ratio_fisher(
            count_a_event, count_a_no_event,
            count_b_event, count_b_no_event
        )

        # Perform statistical test
        try:
            p_value, test_used = perform_fisher_or_chi2(table)
        except Exception as e:
            # Fallback to Fisher's if chi2 fails
            _, p_value = stats.fisher_exact(table)
            test_used = "Fisher's Exact (fallback)"

        # Apply Bonferroni correction
        p_adjusted = min(p_value * n_tests, 1.0)
        is_significant_raw = p_value < 0.05
        is_significant_adjusted = p_adjusted < 0.05

        result = {
            "comparison": f"{drug_a} vs {drug_b}",
            "drug_a": drug_a,
            "drug_b": drug_b,
            "category": category,
            "contingency_table": {
                "drug_a_event": int(count_a_event),
                "drug_a_no_event": int(count_a_no_event),
                "drug_b_event": int(count_b_event),
                "drug_b_no_event": int(count_b_no_event)
            },
            "odds_ratio": float(round(or_value, 4)),
            "ci_95_lower": float(round(ci_lower, 4)),
            "ci_95_upper": float(round(ci_upper, 4)),
            "p_value_raw": float(round(p_value, 6)),
            "p_value_adjusted": float(round(p_adjusted, 6)),
            "test_used": test_used,
            "significant_raw": bool(is_significant_raw),
            "significant_adjusted": bool(is_significant_adjusted)
        }

        all_results.append(result)

        # Progress update
        if comparison_count % 3 == 0:
            print(f"  Processed {comparison_count}/{n_tests} comparisons...")

print()
results["pairwise_results"] = all_results

# Create summary matrices for visualization
print("Creating summary matrices...")
or_matrix = pd.DataFrame(index=categories, columns=[f"{a} vs {b}" for a, b in drug_pairs])
p_matrix = pd.DataFrame(index=categories, columns=[f"{a} vs {b}" for a, b in drug_pairs])

for r in all_results:
    comparison_label = r["comparison"]
    category = r["category"]
    or_matrix.loc[category, comparison_label] = r["odds_ratio"]
    p_matrix.loc[category, comparison_label] = r["p_value_raw"]

# Convert DataFrame values to native Python types for JSON serialization
def convert_to_native(obj):
    """Convert numpy types to native Python types."""
    if isinstance(obj, dict):
        return {k: convert_to_native(v) for k, v in obj.items()}
    elif isinstance(obj, (np.integer, np.floating)):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, (np.bool_,)):
        return bool(obj)
    return obj

results["summary_matrix"] = {
    "odds_ratios": convert_to_native(or_matrix.to_dict()),
    "p_values": convert_to_native(p_matrix.to_dict())
}

# Save results JSON
print("Saving statistical results...")
output_json = RESULTS_DIR / "03_statistical_results.json"
with open(output_json, 'w') as f:
    json.dump(results, f, indent=2)
print(f"  Saved: {output_json}")

# =============================================================================
# Generate Visualizations
# =============================================================================
print()
print("Generating visualizations...")
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# Set publication-quality style
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.size'] = 10
plt.rcParams['axes.linewidth'] = 0.8

# -----------------------------------------------------------------------------
# Forest Plot of Odds Ratios
# -----------------------------------------------------------------------------
print("  Creating forest plot...")

fig, ax = plt.subplots(figsize=(12, 8))

# Prepare data for forest plot
plot_data = []
for r in all_results:
    label = f"{r['category']}\n({r['drug_a']} vs {r['drug_b']})"
    plot_data.append({
        'label': label,
        'or': r['odds_ratio'],
        'ci_lower': r['ci_95_lower'],
        'ci_upper': r['ci_95_upper'],
        'p_value': r['p_value_raw'],
        'significant': r['significant_raw']
    })

# Sort by odds ratio for better visualization
plot_data = sorted(plot_data, key=lambda x: x['or'])

y_positions = np.arange(len(plot_data))
labels = [d['label'] for d in plot_data]
ors = [d['or'] for d in plot_data]
ci_lowers = [d['ci_lower'] for d in plot_data]
ci_uppers = [d['ci_upper'] for d in plot_data]
significants = [d['significant'] for d in plot_data]

# Calculate error bars (asymmetric)
errors_lower = [max(0.01, or_val - ci_lo) for or_val, ci_lo in zip(ors, ci_lowers)]
errors_upper = [ci_up - or_val for or_val, ci_up in zip(ors, ci_uppers)]

# Color coding: significant vs non-significant
colors = ['#D62728' if sig else '#1F77B4' for sig in significants]

# Plot
ax.errorbar(ors, y_positions, xerr=[errors_lower, errors_upper],
            fmt='o', capsize=4, capthick=1.5, markersize=8,
            ecolor='gray', elinewidth=1.5)

# Color the markers
for i, (x, y, color) in enumerate(zip(ors, y_positions, colors)):
    ax.scatter([x], [y], c=color, s=80, zorder=5)

# Add vertical line at OR = 1 (null effect)
ax.axvline(x=1, color='gray', linestyle='--', linewidth=1.5, alpha=0.7)

# Formatting
ax.set_yticks(y_positions)
ax.set_yticklabels(labels, fontsize=9)
ax.set_xlabel('Odds Ratio (95% CI)', fontsize=12, fontweight='bold')
ax.set_title('Forest Plot: Odds Ratios for Adverse Events by Drug Pair\n(Red = p < 0.05, Blue = p ≥ 0.05)',
             fontsize=13, fontweight='bold')

# Set x-axis to log scale for better visualization of ORs
ax.set_xscale('log')
ax.set_xlim(0.1, 20)

# Add grid
ax.grid(axis='x', alpha=0.3, linestyle='-')
ax.set_axisbelow(True)

# Add legend
from matplotlib.lines import Line2D
legend_elements = [
    Line2D([0], [0], marker='o', color='w', markerfacecolor='#D62728', markersize=10, label='Significant (p < 0.05)'),
    Line2D([0], [0], marker='o', color='w', markerfacecolor='#1F77B4', markersize=10, label='Not Significant'),
    Line2D([0], [0], color='gray', linestyle='--', label='Null Effect (OR = 1)')
]
ax.legend(handles=legend_elements, loc='lower right', fontsize=9)

plt.tight_layout()
forest_plot_path = FIGURES_DIR / "or_forest_plot.png"
plt.savefig(forest_plot_path, dpi=300, bbox_inches='tight', facecolor='white')
plt.close()
print(f"    Saved: {forest_plot_path}")

# -----------------------------------------------------------------------------
# Significance Heatmap
# -----------------------------------------------------------------------------
print("  Creating significance heatmap...")

fig, ax = plt.subplots(figsize=(10, 6))

# Convert p-values to -log10(p) for better visualization
p_matrix_numeric = p_matrix.astype(float)
neg_log_p = -np.log10(p_matrix_numeric + 1e-10)  # Add small value to avoid log(0)

# Create heatmap
import matplotlib.colors as mcolors

# Custom colormap: white to red (more significant = more red)
cmap = plt.cm.Reds

im = ax.imshow(neg_log_p.values, cmap=cmap, aspect='auto')

# Set ticks
ax.set_xticks(np.arange(len(neg_log_p.columns)))
ax.set_yticks(np.arange(len(neg_log_p.index)))
ax.set_xticklabels(neg_log_p.columns, fontsize=10, rotation=45, ha='right')
ax.set_yticklabels(neg_log_p.index, fontsize=10)

# Add colorbar
cbar = plt.colorbar(im, ax=ax, shrink=0.8)
cbar.set_label('-log10(p-value)', fontsize=11, fontweight='bold')

# Add text annotations with p-values and significance markers
for i in range(len(neg_log_p.index)):
    for j in range(len(neg_log_p.columns)):
        p_val = p_matrix_numeric.iloc[i, j]
        text_color = 'white' if neg_log_p.iloc[i, j] > 1.5 else 'black'

        if p_val < 0.001:
            text = '***'
        elif p_val < 0.01:
            text = '**'
        elif p_val < 0.05:
            text = '*'
        else:
            text = 'ns'

        ax.text(j, i, f'{p_val:.3f}\n{text}', ha='center', va='center',
                fontsize=9, color=text_color, fontweight='bold')

# Add significance threshold line annotation
ax.set_title('Statistical Significance Heatmap\n(-log10 p-value; * p<0.05, ** p<0.01, *** p<0.001, ns = not significant)',
             fontsize=12, fontweight='bold')
ax.set_xlabel('Drug Comparison', fontsize=11, fontweight='bold')
ax.set_ylabel('Adverse Event Category', fontsize=11, fontweight='bold')

plt.tight_layout()
heatmap_path = FIGURES_DIR / "significance_heatmap.png"
plt.savefig(heatmap_path, dpi=300, bbox_inches='tight', facecolor='white')
plt.close()
print(f"    Saved: {heatmap_path}")

# =============================================================================
# Generate Text Summary
# =============================================================================
print()
print("Generating statistical summary...")

summary_lines = [
    "=" * 70,
    "STATISTICAL SIGNIFICANCE TESTING SUMMARY",
    "GLP-1 Receptor Agonists Adverse Event Analysis",
    f"Analysis Date: {datetime.now().strftime('%Y-%m-%d %H:%M')}",
    "=" * 70,
    "",
    "METHODOLOGY",
    "-" * 70,
    f"• Total comparisons performed: {n_tests}",
    f"• Statistical tests: Fisher's Exact Test (when any expected cell < 5)",
    f"                     Chi-Square Test with Yates correction (otherwise)",
    f"• Significance level (alpha): 0.05",
    f"• Multiple testing correction: Bonferroni (adjusted alpha = {bonferroni_alpha:.6f})",
    "",
    "DRUGS COMPARED",
    "-" * 70,
]

for drug in drugs:
    n_reports = summary_by_drug[drug]['total_reports']
    summary_lines.append(f"• {drug}: {n_reports} total reports")

summary_lines.extend([
    "",
    "ADVERSE EVENT CATEGORIES",
    "-" * 70,
])

for cat in categories:
    cat_count = data['summary_by_category'].get(cat, 0)
    summary_lines.append(f"• {cat}: {cat_count} reports")

summary_lines.extend([
    "",
    "KEY FINDINGS",
    "-" * 70,
])

# Identify significant findings
significant_raw = [r for r in all_results if r['significant_raw']]
significant_adjusted = [r for r in all_results if r['significant_adjusted']]

if significant_raw:
    summary_lines.append(f"\nSignificant differences (p < 0.05, uncorrected):")
    for r in sorted(significant_raw, key=lambda x: x['p_value_raw']):
        direction = "higher" if r['odds_ratio'] > 1 else "lower"
        summary_lines.append(
            f"  • {r['category']}: {r['drug_a']} has {direction} odds vs {r['drug_b']}"
            f"\n    OR = {r['odds_ratio']:.2f} (95% CI: {r['ci_95_lower']:.2f}-{r['ci_95_upper']:.2f}), "
            f"p = {r['p_value_raw']:.4f}"
        )
else:
    summary_lines.append("\nNo significant differences found at p < 0.05 (uncorrected).")

if significant_adjusted:
    summary_lines.append(f"\nSignificant after Bonferroni correction (p < {bonferroni_alpha:.6f}):")
    for r in sorted(significant_adjusted, key=lambda x: x['p_value_adjusted']):
        direction = "higher" if r['odds_ratio'] > 1 else "lower"
        summary_lines.append(
            f"  • {r['category']}: {r['drug_a']} has {direction} odds vs {r['drug_b']}"
            f"\n    OR = {r['odds_ratio']:.2f}, adjusted p = {r['p_value_adjusted']:.4f}"
        )
else:
    summary_lines.append(f"\nNo significant differences remain after Bonferroni correction.")

summary_lines.extend([
    "",
    "DETAILED RESULTS BY COMPARISON",
    "-" * 70,
])

for drug_a, drug_b in drug_pairs:
    summary_lines.append(f"\n{drug_a} vs {drug_b}:")
    pair_results = [r for r in all_results if r['drug_a'] == drug_a and r['drug_b'] == drug_b]

    for r in pair_results:
        sig_marker = "*" if r['significant_raw'] else ""
        adj_marker = "†" if r['significant_adjusted'] else ""
        summary_lines.append(
            f"  {r['category']}: OR = {r['odds_ratio']:.3f} "
            f"(95% CI: {r['ci_95_lower']:.3f}-{r['ci_95_upper']:.3f}), "
            f"p = {r['p_value_raw']:.4f}{sig_marker}{adj_marker}"
        )

summary_lines.extend([
    "",
    "INTERPRETATION NOTES",
    "-" * 70,
    "• OR > 1: Higher odds of event in first drug compared to second",
    "• OR < 1: Lower odds of event in first drug compared to second",
    "• OR = 1: No difference between drugs",
    "• * = significant at p < 0.05 (uncorrected)",
    "• † = significant after Bonferroni correction",
    "",
    "LIMITATIONS",
    "-" * 70,
    "• FAERS is a spontaneous reporting system subject to reporting bias",
    "• Reporting rates do not equal incidence rates",
    "• Differences may reflect market exposure, prescribing patterns, or publicity",
    "• Small sample sizes for some categories may limit statistical power",
    "• Gastroparesis has very few reports (n=2), limiting meaningful comparison",
    "",
    "OUTPUT FILES",
    "-" * 70,
    "• results/03_statistical_results.json - Full statistical results",
    "• figures/or_forest_plot.png - Forest plot of Odds Ratios",
    "• figures/significance_heatmap.png - Heatmap of p-values",
    "=" * 70,
])

summary_text = "\n".join(summary_lines)

summary_path = RESULTS_DIR / "03_stats_summary.txt"
with open(summary_path, 'w') as f:
    f.write(summary_text)
print(f"  Saved: {summary_path}")

# Print summary to console
print()
print(summary_text)

print()
print("=" * 60)
print("Step 3 COMPLETE: Statistical Significance Testing")
print("=" * 60)
print(f"End time: {datetime.now().isoformat()}")
print()
print("Generated outputs:")
print(f"  1. {output_json}")
print(f"  2. {forest_plot_path}")
print(f"  3. {heatmap_path}")
print(f"  4. {summary_path}")