modules/differential.py

"""
Differential Analysis Module
=============================

Differential flux analysis between metabolic domains/groups.
"""

import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import logging
from scipy import stats
from typing import Optional, List
from streamlit_option_menu import option_menu
import spmetatme.plotting as pl
import io
from datetime import datetime

logger = logging.getLogger(__name__)


def display_plot_with_download(fig, plot_name: str = "plot"):
    """
    Display a matplotlib figure with a PDF download button on top right.
    
    Parameters
    ----------
    fig : matplotlib.figure.Figure
        The matplotlib figure to display and download
    plot_name : str
        Name for the downloaded file (without extension)
    """
    # Create layout with download button on top right
    col_space, col_download = st.columns([5.5, 0.5], gap="small")
    
    with col_download:
        # Generate PDF file
        pdf_buffer = io.BytesIO()
        fig.savefig(pdf_buffer, format='pdf', dpi=300, bbox_inches='tight')
        file_data = pdf_buffer.getvalue()
        
        st.download_button(
            label="📥",
            data=file_data,
            file_name=f"{plot_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf",
            mime="application/pdf",
            key=f"download_{plot_name}_{id(fig)}",
            help="Download as PDF",
            use_container_width=False
        )
    
    # Display the plot
    st.pyplot(fig)


def render():
    """Render differential analysis UI with sidebar menu."""
    # Check if we have flux data
    if st.session_state.metabolic_adata is None:
        st.warning("⚠️ No flux data available")
        st.markdown("""
        Please:
        1. **For spatial data**: Complete preprocessing and run flux analysis
        2. **For pre-computed fluxes**: Upload your flux data in the Upload Data tab
        """)
        return
    
    metabolic_adata = st.session_state.metabolic_adata
    
    # Initialize selected differential page
    if 'selected_diff_page' not in st.session_state:
        st.session_state.selected_diff_page = "Differential Reactions"
    
    # Define differential analysis options
    diff_options = [
        "Differential Reactions",
        "Pathway Selection",
        "Differential Pathways",
        "Pathways by Variance"
    ]
    
    diff_icons = [
        "table",
        "fire",
        "diagram-3",
        "graph-up"
    ]
    
    # Get the current index
    try:
        current_index = diff_options.index(st.session_state.selected_diff_page)
    except ValueError:
        current_index = 0
        st.session_state.selected_diff_page = "Differential Reactions"
    
    # Sidebar menu for differential analysis selection
    with st.sidebar:
        selected_diff = option_menu(
            menu_title="Differential Analysis",
            options=diff_options,
            icons=diff_icons,
            default_index=current_index,
            orientation="vertical",
            styles={
                "container": {"padding": "0!important", "background-color": "#ffffff"},
                "icon": {"color": "#1a73e8", "font-size": "18px"},
                "nav-link": {
                    "font-size": "12px",
                    "text-align": "left",
                    "margin": "0px",
                    "padding": "12px 15px",
                    "--hover-color": "#e3f2fd",
                    "color": "#333333"
                },
                "nav-link-selected": {
                    "background-color": "#1a73e8",
                    "color": "#ffffff",
                    "font-weight": "600"
                }
            },
            key="diff_option_menu"
        )
        
        # Only rerun if selection changed
        if selected_diff != st.session_state.selected_diff_page:
            st.session_state.selected_diff_page = selected_diff
            st.rerun()
        
        st.markdown("---")
        
        # Back to home button in sidebar
        if st.button("🏠 Back to Home", use_container_width=True, key="back_to_home_diff_sidebar"):
            st.session_state.adata = None
            st.session_state.metabolic_adata = None
            st.session_state.data_type = None
            st.session_state.preprocessing_done = False
            st.session_state.flux_analysis_done = False
            st.session_state.selected_diff_page = None
            st.rerun()
        
        st.markdown("---")
        
        # Info section in sidebar
        st.markdown("""
        <div style='background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%); padding: 1rem; border-radius: 8px; font-size: 0.85rem; line-height: 1.6; border-left: 3px solid #1a73e8;'>
            <strong style='color: #1a73e8;'>📊 Differential Analysis</strong><br>
            Identify metabolically distinct regions and enriched reactions across domains.
        </div>
        """, unsafe_allow_html=True)
    
    # Main content area
    st.markdown("## 📉 Differential Metabolic Flux Analysis")
    
    st.markdown("""
    Identify metabolic reactions and pathways with significant differences between
    spatial domains and metabolic phenotypes.
    """)
    
    st.markdown("---")
    
    # Render selected differential analysis page
    if st.session_state.selected_diff_page == "Differential Reactions":
        render_differential_reactions(metabolic_adata)
    
    elif st.session_state.selected_diff_page == "Pathway Selection":
        render_pathway_selection(metabolic_adata)
    
    elif st.session_state.selected_diff_page == "Differential Pathways":
        render_differential_pathways(metabolic_adata)
    
    elif st.session_state.selected_diff_page == "Pathways by Variance":
        render_pathways_by_variance(metabolic_adata)


