"""
GWAS and Burden Test Rank Aggregation
======================================

Demonstrates how integrating multiple ranking methods can correct for gene length
bias and improve gene prioritization. Shows a concrete example where a gene improves
from rank ~500 in GWAS-only analysis to top 50 when combined with burden test evidence.

Key Concepts:
- GWAS rankings suffer from length bias (longer genes = more variants = higher ranks)
- Burden tests are less biased but have lower power
- Rank aggregation combines both methods to leverage their complementary strengths

Dependencies: numpy, pandas, matplotlib
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import Tuple

# Set random seed for reproducibility
np.random.seed(42)


def generate_synthetic_genes(n_genes: int = 1000) -> pd.DataFrame:
    """
    Generate synthetic gene data with realistic properties.

    Args:
        n_genes: Number of genes to simulate

    Returns:
        DataFrame with gene properties including length, true importance,
        and simulated ranking signals
    """
    # Generate gene lengths (in kb) - log-normal distribution mimics real data
    gene_lengths = np.random.lognormal(mean=2.5, sigma=1.0, size=n_genes)
    gene_lengths = np.clip(gene_lengths, 1, 200)  # Realistic range 1-200 kb

    # True biological importance (independent of length)
    # Most genes have low importance, few have high importance
    true_importance = np.random.exponential(scale=2.0, size=n_genes)

    # Add some truly important genes (disease-causing)
    n_causal = int(n_genes * 0.05)  # 5% are truly causal
    causal_indices = np.random.choice(n_genes, n_causal, replace=False)
    true_importance[causal_indices] += np.random.uniform(10, 20, n_causal)

    # Create DataFrame
    genes_df = pd.DataFrame({
        'gene_id': [f'Gene_{i:04d}' for i in range(n_genes)],
        'length_kb': gene_lengths,
        'true_importance': true_importance
    })

    return genes_df


def simulate_gwas_scores(genes_df: pd.DataFrame, length_bias: float = 0.6) -> pd.DataFrame:
    """
    Simulate GWAS association scores with length bias.

    GWAS tends to rank longer genes higher because:
    - More variants in longer genes
    - More chances for spurious associations
    - Multiple testing burden not fully corrected

    Args:
        genes_df: DataFrame with gene properties
        length_bias: Strength of length bias (0-1)

    Returns:
        DataFrame with added GWAS scores and ranks
    """
    df = genes_df.copy()

    # GWAS score = true importance + length bias + noise
    df['gwas_score'] = (
        df['true_importance'] * (1 - length_bias) +  # True signal
        np.log(df['length_kb']) * length_bias * 5 +   # Length bias
        np.random.normal(0, 2, len(df))                # Noise
    )

    # Convert to ranks (1 = best)
    df['gwas_rank'] = df['gwas_score'].rank(ascending=False, method='min')

    return df


def simulate_burden_test_scores(genes_df: pd.DataFrame) -> pd.DataFrame:
    """
    Simulate burden test scores with less length bias.

    Burden tests aggregate rare variants and are less susceptible to length bias
    but have lower statistical power (more noise).

    Args:
        genes_df: DataFrame with gene properties

    Returns:
        DataFrame with added burden test scores and ranks
    """
    df = genes_df.copy()

    # Burden test score = true importance + small length effect + more noise
    df['burden_score'] = (
        df['true_importance'] * 0.8 +                  # True signal (strong)
        np.log(df['length_kb']) * 0.1 +                # Minimal length bias
        np.random.normal(0, 4, len(df))                # Higher noise (lower power)
    )

    # Convert to ranks (1 = best)
    df['burden_rank'] = df['burden_score'].rank(ascending=False, method='min')

    return df


def aggregate_ranks(genes_df: pd.DataFrame,
                   gwas_weight: float = 0.5,
                   burden_weight: float = 0.5) -> pd.DataFrame:
    """
    Aggregate GWAS and burden test ranks using weighted average.

    Alternative methods include:
    - Borda count (sum of ranks)
    - Robust Rank Aggregation (RRA)
    - Harmonic mean of p-values

    Here we use weighted average for simplicity and interpretability.

    Args:
        genes_df: DataFrame with GWAS and burden ranks
        gwas_weight: Weight for GWAS ranks (0-1)
        burden_weight: Weight for burden test ranks (0-1)

    Returns:
        DataFrame with aggregated ranks
    """
    df = genes_df.copy()

    # Normalize weights
    total_weight = gwas_weight + burden_weight
    gwas_weight /= total_weight
    burden_weight /= total_weight

    # Calculate weighted average of ranks (lower is better)
    df['aggregated_score'] = (
        df['gwas_rank'] * gwas_weight +
        df['burden_rank'] * burden_weight
    )

    # Convert to final ranks
    df['aggregated_rank'] = df['aggregated_score'].rank(method='min')

    # Calculate rank improvement (positive = improved)
    df['rank_improvement'] = df['gwas_rank'] - df['aggregated_rank']

    return df


def identify_example_gene(genes_df: pd.DataFrame,
                         target_gwas_rank: int = 500,
                         target_final_rank: int = 50) -> str:
    """
    Identify a gene that shows dramatic improvement with rank aggregation.

    Args:
        genes_df: DataFrame with all rankings
        target_gwas_rank: Desired initial GWAS rank
        target_final_rank: Desired final aggregated rank

    Returns:
        Gene ID that best matches the criteria
    """
    # Find genes with GWAS rank around 500 and final rank < 100
    candidates = genes_df[
        (genes_df['gwas_rank'] >= 450) &
        (genes_df['gwas_rank'] <= 550) &
        (genes_df['aggregated_rank'] <= 100)
    ].copy()

    if len(candidates) == 0:
        # Relax constraints if no perfect match
        candidates = genes_df[
            (genes_df['gwas_rank'] >= 400) &
            (genes_df['gwas_rank'] <= 600) &
            (genes_df['aggregated_rank'] <= 150)
        ].copy()

    # Pick the one with best improvement
    if len(candidates) > 0:
        best_idx = candidates['rank_improvement'].idxmax()
        return genes_df.loc[best_idx, 'gene_id']

    # Fallback: just pick best improvement overall
    best_idx = genes_df['rank_improvement'].idxmax()
    return genes_df.loc[best_idx, 'gene_id']


def create_visualization(genes_df: pd.DataFrame,
                        example_gene_id: str,
                        output_path: str):
    """
    Create comprehensive visualization showing rank aggregation effects.

    Args:
        genes_df: DataFrame with all rankings
        example_gene_id: Gene to highlight in plots
        output_path: Where to save the figure
    """
    fig, axes = plt.subplots(2, 2, figsize=(14, 12))
    fig.suptitle('GWAS and Burden Test Rank Aggregation',
                 fontsize=16, fontweight='bold')

    # Get example gene data
    example = genes_df[genes_df['gene_id'] == example_gene_id].iloc[0]

    # Plot 1: Length bias in GWAS rankings
    ax1 = axes[0, 0]
    scatter = ax1.scatter(genes_df['length_kb'],
                         genes_df['gwas_rank'],
                         c=genes_df['true_importance'],
                         cmap='viridis',
                         alpha=0.6,
                         s=20)
    ax1.scatter(example['length_kb'],
               example['gwas_rank'],
               color='red',
               s=200,
               marker='*',
               edgecolors='black',
               linewidths=2,
               label=f'{example_gene_id} (rank {int(example["gwas_rank"])})',
               zorder=5)
    ax1.set_xlabel('Gene Length (kb)', fontsize=11)
    ax1.set_ylabel('GWAS Rank (1 = best)', fontsize=11)
    ax1.set_title('A) GWAS Rankings Show Length Bias', fontsize=12, fontweight='bold')
    ax1.invert_yaxis()
    ax1.legend(loc='upper right')
    ax1.grid(True, alpha=0.3)
    plt.colorbar(scatter, ax=ax1, label='True Importance')

    # Plot 2: Burden test rankings (less bias)
    ax2 = axes[0, 1]
    scatter = ax2.scatter(genes_df['length_kb'],
                         genes_df['burden_rank'],
                         c=genes_df['true_importance'],
                         cmap='viridis',
                         alpha=0.6,
                         s=20)
    ax2.scatter(example['length_kb'],
               example['burden_rank'],
               color='red',
               s=200,
               marker='*',
               edgecolors='black',
               linewidths=2,
               label=f'{example_gene_id} (rank {int(example["burden_rank"])})',
               zorder=5)
    ax2.set_xlabel('Gene Length (kb)', fontsize=11)
    ax2.set_ylabel('Burden Test Rank (1 = best)', fontsize=11)
    ax2.set_title('B) Burden Tests Have Less Length Bias', fontsize=12, fontweight='bold')
    ax2.invert_yaxis()
    ax2.legend(loc='upper right')
    ax2.grid(True, alpha=0.3)
    plt.colorbar(scatter, ax=ax2, label='True Importance')

    # Plot 3: Rank comparison before/after
    ax3 = axes[1, 0]
    # Show top 100 genes only for clarity
    top_genes = genes_df.nsmallest(100, 'aggregated_rank')
    for _, gene in top_genes.iterrows():
        if gene['gene_id'] == example_gene_id:
            ax3.plot([1, 2],
                    [gene['gwas_rank'], gene['aggregated_rank']],
                    'r-',
                    linewidth=3,
                    marker='o',
                    markersize=10,
                    label=example_gene_id,
                    zorder=5)
        else:
            ax3.plot([1, 2],
                    [gene['gwas_rank'], gene['aggregated_rank']],
                    'gray',
                    alpha=0.3,
                    linewidth=0.5)

    ax3.set_xlim(0.8, 2.2)
    ax3.set_xticks([1, 2])
    ax3.set_xticklabels(['GWAS Only', 'Aggregated'], fontsize=11)
    ax3.set_ylabel('Rank (1 = best)', fontsize=11)
    ax3.set_title('C) Rank Changes After Aggregation (Top 100 Genes)',
                 fontsize=12, fontweight='bold')
    ax3.invert_yaxis()
    ax3.legend(loc='best')
    ax3.grid(True, alpha=0.3, axis='y')

    # Plot 4: Rank improvement distribution
    ax4 = axes[1, 1]
    improvements = genes_df['rank_improvement'].values
    ax4.hist(improvements, bins=50, color='steelblue', alpha=0.7, edgecolor='black')
    ax4.axvline(example['rank_improvement'],
               color='red',
               linestyle='--',
               linewidth=2,
               label=f'{example_gene_id}: +{int(example["rank_improvement"])} positions')
    ax4.set_xlabel('Rank Improvement (positive = better)', fontsize=11)
    ax4.set_ylabel('Number of Genes', fontsize=11)
    ax4.set_title('D) Distribution of Rank Improvements', fontsize=12, fontweight='bold')
    ax4.legend(loc='upper right')
    ax4.grid(True, alpha=0.3, axis='y')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"\nVisualization saved to: {output_path}")


def main():
    """
    Main execution function demonstrating rank aggregation workflow.
    """
    print("=" * 70)
    print("GWAS and Burden Test Rank Aggregation Demonstration")
    print("=" * 70)

    # Step 1: Generate synthetic data
    print("\n[1] Generating synthetic gene data (n=1000)...")
    genes_df = generate_synthetic_genes(n_genes=1000)
    print(f"    Created {len(genes_df)} genes")
    print(f"    Length range: {genes_df['length_kb'].min():.1f} - {genes_df['length_kb'].max():.1f} kb")
    print(f"    Mean length: {genes_df['length_kb'].mean():.1f} kb")

    # Step 2: Simulate GWAS rankings (with length bias)
    print("\n[2] Simulating GWAS rankings (with length bias)...")
    genes_df = simulate_gwas_scores(genes_df, length_bias=0.6)

    # Calculate correlation between length and GWAS rank
    length_gwas_corr = genes_df['length_kb'].corr(genes_df['gwas_score'])
    print(f"    Correlation between gene length and GWAS score: {length_gwas_corr:.3f}")
    print(f"    (Positive correlation indicates length bias)")

    # Step 3: Simulate burden test rankings (less bias)
    print("\n[3] Simulating burden test rankings (less biased)...")
    genes_df = simulate_burden_test_scores(genes_df)

    length_burden_corr = genes_df['length_kb'].corr(genes_df['burden_score'])
    print(f"    Correlation between gene length and burden score: {length_burden_corr:.3f}")
    print(f"    (Lower correlation = less length bias)")

    # Step 4: Aggregate ranks
    print("\n[4] Aggregating ranks (50% GWAS, 50% burden test)...")
    genes_df = aggregate_ranks(genes_df, gwas_weight=0.5, burden_weight=0.5)

    n_improved = (genes_df['rank_improvement'] > 0).sum()
    n_worsened = (genes_df['rank_improvement'] < 0).sum()
    print(f"    Genes improved: {n_improved} ({n_improved/len(genes_df)*100:.1f}%)")
    print(f"    Genes worsened: {n_worsened} ({n_worsened/len(genes_df)*100:.1f}%)")

    # Step 5: Identify example gene
    print("\n[5] Identifying example gene with dramatic improvement...")
    example_gene_id = identify_example_gene(genes_df, target_gwas_rank=500, target_final_rank=50)
    example = genes_df[genes_df['gene_id'] == example_gene_id].iloc[0]

    print(f"\n    *** EXAMPLE GENE: {example_gene_id} ***")
    print(f"    Gene Length: {example['length_kb']:.1f} kb")
    print(f"    True Importance: {example['true_importance']:.2f}")
    print(f"    GWAS Rank: {int(example['gwas_rank'])} (buried due to short length)")
    print(f"    Burden Test Rank: {int(example['burden_rank'])} (detected by rare variants)")
    print(f"    Aggregated Rank: {int(example['aggregated_rank'])} (RESCUED!)")
    print(f"    Improvement: +{int(example['rank_improvement'])} positions")

    # Step 6: Show top 20 genes from each method
    print("\n[6] Top 20 genes by each ranking method:")
    print("\n    GWAS Top 20:")
    gwas_top = genes_df.nsmallest(20, 'gwas_rank')[['gene_id', 'length_kb', 'gwas_rank']]
    print(gwas_top.to_string(index=False))

    print("\n    Aggregated Top 20:")
    agg_top = genes_df.nsmallest(20, 'aggregated_rank')[['gene_id', 'length_kb', 'aggregated_rank', 'gwas_rank', 'burden_rank']]
    print(agg_top.to_string(index=False))

    # Step 7: Summary statistics
    print("\n[7] Summary Statistics:")
    print(f"\n    Rank correlation between methods:")
    print(f"    - GWAS vs Burden: {genes_df['gwas_rank'].corr(genes_df['burden_rank']):.3f}")
    print(f"    - GWAS vs Aggregated: {genes_df['gwas_rank'].corr(genes_df['aggregated_rank']):.3f}")
    print(f"    - Burden vs Aggregated: {genes_df['burden_rank'].corr(genes_df['aggregated_rank']):.3f}")

    print(f"\n    Genes with >100 position improvement: {(genes_df['rank_improvement'] > 100).sum()}")
    print(f"    Genes with >200 position improvement: {(genes_df['rank_improvement'] > 200).sum()}")
    print(f"    Maximum improvement: {int(genes_df['rank_improvement'].max())} positions")

    # Step 8: Create visualization
    print("\n[8] Creating visualization...")
    output_dir = '/Users/timrichardson/Documents/projects/personal/blog/blog-article-generator/blog-post-gwas-gene-rankings-20251106-143910/assets/code/'
    viz_path = f'{output_dir}01-chart-rank-comparison.png'
    create_visualization(genes_df, example_gene_id, viz_path)

    # Step 9: Key takeaway
    print("\n" + "=" * 70)
    print("KEY TAKEAWAY")
    print("=" * 70)
    print(f"""
The example gene {example_gene_id} demonstrates the power of rank aggregation:

- In GWAS-only analysis, it ranked #{int(example['gwas_rank'])} (buried in noise)
- Its low GWAS rank was due to SHORT LENGTH ({example['length_kb']:.1f} kb)
- Burden tests correctly identified it at rank #{int(example['burden_rank'])}
- After aggregation, it improved to rank #{int(example['aggregated_rank'])}

This represents a {int(example['rank_improvement'])}-position improvement, moving the gene from
"not interesting" to "high priority for follow-up validation."

Length bias in GWAS: r = {length_gwas_corr:.3f}
Length bias in burden: r = {length_burden_corr:.3f}

By combining methods, we leverage the power of GWAS while correcting for its
systematic biases, resulting in more reliable gene prioritization.
    """)

    print("\n" + "=" * 70)
    print("EXECUTION COMPLETE")
    print("=" * 70)


if __name__ == "__main__":
    main()