src/ui/plots/umap_embedding.py

import streamlit as st
import scanpy as sc
import matplotlib.pyplot as plt
import numpy as np
import textwrap
from .utils import display_plot_with_download, display_interactive_spatial_plot, display_plotly_with_download

def render_umap_embedding(metabolic_adata):
    """Render UMAP embedding with Red theme."""
    st.markdown("<h2 style='color: #d32f2f;'><i class='fas fa-palette'></i> UMAP Analysis</h2>", unsafe_allow_html=True)
    
    umap_viz_type = st.session_state.get("u_v_t", "Domain")
    
    if umap_viz_type == "Domain":
        c1, c2 = st.columns([1.5, 1.5])
    else:
        c1, c2, c3 = st.columns([1.2, 1.8, 1.2])
    
    with c1:
        umap_viz_type = st.selectbox("Color By:", options=["Domain", "Reaction", "Pathway"], key="u_v_t")
    
    with (c2 if umap_viz_type == "Domain" else c3):
        plot_mode = st.radio("Plot Mode:", ["Static", "Interactive"], horizontal=True, key="u_mode")
    
    selected_items = []
    
    if umap_viz_type != "Domain":
        with c2:
            if umap_viz_type == "Reaction":
                if 'rxn_full_names' in metabolic_adata.var.columns:
                    # Map full name to ID for user selection
                    unique_names = {}
                    for idx, row in metabolic_adata.var.iterrows():
                        f_name = str(row['rxn_full_names'])
                        if f_name not in unique_names:
                            unique_names[f_name] = idx
                    rx_options = sorted(list(unique_names.keys()))
                    
                    if plot_mode == "Interactive":
                        sel_name = st.selectbox("Select Reaction:", options=rx_options, key="u_rx_single")
                        selected_items = [unique_names[sel_name]] if sel_name else []
                    else:
                        sel_names = st.multiselect("Select Reactions:", options=rx_options, default=rx_options[:1], key="u_rx_multi")
                        selected_items = [unique_names[n] for n in sel_names if n in unique_names]
                else:
                    rx_options = metabolic_adata.var_names.tolist()
                    if plot_mode == "Interactive":
                        sel = st.selectbox("Select Reaction:", options=rx_options, key="u_rx_single")
                        selected_items = [sel] if sel else []
                    else:
                        selected_items = st.multiselect("Select Reactions:", options=rx_options, default=rx_options[:1], key="u_rx_multi")
            
            elif umap_viz_type == "Pathway":
                if 'subsystems' in metabolic_adata.var.columns:
                    import pandas as pd
                    path_options = sorted([p for p in metabolic_adata.var['subsystems'].unique() if pd.notna(p)])
                    if plot_mode == "Interactive":
                        sel = st.selectbox("Select Pathway:", options=path_options, key="u_path_single")
                        selected_items = [sel] if sel else []
                    else:
                        selected_items = st.multiselect("Select Pathways:", options=path_options, default=path_options[:1], key="u_path_multi")
                else:
                    st.warning("No pathway data.")

    if 'X_umap' not in metabolic_adata.obsm:
        with st.spinner("Calculating UMAP..."):
            sc.pp.pca(metabolic_adata, n_comps=50)
            sc.pp.neighbors(metabolic_adata, n_neighbors=15, n_pcs=50)
            sc.tl.umap(metabolic_adata)
    
    try:
        if plot_mode == "Interactive" and (umap_viz_type == "Domain" or selected_items):
            import plotly.express as px
            import pandas as pd
            
            umap_coords = metabolic_adata.obsm['X_umap']
            target = selected_items[0] if selected_items else "Domain"
            display_title = target
            if umap_viz_type == "Reaction" and 'rxn_full_names' in metabolic_adata.var.columns:
                if target in metabolic_adata.var_names:
                    display_title = metabolic_adata.var.loc[target, 'rxn_full_names']
            
            if umap_viz_type == "Domain":
                vals = metabolic_adata.obs["domain"].astype(str).values
                color_scale = None # Use default qualitative for domain
                color_label = "Domain"
            elif target in metabolic_adata.var_names:
                idx = metabolic_adata.var_names.get_loc(target)
                raw = metabolic_adata.X[:, idx]
                vals = raw.toarray().flatten() if hasattr(raw, "toarray") else np.asarray(raw).flatten()
                color_scale = "Jet"
                color_label = "Flux"
            else:
                # Pathway
                rx_list = metabolic_adata.var[metabolic_adata.var['subsystems'] == target].index.tolist()
                X_sub = metabolic_adata[:, rx_list].X
                vals = np.array(X_sub.mean(axis=1)).flatten() if not hasattr(X_sub, "toarray") else np.array(X_sub.toarray().mean(axis=1)).flatten()
                color_scale = "Jet"
                color_label = "Flux"

            df_umap = pd.DataFrame({
                "UMAP1": umap_coords[:, 0],
                "UMAP2": umap_coords[:, 1],
                "color": vals,
                "Domain": metabolic_adata.obs["domain"].values if "domain" in metabolic_adata.obs.columns else "N/A",
                "Spot": metabolic_adata.obs_names
            })
            
            fig = px.scatter(df_umap, x="UMAP1", y="UMAP2", color="color", 
                             hover_data=["Domain", "Spot"],
                             color_continuous_scale=color_scale if color_scale else None,
                             title=f"UMAP Analysis: {display_title}")
            
            fig.update_layout(
                template="simple_white",
                coloraxis_colorbar=dict(title=color_label) if color_scale else None,
                legend_title_text="Domain" if umap_viz_type == "Domain" else None,
                yaxis=dict(scaleanchor="x", scaleratio=1),
                width=700, height=700,
                xaxis=dict(showgrid=False, zeroline=False),
                yaxis_showgrid=False, yaxis_zeroline=False
            )
            # Dynamic help text
            help_msg = f"Uniform Manifold Approximation and Projection (UMAP) is used for dimensionality reduction. "
            if umap_viz_type == "Reaction":
                help_msg += f"This plot shows the flux distribution of **{display_title}** in the reduced feature space."
            elif umap_viz_type == "Pathway":
                 help_msg += f"Across the UMAP manifold, we visualize the average flux for the **{target}** pathway."
            else:
                help_msg += "Spots are colored by metabolic domain to visualize global functional clustering."

            display_plotly_with_download(
                fig, 
                f"umap_{umap_viz_type}",
                help_text=help_msg
            )


        elif umap_viz_type == "Domain":
            # Static Domain
            fig, ax = plt.subplots(figsize=(8, 8))
            sc.pl.umap(metabolic_adata, color=['domain'], show=False, ax=ax, size=100)
            display_plot_with_download(
                fig, 
                "umap_domain",
                help_text="This static UMAP shows the distribution of metabolic domains in lower-dimensional space. Spots colored by domain help visualize how well-separated the clustered metabolic regions are."
            )

            plt.close(fig)

        elif selected_items:
            per_page = 8
            total = len(selected_items)
            pages = (total + per_page - 1) // per_page
            if "umap_page" not in st.session_state: st.session_state.umap_page = 1
            if st.session_state.umap_page > pages: st.session_state.umap_page = 1
            
            curr = selected_items[(st.session_state.umap_page-1)*per_page : st.session_state.umap_page*per_page]
            n_cols = min(2, len(curr))
            n_rows = (len(curr) + n_cols - 1) // n_cols
            fig, axes = plt.subplots(n_rows, n_cols, figsize=(5*n_cols, 4.5*n_rows))
            
            if len(curr) == 1: axes = np.array([[axes]])
            elif n_rows == 1: axes = axes.reshape(1, -1)
            elif n_cols == 1: axes = axes.reshape(-1, 1)

            for i, target in enumerate(curr):
                r, c = i // n_cols, i % n_cols
                if target in metabolic_adata.var_names:
                    sc.pl.umap(metabolic_adata, color=[target], cmap='jet', show=False, ax=axes[r,c], size=80)
                    if 'rxn_full_names' in metabolic_adata.var.columns:
                        full_name = str(metabolic_adata.var.loc[target, 'rxn_full_names'])
                        wrapper = textwrap.TextWrapper(width=40)
                        axes[r,c].set_title(wrapper.fill(text=full_name), fontsize=10)
                else:
                    # Pathway aggregate
                    rx_list = metabolic_adata.var[metabolic_adata.var['subsystems'] == target].index.tolist()
                    metabolic_adata.obs['tmp_u'] = np.array(metabolic_adata[:, rx_list].X.mean(axis=1)).flatten()
                    sc.pl.umap(metabolic_adata, color=['tmp_u'], cmap='jet', show=False, ax=axes[r,c], size=80)
                    wrapper = textwrap.TextWrapper(width=40)
                    axes[r,c].set_title(wrapper.fill(text=str(target)), fontsize=10)
                    if 'tmp_u' in metabolic_adata.obs: del metabolic_adata.obs['tmp_u']
                axes[r,c].axis('off')
            
            for j in range(len(curr), n_rows*n_cols): axes[j//n_cols, j%n_cols].axis('off')
            plt.tight_layout()
            # Dynamic help text for static panels
            static_names = []
            for t in curr:
                if t in metabolic_adata.var_names and 'rxn_full_names' in metabolic_adata.var.columns:
                    static_names.append(metabolic_adata.var.loc[t, 'rxn_full_names'])
                else:
                    static_names.append(str(t))
            static_names_str = ", ".join(static_names)

            display_plot_with_download(
                fig, 
                f"umap_p{st.session_state.umap_page}",
                help_text=f"These static UMAP panels show the flux distribution for: **{static_names_str}**. It helps identify metabolic hotspots for these specific processes within the reduced manifold."
            )

            plt.close(fig)

            if pages > 1:
                cx1, cx2, cx3 = st.columns([1,2,1])
                if cx1.button("Prev UMAP Page", key="u_prev"): st.session_state.umap_page -= 1; st.rerun()
                cx2.markdown(f"<center>Page {st.session_state.umap_page} / {pages}</center>", unsafe_allow_html=True)
                if cx3.button("Next UMAP Page", key="u_next"): st.session_state.umap_page += 1; st.rerun()

    except Exception as e:
        st.error(f"Error during UMAP visualization: {e}")