def render_differential_reactions(metabolic_adata):
    """Render differential reactions analysis with tabs for different heatmap types."""
    st.markdown("### Differential Metabolic Reactions Analysis")
    
    st.markdown("""
    Analyze differentially enriched metabolic reactions across spatial domains
    using different visualization approaches.
    """)
    
    # Create tabs for different analysis types
    tab1, tab2, tab3 = st.tabs([
        "Pathway-Specific Reactions",
        "All Differential Reactions",
        "Pathways by Variance"
    ])
    
    # TAB 1: Pathway-Specific Reactions (plot_differential_reactions_by_pathway_heatmap)
    with tab1:
        st.markdown("#### Pathway-Specific Differential Analysis")
        
        if 'subsystems' not in metabolic_adata.var.columns:
            st.error("Pathway information (subsystems) not found in data")
        else:
            available_pathways = sorted(metabolic_adata.var['subsystems'].unique().tolist())
            
            # Controls
            col1, col2, col3 = st.columns(3)
            
            with col1:
                selected_pathway = st.selectbox(
                    "Select pathway:",
                    options=available_pathways,
                    key="tab1_pathway_dropdown"
                )
            
            with col2:
                top_n_pathway = st.slider(
                    "Top N reactions",
                    min_value=5,
                    max_value=50,
                    value=15,
                    step=1,
                    key="tab1_pathway_top_n"
                )
            
            with col3:
                row_cluster = st.checkbox("Cluster rows", value=True, key="tab1_row_cluster")
            
            try:
                with st.spinner(f"Analyzing {selected_pathway}..."):
                    # Generate heatmap
                    df_pathway = pl.plot_differential_reactions_by_pathway_heatmap(
                        metabolic_adata, 
                        selected_pathway, 
                        row_cluster=row_cluster, 
                        return_marker_df=True, 
                        save_path=None,
                        top_n=top_n_pathway
                    )
                    
                    fig = plt.gcf()
                    
                    # Two-column layout: Heatmap and Table
                    col_plot, col_table = st.columns([1, 1], gap="large")
                    
                    with col_plot:
                        display_plot_with_download(fig, f"{selected_pathway.replace(' ', '_')}_Heatmap")
                    
                    with col_table:
                        st.write("")  
                        st.markdown("##### Reactions Data")
                        if df_pathway is not None:
                            st.dataframe(df_pathway, use_container_width=True)
                            
                            # Download button
                            csv = df_pathway.to_csv(index=False)
                            st.download_button(
                                label="📥 Download Table (CSV)",
                                data=csv,
                                file_name=f"pathway_{selected_pathway.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                                mime="text/csv",
                                key="tab1_download_table"
                            )
                        else:
                            st.info("No data available")
                    
            except Exception as e:
                st.error(f"Error: {str(e)}")
                logger.error(f"Tab1 error: {str(e)}", exc_info=True)
    
    # TAB 2: All Differential Reactions (plot_differential_reactions_heatmap)
    with tab2:
        st.markdown("#### All Differential Reactions Heatmap")
        
        # Controls
        col1, col2 = st.columns(2)
        
        with col1:
            top_n_reactions = st.slider(
                "Top N reactions to show",
                min_value=5,
                max_value=100,
                value=20,
                step=5,
                key="tab2_top_n_reactions"
            )
        
        with col2:
            st.write("")  # Spacer
        
        try:
            with st.spinner("Analyzing all differential reactions..."):
                # Generate heatmap
                df_reactions = pl.plot_differential_reactions_heatmap(
                    metabolic_adata, 
                    save_path=None,
                    top_n=top_n_reactions,
                    return_marker_df=True
                )
                
                fig = plt.gcf()
                                            
                # Two-column layout: Heatmap and Table
                col_plot, col_table = st.columns([1, 1], gap="large")
                
                with col_plot:
                    display_plot_with_download(fig, "Differential_Reactions_Heatmap")
                
                with col_table:
                    st.write("")  
                    st.markdown("##### Reactions Data")
                    if df_reactions is not None:
                        st.dataframe(df_reactions, use_container_width=True)
                        
                        # Download button
                        csv = df_reactions.to_csv(index=False)
                        st.download_button(
                            label="📥 Download Table (CSV)",
                            data=csv,
                            file_name=f"differential_reactions_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                            mime="text/csv",
                            key="tab2_download_table"
                        )
                    else:
                        st.info("No data available")
                                        
        except Exception as e:
            st.error(f"Error: {str(e)}")
            logger.error(f"Tab2 error: {str(e)}", exc_info=True)
    
    # TAB 3: Pathways by Variance (plot_pathways_flux_heatmap)
    with tab3:
        st.markdown("#### Pathways by Variance")
        
        # Controls
        col1, col2, col3 = st.columns(3)
        
        with col1:
            top_n = st.slider(
                "Top N pathways",
                min_value=5,
                max_value=30,
                value=20,
                step=1,
                key="tab3_top_n"
            )
        
        with col2:
            sort_by = st.selectbox(
                "Sort by",
                options=["variance", "mean"],
                key="tab3_sort_by"
            )
        
        with col3:
            st.write("")  # Spacer
        
        
        try:
            with st.spinner(f"Analyzing top {top_n} pathways by {sort_by}..."):
                # Generate heatmap
                df_pathways_var = pl.plot_pathways_flux_heatmap(
                    metabolic_adata, 
                    group_key="domain", 
                    pathway_key="subsystems", 
                    top_n=top_n, 
                    sort_by=sort_by
                )
                
                fig = plt.gcf()
                                
                # Two-column layout: Heatmap and Table
                col_plot, col_table = st.columns([1, 1], gap="large")
                
                with col_plot:
                    display_plot_with_download(fig, f"Pathways_Variance_Top{top_n}")
                
                with col_table:
                    st.markdown("##### Pathways Data")
                    if df_pathways_var is not None:
                        st.dataframe(df_pathways_var, use_container_width=True)
                        
                        # Download button
                        csv = df_pathways_var.to_csv(index=False)
                        st.download_button(
                            label="📥 Download Table (CSV)",
                            data=csv,
                            file_name=f"pathways_variance_top{top_n}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                            mime="text/csv",
                            key="tab3_download_table"
                        )
                    else:
                        st.info("No data available")
                
        except Exception as e:
            st.error(f"Error: {str(e)}")
            logger.error(f"Tab3 error: {str(e)}", exc_info=True)


