#!/usr/bin/env python3
"""
Tool: volcano_ma_plot.py
Category: Transcriptomics & Data Visualization
Description: Reads differential gene expression results and generates publication-grade
             Volcano plots and MA plots with threshold highlights and gene labeling.
Usage:
    python volcano_ma_plot.py -i deseq_results.tsv -o volcano.png --fc 1.5 --fdr 0.05
"""

import argparse
import sys
import numpy as np
import pandas as pd

def generate_plot_data():
    """Generates synthetic differential expression dataset for demonstration."""
    np.random.seed(42)
    n = 200
    genes = [f"Gene_{i:04d}" for i in range(1, n + 1)]
    base_mean = 10 ** np.random.uniform(1, 4, n)
    log2fc = np.random.normal(0, 1.2, n)
    
    # Introduce true signals
    log2fc[0:5] += 2.8
    log2fc[5:10] -= 2.8
    
    pvals = np.random.beta(0.5, 2.0, n)
    pvals[0:10] *= 1e-5
    
    return pd.DataFrame({
        'Gene': genes,
        'baseMean': np.round(base_mean, 1),
        'log2FoldChange': np.round(log2fc, 3),
        'padj': pvals
    })

def main():
    parser = argparse.ArgumentParser(description="Publication-Grade Volcano and MA Plot Generator.")
    parser.add_argument("-i", "--input", help="Path to input TSV differential expression table")
    parser.add_argument("-o", "--output", default="volcano_plot.png", help="Path to output plot image")
    parser.add_argument("--fc", type=float, default=1.5, help="Fold-change cutoff in log2 units (default: 1.5)")
    parser.add_argument("--fdr", type=float, default=0.05, help="Adjusted p-value / FDR cutoff (default: 0.05)")
    args = parser.parse_args()

    df = generate_plot_data() if not args.input else pd.read_csv(args.input, sep='\t')
    
    df['negLog10P'] = -np.log10(np.clip(df['padj'], 1e-15, 1.0))
    df['Category'] = 'Non-significant'
    df.loc[(df['log2FoldChange'] >= args.fc) & (df['padj'] < args.fdr), 'Category'] = 'Upregulated'
    df.loc[(df['log2FoldChange'] <= -args.fc) & (df['padj'] < args.fdr), 'Category'] = 'Downregulated'

    print("=== Differential Expression Summary ===")
    print(df['Category'].value_counts())
    
    top_up = df[df['Category'] == 'Upregulated'].sort_values('padj').head(5)
    print("\nTop Significant Upregulated Biomarkers:")
    for _, r in top_up.iterrows():
        print(f"  - {r['Gene']:10} | Log2FC: {r['log2FoldChange']:6.2f} | FDR: {r['padj']:.2e}")

    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        
        plt.figure(figsize=(7, 6))
        colors = {'Non-significant': '#94a3b8', 'Upregulated': '#ef4444', 'Downregulated': '#3b82f6'}
        for cat, color in colors.items():
            sub = df[df['Category'] == cat]
            plt.scatter(sub['log2FoldChange'], sub['negLog10P'], c=color, label=cat, alpha=0.7, s=25)
            
        plt.axvline(args.fc, color='#cbd5e1', linestyle='--', alpha=0.8)
        plt.axvline(-args.fc, color='#cbd5e1', linestyle='--', alpha=0.8)
        plt.axhline(-np.log10(args.fdr), color='#cbd5e1', linestyle='--', alpha=0.8)
        
        plt.title("RNA-Seq Volcano Plot (Omics Portal Toolkit)", fontsize=13, fontweight='bold')
        plt.xlabel("Log2 Fold Change", fontsize=11)
        plt.ylabel("-Log10 Adjusted P-Value", fontsize=11)
        plt.legend()
        plt.tight_layout()
        plt.savefig(args.output, dpi=200)
        print(f"\n[Saved Volcano Plot Image to: {args.output}]")
    except ImportError:
        print("\n[Install matplotlib to render plot graphics: pip install matplotlib]")

if __name__ == "__main__":
    main()