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
from .utils import display_plot_with_download, display_formatted_table

logger = logging.getLogger(__name__)



def render_differential_reactions(metabolic_adata):
    """Replicated original render_differential_reactions."""
    st.markdown("<h2 style='color: #d32f2f;'><i class='fas fa-chart-bar'></i> Differential Metabolic Reactions Analysis</h2>", unsafe_allow_html=True)
    
    tab1, tab2, tab3, tab4 = st.tabs([
        "Pathway-Specific Reactions",
        "All Differential Reactions",
        "Pathways by Variance",
        "Differential Pathways"
    ])
    
    with tab1:
        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())
            col1, col2, col3 = st.columns(3)
            with col1:
                selected_pathway = st.selectbox("Select pathway:", options=available_pathways, key="tab1_path_sel", help="Select a metabolic pathway for differential reaction analysis.")
            with col2:
                top_n = st.slider("Top N reactions", 5, 20, 10, key="tab1_top_n", help="Filter the number of top significant reactions to display.")
            with col3:
                row_cluster = st.checkbox("Cluster rows", value=True, key="tab1_cluster")
            
            try:
                with st.spinner(f"Analyzing {selected_pathway}..."):
                    df = pl.plot_differential_reactions_by_pathway_heatmap(
                        metabolic_adata, selected_pathway, row_cluster=row_cluster, return_marker_df=True, top_n=top_n
                    )
                    fig = plt.gcf()
                    col_p, col_t = st.columns([1.5, 1], gap="small")
                    with col_p:
                        display_plot_with_download(fig, f"{selected_pathway.replace(' ', '_')}_Heatmap")
                    with col_t:
                        if df is not None:
                            display_formatted_table(df, "Differential Reactions")
                            csv = df.to_csv(index=False)
                            st.download_button("Download CSV", data=csv, file_name=f"{selected_pathway}.csv", mime="text/csv", icon=":material/download:")
            except Exception as e:
                st.error(f"Error: {e}")

    with tab2:
        col1, col2 = st.columns(2)
        with col1:
            top_n = st.slider("Top N reactions:", 5, 20, 10, key="tab2_top_n", help="Number of differentially active reactions to display across all pathways.")
        with col2:
            row_cluster = st.checkbox("Cluster rows", value=False, key="tab2_cluster")
        
        try:
            with st.spinner("Analyzing..."):
                df = pl.plot_differential_reactions_heatmap(metabolic_adata, top_n=top_n, row_cluster=row_cluster, return_marker_df=True)
                fig = plt.gcf()
                col_p, col_t = st.columns([1.5, 1], gap="small")
                with col_p: display_plot_with_download(fig, "Diff_Reactions_Heatmap")
                with col_t:
                    if df is not None:
                        display_formatted_table(df, "Differential Reactions")
        except Exception as e:
            st.error(f"Error: {e}")

    with tab3:
        col1, col2 = st.columns(2)
        with col1: top_n = st.slider("Top N pathways", 5, 20, 10, key="tab3_top_n", help="Filter top pathways based on the selected metric.")
        with col2: sort_by = st.selectbox("Sort by", ["variance", "mean"], key="tab3_sort", help="Metric to rank pathways.")
        try:
            with st.spinner("Analyzing..."):
                df = pl.plot_pathways_flux_heatmap(metabolic_adata, group_key="domain", pathway_key="subsystems", top_n=top_n, sort_by=sort_by)
                fig = plt.gcf()
                col_p, col_t = st.columns([1.5, 1], gap="small")
                with col_p: display_plot_with_download(fig, "Pathways_Var")
                with col_t:
                    if df is not None:
                        display_formatted_table(df, "Pathways Data")
        except Exception as e:
            st.error(f"Error: {e}")
    with tab4:
        col1, col2 = st.columns(2)
        with col1: 
            top_n = st.slider("Top N pathways", 5, 20, 10, key="tab4_top_n", help="Filter top pathways.")
        with col2: 
            row_cluster = st.checkbox("Cluster rows", value=True, key="tab4_cluster")
        try:
            with st.spinner("Analyzing..."):
                df = pl.plot_differential_pathways_heatmap(metabolic_adata, row_cluster=row_cluster, top_n=top_n, return_marker_df= True)
                fig = plt.gcf() 
                col_p, col_t = st.columns([1.5, 1], gap="small")
                with col_p: display_plot_with_download(fig, "Pathways_Var")
                with col_t:
                    if df is not None:
                        display_formatted_table(df, "Differential Pathways")
                        csv = df.to_csv(index=False)
                        st.download_button("Download CSV", data=csv, file_name=f"Pathways_Diff.csv", mime="text/csv", icon=":material/download:")
        except Exception as e:
            st.error(f"Error: {e}")