def render_pathway_selection(metabolic_adata):
    """Render interactive pathway selection with dropdown for differential analysis."""
    st.markdown("### Pathway-Specific Differential Analysis")
    
    st.markdown("""
    Select any metabolic pathway to investigate differential enrichment of reactions
    within that pathway across spatial metabolic domains.
    """)
    
    # Get all available pathways
    if 'subsystems' not in metabolic_adata.var.columns:
        st.error("Pathway information (subsystems) not found in data")
        return
    
    available_pathways = sorted(metabolic_adata.var['subsystems'].unique().tolist())
    
    # Pathway selection
    col1, col2 = st.columns(2)
    
    with col1:
        selected_pathway = st.selectbox(
            "Select pathway to analyze:",
            options=available_pathways,
            key="pathway_dropdown"
        )
    
    with col2:
        top_n_pathway = st.slider(
            "Top N reactions to display",
            min_value=5,
            max_value=50,
            value=15,
            step=1,
            key="pathway_top_n"
        )
    
    # Analysis options
    col1, col2, col3 = st.columns(3)
    
    with col1:
        row_cluster = st.checkbox("Cluster rows", value=True, key="pathway_row_cluster")
    
    with col2:
        show_table = st.checkbox("Show data table", value=True, key="pathway_show_table")
    
    with col3:
        show_stats = st.checkbox("Show statistics", value=True, key="pathway_show_stats")
    
    if st.button(f"📊 Analyze {selected_pathway}", key="pathway_analyze_btn"):
        try:
            with st.spinner(f"Analyzing {selected_pathway}..."):
                
                # Generate the heatmap
                df_pathway = pl.plot_differential_reactions_by_pathway_heatmap(
                    metabolic_adata, 
                    selected_pathway, 
                    row_cluster=row_cluster, 
                    return_marker_df=True, 
                    save_path=None,
                    top_n=top_n_pathway
                )
                
                # Get the current figure
                fig = plt.gcf()
                
                st.success(f"✓ {selected_pathway} analysis completed!")
                
                # Display with download option
                display_plot_with_download(fig, f"Pathway_{selected_pathway.replace(' ', '_')}_Heatmap")
                
                st.markdown("---")
                
                # Display statistics if requested
                if show_stats:
                    col1, col2, col3 = st.columns(3)
                    
                    with col1:
                        reactions_in_pathway = len(df_pathway) if df_pathway is not None else 0
                        st.metric("Reactions in Pathway", reactions_in_pathway)
                    
                    with col2:
                        if 'domain' in metabolic_adata.obs.columns:
                            n_domains = metabolic_adata.obs['domain'].nunique()
                            st.metric("Number of Domains", n_domains)
                    
                    with col3:
                        st.metric("Spatial Spots", metabolic_adata.n_obs)
                    
                    st.markdown("---")
                
                # Show data table if requested
                if show_table and df_pathway is not None:
                    st.markdown(f"#### {selected_pathway} - Reactions Data")
                    st.dataframe(df_pathway, use_container_width=True)
                    
                    # Download button for table
                    csv = df_pathway.to_csv(index=False)
                    st.download_button(
                        label="📥 Download Table (CSV)",
                        data=csv,
                        file_name=f"pathway_{selected_pathway.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        mime="text/csv",
                        key="download_pathway_table"
                    )
                
                st.info(f"💡 Tip: This heatmap shows the {top_n_pathway} most differential reactions in the {selected_pathway} pathway")
                
        except Exception as e:
            st.error(f"Error analyzing {selected_pathway}: {str(e)}")
            logger.error(f"Pathway selection error for {selected_pathway}: {str(e)}", exc_info=True)



def render_differential_pathways(metabolic_adata):
    """Render differential pathways heatmap (top N pathways)."""
    st.markdown("### Differential Pathways Heatmap")
    
    st.markdown("""
    This visualization shows metabolic pathways with the largest differences
    in mean flux between spatial domains. Each pathway is aggregated from its constituent reactions.
    """)
    
    # Options
    col1, col2 = st.columns(2)
    
    with col1:
        top_n_pathways = st.slider(
            "Number of top pathways to show",
            min_value=5,
            max_value=20,
            value=15,
            step=1,
            key="diff_pathway_top_n"
        )
    
    with col2:
        show_table = st.checkbox("Show data table", value=True, key="diff_pathway_show_table")
    
    if st.button("📊 Generate Differential Pathways Heatmap", key="diff_pathway_btn"):
        try:
            with st.spinner("Generating differential pathways heatmap..."):
                
                # Generate the heatmap
                fig = plt.figure(figsize=(14, 10))
                df_pathways = pl.plot_differential_pathways_heatmap(
                    metabolic_adata, 
                    save_path=None,
                    top_n=top_n_pathways
                )
                
                # Get the current figure
                fig = plt.gcf()
                
                st.success("✓ Differential pathways heatmap generated successfully!")
                
                # Display with download option
                display_plot_with_download(fig, "Differential_Pathways_Heatmap")
                
                st.markdown("---")
                
                # Display statistics
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    st.metric("Top Pathways Shown", top_n_pathways)
                
                with col2:
                    if 'domain' in metabolic_adata.obs.columns:
                        n_domains = metabolic_adata.obs['domain'].nunique()
                        st.metric("Number of Domains", n_domains)
                
                with col3:
                    if 'subsystems' in metabolic_adata.var.columns:
                        n_pathways = metabolic_adata.var['subsystems'].nunique()
                        st.metric("Total Pathways", n_pathways)
                
                st.info("💡 Tip: Pathways ranked by the sum of absolute flux differences across domains")
                
                # Show data table if requested
                if show_table and df_pathways is not None:
                    st.markdown("---")
                    st.markdown("#### Differential Pathways Data")
                    st.dataframe(df_pathways, use_container_width=True)
                    
                    # Download button for table
                    csv = df_pathways.to_csv(index=False)
                    st.download_button(
                        label="📥 Download Table (CSV)",
                        data=csv,
                        file_name=f"differential_pathways_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        mime="text/csv",
                        key="download_diff_pathways_table"
                    )
                
        except Exception as e:
            st.error(f"Error generating differential pathways heatmap: {str(e)}")
            logger.error(f"Differential pathways error: {str(e)}", exc_info=True)


def render_pathways_by_variance(metabolic_adata):
    """Render pathways ranked by variance (top N)."""
    st.markdown("### Pathways by Variance")
    
    st.markdown("""
    This visualization shows metabolic pathways with the highest variance
    in flux values across the tissue. High variance indicates heterogeneous metabolic activity
    and potential metabolic specialization across domains.
    """)
    
    # Options
    col1, col2, col3 = st.columns(3)
    
    with col1:
        top_n = st.slider(
            "Number of pathways to show",
            min_value=5,
            max_value=30,
            value=20,
            step=1,
            key="pathway_variance_n"
        )
    
    with col2:
        sort_by = st.selectbox(
            "Sort by",
            options=["variance", "mean"],
            key="pathway_sort_by"
        )
    
    with col3:
        show_table = st.checkbox("Show data table", value=True, key="pathway_var_show_table")
    
    if st.button("📊 Generate Pathways by Variance Heatmap", key="pathway_var_btn"):
        try:
            with st.spinner(f"Generating top {top_n} pathways by {sort_by} heatmap..."):
                
                # Generate the heatmap
                fig = plt.figure(figsize=(14, 10))
                df_pathways_var = pl.plot_pathways_flux_heatmap(
                    metabolic_adata, 
                    group_key="domain", 
                    pathway_key="subsystems", 
                    top_n=top_n, 
                    sort_by=sort_by
                )
                
                # Get the current figure
                fig = plt.gcf()
                
                st.success(f"✓ Pathways by {sort_by} heatmap generated successfully!")
                
                # Display with download option
                display_plot_with_download(fig, f"Pathways_Variance_Top{top_n}")
                
                st.markdown("---")
                
                # Display statistics
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    st.metric("Top Pathways Shown", top_n)
                
                with col2:
                    st.metric("Sort Metric", sort_by.capitalize())
                
                with col3:
                    if 'domain' in metabolic_adata.obs.columns:
                        n_domains = metabolic_adata.obs['domain'].nunique()
                        st.metric("Number of Domains", n_domains)
                
                st.info(f"💡 Tip: Shows {top_n} most variable pathways across spatial domains, highlighting metabolic hotspots")
                
                # Show data table if requested
                if show_table and df_pathways_var is not None:
                    st.markdown("---")
                    st.markdown(f"#### Top {top_n} Pathways by {sort_by.title()}")
                    st.dataframe(df_pathways_var, use_container_width=True)
                    
                    # Download button for table
                    csv = df_pathways_var.to_csv(index=False)
                    st.download_button(
                        label="📥 Download Table (CSV)",
                        data=csv,
                        file_name=f"pathways_variance_top{top_n}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                        mime="text/csv",
                        key="download_pathways_var_table"
                    )
                
        except Exception as e:
            st.error(f"Error generating pathways by variance heatmap: {str(e)}")
            logger.error(f"Pathways by variance error: {str(e)}", exc_info=True)