Protein_Dendrogram.py

import base64
from copy import deepcopy
import streamlit as st
import pandas as pd
import requests
import numpy as np
import io
import plotly.figure_factory as ff
import plotly
# Now lets do pairwise cosine distance
from sklearn.metrics.pairwise import cosine_distances, euclidean_distances
from scipy.cluster.hierarchy import linkage, to_tree
from scipy.spatial.distance import squareform
import time
import textwrap
import traceback

import plotly.graph_objects as go
from plotly.graph_objs import graph_objs


import numpy as np

from utils import write_job_params, write_warnings, enrich_genbank_metadata, metadata_validation, custom_css
from Protein_Dendrogram_Components import draw_protein_heatmap, _Dendrogram
from streamlit.components.v1 import html
import ete3

# Set Page Configuration
st.set_page_config(page_title="IDBac - Dendrogram", page_icon="assets/idbac_logo_square.png", layout="centered", initial_sidebar_state="auto", menu_items=None)
# Add tracking
html('<script async defer data-website-id="4611e28d-c0ff-469d-a2f9-a0b54c0c8ee0" src="https://analytics.gnps2.org/umami.js"></script>', width=0, height=0)
custom_css()

# DEFAULT_TASK_ID = "b06ac61f27e346e3a443234d80e16ced"    # ML TASK
OLD_DEFAULT_TASK_IDS = ["8e0cb0c6a3c04ae1991bbc1dca2882b5",]
DEFAULT_TASK_ID = "4c43a2ca6f3541938e491b3c52442721"
# DEFAULT_TASK_ID = "dd792e0180cb4ef2950077a8ba485790"
# Microbiome Presence: 634aa977f11148da8bada35f1ef2ff06

# Alert banner
# st.warning(
#     "⚠️⚠️ **Maintenance Status Update**: The knowledgebase is currently under maintenance and dependent results may change. "
#     "Workflows executed during this time will be **permanently** affected, even after maintenance is complete. "
#     "Anticipated completion **has been extended to** to 5/12/2025 at 12:00 Noon PST. If you have any questions, please contact us at nkrull@uic.edu."
# )


def assemble_complete_distance_matrix(
    db_distance_dict,
    labels
):
    """Generate a complete distance matrix of (n+m) x (n+m) size from the three distance matrices.
    
    Parameters:
    - db_distance_dict (dict): The dictionary containing the database distance information.
    - all_spectra_df (pandas.DataFrame): The dataframe containing all spectra data relevant to what we're plotting.

    Returns:
    - complete_distance_matrix (numpy.ndarray): The complete distance matrix.
    """
    # all_spectra_df = deepcopy(all_spectra_df)
    # num_inputs = all_spectra_df[all_spectra_df['db_search_result'] == False].shape[0]
    # num_db_search_results = all_spectra_df[all_spectra_df['db_search_result'] == True].shape[0]
    # complete_distance_matrix = np.ones((num_inputs + num_db_search_results, num_inputs + num_db_search_results)) * np.inf    # Multiply by inf to allow for a sanity check
    
    # # 'filename' contains database_id for db search results
    # all_filenames = all_spectra_df['filename']

    num_labels = np.unique(labels).shape[0]
    complete_distance_matrix = np.ones((num_labels, num_labels)) * np.inf

    # Fill in input-input distances
    for i in range(len(labels)):
        fi = labels[i]
        for j in range(i, len(labels)):  # Start from i to cover the diagonal and upper triangle
            fj = labels[j]
            
            if i == j:
                dist = 0.0
            else:
                # Check both directions in the dict efficiently
                dist = db_distance_dict.get(fi, {}).get(fj)
                if dist is None:
                    dist = db_distance_dict.get(fj, {}).get(fi)

            if dist is None:
                raise ValueError(f"Missing distance for {fi} and {fj}")

            complete_distance_matrix[i, j] = dist
            complete_distance_matrix[j, i] = dist

    # Ensure no inf vals left
    if np.isinf(complete_distance_matrix).any():
        raise ValueError("Some distances are missing in the complete distance matrix")

    return complete_distance_matrix

def create_dendrogram(data_np, all_spectra_df, db_distance_dict, 
                      plotted_metadata=[],
                      db_label_column=None,
                      metadata_df=None,
                      db_search_columns=None,
                      cluster_method="average",
                      coloring_threshold=None,
                      cutoff=None,
                      show_annotations=True,
                      ):
    """
    Create a dendrogram using the given data and parameters.

    Parameters:
    - data_np (numpy.ndarray): The input data as a numpy array.
    - all_spectra_df (pandas.DataFrame): The dataframe containing all spectra data.
    - db_distance_dict (dict): The dictionary containing the database distance information.
    - plotted_metadata (list of str, optional): The column name to be shown in the scatter plot. Defaults to "filename".
    - metadata_df (pandas.DataFrame, optional): The dataframe containing metadata information. Defaults to None.
    - db_search_columns (list, optional): The list of columns to be used for displaying database search result metadata. Defaults to None.
    - cluster_method (str, optional): The clustering method to be used for clustering the data. Defaults to "ward".
    - coloring_threshold (float, optional): The threshold for coloring the dendrogram. Defaults to None.
    - cutoff (float, optional): The cutoff line for the dendrogram. Defaults to None.
    - show_annotations (bool, optional): Whether to show annotations on the dendrogram. Defaults to True.

    Returns:
    - dendro (plotly.graph_objs._figure.Figure): The generated dendrogram as a Plotly figure.
    """
    # General Metadata Prep
    if metadata_df is not None:
        original_all_spectra_cols = all_spectra_df.columns
        all_spectra_df = all_spectra_df.merge(metadata_df, how="left", left_on="filename", right_on="Filename", suffixes=("", "_metadata"))
    
    # Prepare text metadata
    if metadata_df is not None and st.session_state["metadata_label"] != "None":
        metadata_column = st.session_state["metadata_label"]
        if metadata_column not in metadata_df.columns:
            st.error("Metadata file does not have the specified column")
        else:
            has_metadata = ~all_spectra_df[metadata_column].isna()
            all_spectra_df.loc[has_metadata, metadata_column] = all_spectra_df.loc[has_metadata, metadata_column].astype(str) + ' - ' # Prepend dash only if we have metadata
            all_spectra_df.loc[:, metadata_column] = all_spectra_df.loc[:, metadata_column].fillna("")
            db_result_mask = (all_spectra_df["db_search_result"] == False)
            all_spectra_df.loc[db_result_mask, "label"] = all_spectra_df.loc[db_result_mask][metadata_column].astype(str) + all_spectra_df.loc[db_result_mask]['label'].astype(str)
    
    # if metadata_df is not None and label_column != 'None':
    if metadata_df is not None and len(plotted_metadata) > 0:
        # Attempt to fall back to lowercase filename if uppercase filename is not present
        if 'Filename' not in metadata_df.columns and 'filename' in metadata_df.columns:
            metadata_df['Filename'] = metadata_df['filename']
        # Raise an error if there is not filename column
        if 'Filename' not in metadata_df.columns and 'filename' not in metadata_df.columns:
            st.error("Metadata file does not have a 'Filename' column")
        
        # If the label column is in the original dataframe, a suffix is added
        plotted_metadata = [metadata_column if metadata_column not in original_all_spectra_cols else metadata_column+"_metadata" for metadata_column in plotted_metadata]
        
    # Add metadata for db search results
    if db_label_column != "No Database Search Results" and sum(all_spectra_df["db_search_result"]) > 0:
        # all_spectra_df.loc[all_spectra_df["db_search_result"] == True, db_metadata_column].fillna("No Metadata", inplace=True)
        # all_spectra_df.loc[all_spectra_df["db_search_result"] == True, "label"] = 'KB Result - ' + all_spectra_df.loc[all_spectra_df["db_search_result"] == True][db_label_column].astype(str)
        if db_search_columns != 'None':
            all_spectra_df.loc[all_spectra_df["db_search_result"] == True, "label"] = 'KB Result - ' \
                + all_spectra_df.loc[all_spectra_df["db_search_result"] == True][db_search_columns].astype(str) \
                + ' - ' + all_spectra_df.loc[all_spectra_df["db_search_result"] == True]['db_strain_name'].astype(str)
        else:
            all_spectra_df.loc[all_spectra_df["db_search_result"] == True, "label"] = 'KB Result - ' \
                + all_spectra_df.loc[all_spectra_df["db_search_result"] == True]['db_strain_name'].astype(str)
        
    # all_spectra_df["label"] = all_spectra_df["label"].astype(str) + " - " + all_spectra_df["filename"].astype(str)
    all_labels_list = all_spectra_df["label"].to_list()
    
    # Reset index to use as unique identifier
    all_spectra_df.reset_index(drop=True, inplace=True)

    # We only have (and only need to input) database_id when there are database search results
    if 'database_id' in all_spectra_df.columns:
        all_spectra_input = all_spectra_df[['filename','db_search_result','database_id']]
    else:
        all_spectra_input = all_spectra_df[['filename','db_search_result']]
        
    if len(all_spectra_input) <= 1:
        st.error("There are not enough spectra to create a dendrogram. Please check number of input spectra and database search results file.")
        return None, None, None

    complete_distance_matrix = assemble_complete_distance_matrix(
        db_distance_dict,
        all_spectra_df.filename.values
    )
    complete_distance_matrix = squareform(complete_distance_matrix, force='tovector')
    _dendro = _Dendrogram(
            complete_distance_matrix, 
            orientation='left',
            labels=all_spectra_df.index.values, # We will use the labels as a unique identifier
            distfun=None,                       # This is a distance matrix
            linkagefun=lambda x: linkage(x, method=cluster_method),
            color_threshold=coloring_threshold
        )
    dendro = graph_objs.Figure(
        data = _dendro.data,
        layout = _dendro.layout
    )

    linkage_matrix = linkage(complete_distance_matrix,
                            method=cluster_method)
    dist_matrix_relative_labels = all_spectra_df.index.values

    dendrogram_width = 800
    dendrogram_height = max(20*len(all_labels_list), 350)
    
    # Add labels for each intersection
    if show_annotations:
        for dd in dendro.data:
            # Get middle two x's and y's
            # (it's actually drawing rectangles and you can think of these as the top right and bottom right corners)
            x = (dd.x[1] + dd.x[2]) / 2
            y = (dd.y[1] + dd.y[2]) / 2
            # Add a text element:
            dendro.add_trace(go.Scatter
                                (x=[x-0.005],
                                y=[y],
                                text=f'{x:.2f}',
                                mode='text',
                                showlegend=False,
                                textposition='middle left',
                                textfont=dict(size=10),
                                hoverinfo='none',
                                xaxis='x',
                                yaxis='y'
                                )
                            )
    # Remove Gridlines from Dendrogram
    dendro.update_xaxes(showgrid=False)
    dendro.update_yaxes(showgrid=False)
    
    # Get y values for the axis labels
    y_values            = dendro['layout']['yaxis']['tickvals']
    y_axis_identifiers  = dendro['layout']['yaxis']['ticktext']
    y_labels            = [all_labels_list[int(identifier)] for identifier in y_axis_identifiers]
    
    # Add a scatter to show metadata
    if metadata_df is not None and plotted_metadata != []:
        num_cols = len(plotted_metadata) + 1
        
        # Calculate ideal column widths based on the number of unique values.
        column_widths = []
        for col in plotted_metadata:
            # Use the unique count if not a numeric type
            if metadata_df[col].dtype == 'object':
                num_unique = len(all_spectra_df[col].unique())
                column_widths.append(max(0.02 * num_unique, 0.15))    # Max width is 0.3 for any given column
            else:
                column_widths.append(0.15)                          # Default width for numeric types
            
        column_widths.append(1 - sum(column_widths))            # Add the remaining width for the dendrogram
        fig = plotly.subplots.make_subplots(rows=1, cols=num_cols,
                                            shared_yaxes=True,
                                            column_widths=column_widths,
                                            horizontal_spacing=0.0)
        fig.update_layout(width=dendrogram_width, height=dendrogram_height, margin=dict(l=0, r=0, b=0, t=0, pad=0))
        
        col_counter = 1
        for col_name in plotted_metadata:
            # Reorder metadata array to be consistent with histogram axis
            consistently_ordered_metadata = all_spectra_df[col_name].loc[y_axis_identifiers]    # Contains list of x-values ordered by y-axis

            # Create the scatter on the new axis
            metadata_scatter = go.Scatter(x=consistently_ordered_metadata, y=y_values, mode='markers')

            # Show all x ticks
            fig.update_xaxes(tickvals=consistently_ordered_metadata,
                             ticktext=consistently_ordered_metadata.astype(str).str.strip(' - '),   # strip in case it's also text metadata
                             row=1,
                             col=col_counter,
                             tickangle=90,
                             ticks="outside",
                             showgrid=True,
                             type='category',
                             categoryorder='category ascending')
            
            fig.add_trace(metadata_scatter, row=1, col=col_counter)
            
            # ylim must be set for each axis, otherwise we get blank space
            fig.update_yaxes(tickvals=y_values, ticktext=y_labels, range=[min(y_values)-10, max(y_values)+25], row=1, col=col_counter, showgrid=True, automargin=True)
            
            # Add title
            wrapped_title = col_name.replace(" ", "<br>") # TODO: Better wrapping
            fig.add_annotation(xref="x domain",yref="y domain",x=0.5, y=1, showarrow=False,
                   text=f"<b>{wrapped_title}</b>", row=1, col=col_counter)
            
            col_counter+=1
            
        # Add a border around the scatter plots
        rectangles_to_add = []
        for x,y in zip(range(1,num_cols), range(1, num_cols)):
            if x == 1:
                x = ""
            if y == 1:
                y = ""
            rectangles_to_add.append(go.layout.Shape(
                    type="rect",
                    xref=f"x{x} domain",
                    yref=f"y domain",
                    x0= 0.0,
                    y0 = 0,
                    x1= 1.,
                    y1= 1,
                    line={'width':1, 'color':"rgb(250, 250, 250)"})
                    )
        
        fig.update_layout(shapes=rectangles_to_add)
        
        # Add the dendrogram to the figure
        for trace in dendro.data:
            fig.add_trace(trace, row=1, col=col_counter)
        
        # Remove gridlines from dendrogram (again)
        fig.update_yaxes(showgrid=False, row=1, col=col_counter)
        
        # Remove legend from dendrogram
        fig.update_layout(showlegend=False)
        
        # Label y axis
        fig.update_yaxes(tickvals=y_values, ticktext=y_labels, row=1, col=1, autorange=False, ticksuffix=" ", ticks="outside")
        # # Set ylim
        #     fig.update_yaxes(range=[min(y_values)-10, max(y_values)+10], row=1, col=col_counter)
        
        translated_labels = [all_labels_list[int(identifier)] for identifier in y_axis_identifiers]
        dist_matrix_relative_labels = [all_labels_list[int(identifier)] for identifier in dist_matrix_relative_labels]

        if cutoff is not None:
            # add to dendrogram col only
            fig.add_vline(x=cutoff, line_width=1, line_color='grey', row=1, col=col_counter)

        return fig, linkage_matrix, dist_matrix_relative_labels
    
    else:

        # Prevent x,y axis labels from being cut off during export
        dendro.update_xaxes(automargin=True, title="")  # I can't believe the title is actually required here lol
        dendro.update_yaxes(automargin=True, title="")

        dendro.update_layout(width=dendrogram_width, height=dendrogram_height) #margin=dict(l=10, r=10, b=10, t=10, pad=0))
        # Set ylim
        dendro.update_yaxes(range=[min(y_values)-5, max(y_values)+5], tickvals=y_values, ticktext=y_labels, autorange=False, ticksuffix=" ", ticks="outside")
        
        translated_labels = [all_labels_list[int(identifier)] for identifier in y_axis_identifiers]
        dist_matrix_relative_labels = [all_labels_list[int(identifier)] for identifier in dist_matrix_relative_labels]

        if cutoff is not None:
            dendro.add_vline(x=cutoff, line_width=1, line_color='grey')
        
        return dendro, linkage_matrix, dist_matrix_relative_labels


@st.cache_data
def collect_database_search_results(task, base_url):
    """
    Collect the database search results from the IDBAC Database. If the database search results are not available, then None is returned.

    Parameters:
    - task (str): The GNPS2 task ID.

    Returns:
    - database_search_results_df (pandas.DataFrame): The dataframe containing the database search results.
    - database_distance_table (pandas.DataFrame): The dataframe containing the database distance tabl (contains original fikle)
    """
    try:
        # Getting the database search results
        # database_search_results_url = f"{base_url}/resultfile?task={task}&file=nf_output/search/enriched_db_results.tsv"
        database_search_results_url = f"{base_url}/resultfile?task={task}&file=nf_output/search/complete_enriched_db_results.tsv"
        print("Knowledgebase Search Results URL", database_search_results_url, flush=True)
        database_search_results_df = pd.read_csv(database_search_results_url, sep="\t")
    except:
        st.error("This is GNPS task is now out of date. Please clone it to use the interactive dashboard.")
        st.stop()
        database_search_results_df = None
        
    try:
        # Get the DB-DB distances
        database_distance_url = f"{base_url}/resultfile?task={task}&file=nf_output/search/db_db_distance.tsv"
        print("Database Distance URL", database_distance_url, flush=True)
        database_database_distance_table = pd.read_csv(database_distance_url, sep="\t")
    except Exception:
        database_database_distance_table = None
        
        if database_search_results_df is not None:
            if 'distance' not in database_search_results_df.columns:
                st.warning("This is GNPS task is now out of date. Please clone it to use the interactive dashboard.")
                st.stop()

    try:
        # Get the query-query distances
        query_query_distance_url = f"{base_url}/resultfile?task={task}&file=nf_output/search/query_query_distances.tsv"
        print("Query-Query Distance URL", query_query_distance_url, flush=True)
        query_query_distance_table = pd.read_csv(query_query_distance_url, sep="\t")
    except Exception:
        st.warning("This is GNPS task is now out of date. Please clone it to use the interactive dashboard.")
        st.stop()
        
    return database_search_results_df, database_database_distance_table, query_query_distance_table




def integrate_database_search_results(
        all_spectra_df:pd.DataFrame,
        database_search_results_df:pd.DataFrame,
        query_query_distances:pd.DataFrame,
        database_database_distances:pd.DataFrame,
        session_state,
    ):
    """
    Integrate the database search results into the original data. Adds unique database search results to the original data and returns a dictionary of database distances.
    Only the database_id column is considered for uniqueness.
    
    Parameters:
    - all_spectra_df (pandas.DataFrame): The dataframe containing all spectra data.
    - database_search_results_df (pandas.DataFrame): The dataframe containing the database search results.
    - session_state (dict): The session state containing the display parameters.
    
    Returns:
    - all_spectra_df (pandas.DataFrame): The dataframe containing all spectra data with database search results added.
    - database_distance_dict (dict): The dictionary containing the database distances.
    """
    db_taxonomy_filter   = session_state["db_taxonomy_filter"]
    distance_threshold = session_state["db_distance_threshold"]
    maximum_db_results   = session_state["max_db_results"]

    filtered_filenames = set(all_spectra_df.filename)

    # Holds ((n+m) * (n+m) / 2 size) matrix of distances
    distances_dict = {}
    ## Query - Query Distances
    for _, row in query_query_distances.iterrows():
        if distances_dict.get(row['query_filename_left']) is None:
            distances_dict[row['query_filename_left']] = {row["query_filename_right"]: row['distance']}
        else:
            distances_dict[row['query_filename_left']][row["query_filename_right"]] = row['distance']


    # If there are no database search results, mark everything as not a database search result and return
    if database_search_results_df is None:
        all_spectra_df["db_search_result"] = False
        return all_spectra_df, distances_dict
    
    # Apply DB Taxonomy Filter
    split_taxonomy = database_search_results_df['db_taxonomy'].str.split(";")
    trimmed_search_results_df = database_search_results_df.loc[[any([x in db_taxonomy_filter for x in y]) for y in split_taxonomy]]
    
    # Reduce visible taxonomy to only Genus, Family, Species
    trimmed_search_results_df['db_taxonomy'] = trimmed_search_results_df['db_taxonomy'].str.split(";").str[-3:].str.join(" - ")
    
    # Apply distance Filter
    trimmed_search_results_df = trimmed_search_results_df[trimmed_search_results_df["distance"] <= distance_threshold]
       
    # Apply Maximum KB Results Filter
    #best_ids = trimmed_search_results_df.sort_values(by="distance", ascending=False).database_id.drop_duplicates(keep="first")
    # Get maximum_db_results database_ids per each query_filename
    if maximum_db_results != -1:
        best_ids = []
        for query_filename in trimmed_search_results_df.query_filename.unique():
            query_results_for_filename = trimmed_search_results_df.loc[trimmed_search_results_df.query_filename == query_filename]
            best_ids_for_filename = query_results_for_filename.sort_values(by="distance", ascending=True).database_id
            best_ids_for_filename = best_ids_for_filename.iloc[:maximum_db_results] # Safe out of bounds
            best_ids.extend(best_ids_for_filename.values)
            
        best_ids = np.unique(best_ids)
        trimmed_search_results_df = trimmed_search_results_df.loc[trimmed_search_results_df["database_id"].isin(best_ids)]
        if database_database_distances is not None:        
            database_database_distances = database_database_distances[(database_database_distances["database_id_left"].isin(best_ids)) &
                                                                            (database_database_distances["database_id_right"].isin(best_ids))]
    
    # We will abuse filename because during display, we display "metadata - filename"
    trimmed_search_results_df["filename"] = trimmed_search_results_df["database_id"].astype(str)
    
    all_spectra_df["db_search_result"] = False
    trimmed_search_results_df["db_search_result"] = True
    
    # Concatenate DB search results
    to_concat = trimmed_search_results_df.drop_duplicates(subset=["database_id"])   # Get unique database hits, assuming database_id is unique
    to_concat = to_concat.drop(columns=['query_filename','distance'])               # Remove distance info 
    all_spectra_df = pd.concat((all_spectra_df, to_concat), axis=0)
    

    ## Query - DB HitsDistances
    for _, row in database_search_results_df.iterrows():
        if distances_dict.get(row['query_filename']) is None:
            distances_dict[row['query_filename']] = {row["database_id"]: row['distance']}
        else:
            distances_dict[row['query_filename']][row["database_id"]] = row['distance']
    
    ## DB Hits - DB Hits Distances
    if database_database_distances is not None:    
        # Add the DB-DB distances to the distance dictionary
        for _, row in database_database_distances.iterrows():    # This is known to be square, no need to flip the indices
            if distances_dict.get(row['database_id_left']) is None:
                distances_dict[row['database_id_left']] = {row['database_id_right']: row['distance']}
            else:
                distances_dict[row['database_id_left']][row['database_id_right']] = row['distance']
    
    return all_spectra_df, distances_dict


# Here we will add an input field for the GNPS2 task ID
url_parameters = st.query_params

if "task" in url_parameters and "task_id" not in st.session_state:
    task_id = url_parameters["task"]
    if task_id in OLD_DEFAULT_TASK_IDS: # Redirect Old Default Task to New Default Task
        st.session_state["task_id"] = DEFAULT_TASK_ID
        st.rerun()
    else:
        st.session_state["task_id"] = task_id

elif "task_id" not in st.session_state:
    st.session_state["task_id"] = DEFAULT_TASK_ID
    
# Add other items to session state if available
if "metadata_label" in url_parameters:
    st.session_state["metadata_label"] = url_parameters["metadata_label"]
if "metadata_scatter" in url_parameters:
    st.session_state["metadata_scatter"] = url_parameters["metadata_scatter"]
if "db_search_result_label" in url_parameters:
    st.session_state["db_search_result_label"] = url_parameters["db_search_result_label"]
if "db_distance_threshold" in url_parameters:
    st.session_state["db_distance_threshold"] = float(url_parameters["db_distance_threshold"])
if "max_db_results" in url_parameters:
    st.session_state["max_db_results"] = int(url_parameters["max_db_results"])
if "db_taxonomy_filter" in url_parameters:
    st.session_state["db_taxonomy_filter"] = url_parameters["db_taxonomy_filter"].split(",")
if "clustering_method" in url_parameters:
    st.session_state["clustering_method"] = url_parameters["clustering_method"]
if "coloring_threshold" in url_parameters:
    st.session_state["coloring_threshold"] = float(url_parameters["coloring_threshold"])
if "hide_isolates" in url_parameters:
    st.session_state["hidden_isolates"] = url_parameters["hide_isolates"].split(",")
if "cutoff" in url_parameters:
    if url_parameters["cutoff"] == 'None':
        st.session_state["cutoff"] = None
    else:
        st.session_state["cutoff"] = float(url_parameters["cutoff"])
if "show_annotations" in url_parameters:
    st.session_state["show_annotations"] = bool(url_parameters["show_annotations"])


# IDBac Logo (+ janky image centering)
col1, col2, col3 = st.columns([1, 3, 1])
with col1:
    st.write(' ')
with col2:
    st.image("assets/idbac_logo.png", )
with col3:
    st.write(' ')

# Welcome message
st.markdown("""
            Welcome to the IDBac Interactive Interface! Here, you can analyze MALDI-TOF MS data to organize, identify, and visualize relationships between bacteria based on protein MS signals;
            associate clustering and biomarker behavior with metadata (e.g., culture conditions); and explore relationships between taxonomy and specialized metabolite production.\
            
            Please see our documentation [here](https://wang-bioinformatics-lab.github.io/GNPS2_Documentation/idbacanalysisplatform/) for more information on how to use this tool.
            """)

if st.session_state.get("task_id", '') == DEFAULT_TASK_ID:
    st.info("""
    ⚠️ You're currently viewing the demo task to explore the features of this dashboard. ⚠️ \
    
    Looking to analyze your own data? [Sign up for an account.](https://wang-bioinformatics-lab.github.io/GNPS2_Documentation/idbac-request-an-account/)
    """)

st.session_state["task_id"] = st.text_input(
    'GNPS2 Task ID',
    st.session_state["task_id"],
    help="Enter your GNPS2 Task ID here to select a different task. If you're coming from GNPS2, this is automatically populated."
)
if st.session_state["task_id"] == '':
    st.error("Please input a valid GNPS2 Task ID")
st.write("Task Loaded:", st.session_state["task_id"])

task_id = st.session_state["task_id"]
base_url = None

# Now we will get all the relevant data from GNPS2 for plotting
if st.session_state["task_id"].startswith("DEV-"):
    base_url = "http://dev.gnps2.org:4000"
    task_id = st.session_state['task_id'][4:]
elif st.session_state["task_id"].startswith("BETA-"):
    base_url = "https://beta.gnps2.org"
    task_id = st.session_state['task_id'][5:]
else:
    base_url = "https://gnps2.org"
    task_id = st.session_state['task_id']

task_status_url = f"{base_url}/status.json?task={task_id}"
labels_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/output_histogram_data_directory/labels_spectra.tsv"
numpy_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/output_histogram_data_directory/numerical_spectra.npy"
bin_counts_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/bin_counts/bin_counts.csv"
replicate_count_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/bin_counts/replicates.csv"
protein_heatmap_binned_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/bin_counts/binned_spectra.csv"
warnings_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/errors.csv"

#### Verify that we can access the data before proceeding with plotting ####
try:
    gnps2_task_status = requests.get(task_status_url, timeout=60)
    if gnps2_task_status.status_code != 200:
        st.error("Unable to access the data for this task. Please check the Task ID and try again.")
        st.stop()
    else:
        gnps2_task_status_json = gnps2_task_status.json()
        gnps2_task_status = gnps2_task_status_json.get("task_status", "unknown")
        gnps2_task_status = str(gnps2_task_status).lower().strip()
        if gnps2_task_status in ["pending", "running"]:
            st.warning("This task is still running. Please check again later.")
            st.stop()
        elif gnps2_task_status in ["failed", "error"]:
            st.error("This task has failed. Please check the workflow and try again.")
            st.stop()
        elif gnps2_task_status == "unknown":
            st.warning("Unable to determine task status. Proceed with caution.")
        else:
            pass
except Exception as e:
    st.warning("Unable to determine task status. Proceed with caution.")

st.session_state['workflow_params'] = write_job_params(st.session_state['task_id'])
if st.checkbox("Show Warnings", value=True, key="show_warnings"):
    write_warnings(warnings_url)

# If workflow parameters specfiy a similarity function, use it. Otherwise, default to cosine
if "distance" in st.session_state['workflow_params'] and st.session_state['workflow_params'] is not None:
    given_distance_measure = st.session_state['workflow_params'].get('distance', 'cosine')
    assert given_distance_measure in {'presence', 'cosine', 'euclidean', 'reverse_cosine', 'reverse_presence'}
    st.session_state['given_distance_measure'] = given_distance_measure
    if given_distance_measure == 'cosine':
        st.session_state['distance_measure'] = cosine_distances
    elif given_distance_measure == 'euclidean':
        st.session_state['distance_measure'] = euclidean_distances
    elif given_distance_measure == 'presence':
        st.session_state['distance_measure'] = cosine_distances # TODO: Ensure array is binary
    else:
        st.session_state['distance_measure'] = None
else:
    st.warning("**Warning:** Unable to find a distance function in the workflow parameters. This may be an old task. Please rerun it. \
               Defaulting to cosine similarity.")
    st.session_state['given_distance_measure'] = 'cosine'
    st.session_state['distance_measure'] = cosine_distances
# By request, no longer displaying labels url
if False:
    st.write(labels_url)

# read numpy from url into a numpy array
try:
    numpy_file = requests.get(numpy_url, 60)
    numpy_file.raise_for_status()
    numpy_array = np.load(io.BytesIO(numpy_file.content))
except:
    st.warning("No Spectra found for this task. Please check the workflow inputs.")
    numpy_array = None
st.session_state['query_spectra_numpy_data'] = numpy_array

############ Protein Heatmap I/O ############
# read bin counts from url
bin_counts_df = None
try:
    bin_counts_csv = requests.get(bin_counts_url, 60)
    bin_counts_csv.raise_for_status()
    bin_counts_df = pd.read_csv(io.StringIO(bin_counts_csv.text), index_col=0)
except:
    if numpy_array is not None:
        st.warning("Unable to retrieve bin counts, this may be an old task.")
    bin_counts_df = None
st.session_state['bin_counts_df'] = bin_counts_df

replicate_count_df = None
try:
    replicate_count_csv = requests.get(replicate_count_url, 60)
    replicate_count_csv.raise_for_status()
    replicate_count_df = pd.read_csv(io.StringIO(replicate_count_csv.text), index_col=1)
except:
    if numpy_array is not None:
        st.warning("Unable to retrieve replicate counts, this may be an old task.")
    replicate_count_df = None
st.session_state['replicate_count_df'] = replicate_count_df

heatmap_binned_spectra = None
try:
    heatmap_binned_spectra_csv = requests.get(protein_heatmap_binned_url, 60)
    heatmap_binned_spectra_csv.raise_for_status()
    heatmap_binned_spectra = pd.read_csv(io.StringIO(heatmap_binned_spectra_csv.text), index_col=0)
except:
    if heatmap_binned_spectra is not None:
        st.warning("Unable to retrieve binned spectra, this may be an old task.")
    heatmap_binned_spectra = None
st.session_state['heatmap_binned_spectra'] = heatmap_binned_spectra

################################################


# read pandas dataframe from url
st.session_state['all_spectra_df'] = None
try:
    all_spectra_df = pd.read_csv(labels_url, sep="\t")
    st.session_state['all_spectra_df'] = all_spectra_df
except:
    if numpy_array is not None:
        st.warning("No Spectra found for this task. Please check the workflow inputs.")
    all_spectra_df = None

# By request, no longer displaying dataframe table
if False:
    st.write(all_spectra_df) # Currently, we're not displaying db search results

# Collect the database search results
db_search_results, db_db_distance_table, query_query_distance_table = collect_database_search_results(task_id, base_url)

if db_db_distance_table is None and db_search_results is not None:
    st.warning("""Knowledgebase-knowledgebase distances are not available for this task, perhaps this is an old task?  
                Please clone and rerun the task on GNPS2 for proper visualization. The distances between these examples
                will be represented as 1.0 to the dendrogram. """)

# Get displayable metadata columns for the database search results
if db_search_results is not None:
    # Remove database search result columns we don't want displayed
    invisible_cols = ['query_filename','distance','query_index','database_index','row_count','database_id','database_scan']
    db_search_columns = [x for x in db_search_results.columns if x not in invisible_cols]
    db_search_results['db_taxonomy'] = db_search_results['db_taxonomy'].fillna("No Taxonomy")
    db_taxonomies = db_search_results['db_taxonomy'].str.split(";").to_list()
    # Flatten
    db_taxonomies = [item for sublist in db_taxonomies for item in sublist]
    db_taxonomies = list(set(db_taxonomies))
    # Sort alphabetically
    db_taxonomies.sort()
    st.session_state['db_search_results'] = db_search_results
    
else:
    db_search_columns = []
    st.session_state['db_search_results'] = None
##### Create Session States #####
# Create a session state for the clustering method
if "clustering_method" not in st.session_state:
    st.session_state["clustering_method"] = "average"
# Create a session state for the coloring threshold
if "coloring_threshold" not in st.session_state:
    st.session_state["coloring_threshold"] = 0.70
# Create a plot for the metadata text labels
if "metadata_label" not in st.session_state:
    st.session_state["metadata_label"] = "None"
# Create a session state for the metadata scatter    
if "metadata_scatter" not in st.session_state:
    st.session_state["metadata_scatter"] = []
# Create a session state for the db search result label
if "db_search_result_label" not in st.session_state and db_search_results is not None:
    st.session_state["db_search_result_label"] = []
elif "db_search_result_label" not in st.session_state and db_search_results is None:
    st.session_state["db_search_result_label"] = "No Database Search Results"
# Create a session state for the db distance threshold
if "db_distance_threshold" not in st.session_state:
    st.session_state["db_distance_threshold"] = 0.50
# Check if db_distance_threshold is greater than the workflow threshold
if float(st.session_state["db_distance_threshold"]) > float(st.session_state['workflow_params']["database_search_threshold"]):
    st.session_state["db_distance_threshold"] = float(st.session_state['workflow_params']["database_search_threshold"])
# Create a session state for the maximum number of database results shown
if "max_db_results" not in st.session_state:
    st.session_state["max_db_results"] = 1
# Create a session state to filter by db taxonomy
if "db_taxonomy_filter" not in st.session_state:
    st.session_state["db_taxonomy_filter"] = None
# Create a session state for alternate metadata
if "upload_metadata" not in st.session_state:
    st.session_state["upload_metadata"] = False
# Create a session state for hidden isolates
if "hidden_isolates" not in st.session_state:
    st.session_state["hidden_isolates"] = []
if "cutoff" not in st.session_state:
    st.session_state["cutoff"] = None
if "show_annotations" not in st.session_state:
    st.session_state["show_annotations"] = False

# Add checkbox for manual metadata upload
if st.checkbox("Upload Metadata", help="If left unchecked, the metadata associated with the task will be used."):
    st.session_state["upload_metadata"] = True
    # Add file uploader
    metadata_file = st.file_uploader("Upload Metadata File", type=["csv", "tsv", "txt", "xlsx"])
    if metadata_file is not None:
        if metadata_file.name.endswith(".txt"):
            metadata_df = pd.read_table(metadata_file, index_col=False)
        elif metadata_file.name.endswith(".csv"):
            metadata_df = pd.read_csv(metadata_file, sep=",", index_col=False)
        elif metadata_file.name.endswith(".tsv"):
            metadata_df = pd.read_csv(metadata_file, sep="\t", index_col=False)
        elif metadata_file.name.endswith(".xlsx"):
            metadata_df =  pd.read_excel(metadata_file, sheet_name=None)
            if "Metadata sheet" in metadata_df:
                metadata_df = metadata_df["Metadata sheet"]
            else:
                metadata_df = metadata_df[list(metadata_df.keys())[0]]
        else:
            st.error("Please upload a .csv, .tsv, or .txt file")
    else:
        metadata_df = None
else:
    # Getting the metadata
    if st.session_state['task_id'].startswith("DEV-"):
        base_url = "http://dev.gnps2.org:4000"
    elif st.session_state['task_id'].startswith("BETA-"):
        base_url = "https://beta.gnps2.org"
    else:
        base_url = "https://gnps2.org"
    try:
        metadata_url = f"{base_url}/resultfile?task={task_id}&file=nf_output/output_histogram_data_directory/metadata.tsv"
        print("metadata_url", metadata_url, flush=True)
        metadata_df = pd.read_csv(metadata_url, sep="\t", index_col=False)
    except:
        metadata_df = None 

if metadata_df is not None:
    # Drop anything with a nan filename
    metadata_df = metadata_df.dropna(subset=["Filename"], axis=0)
    metadata_validation(metadata_df, all_spectra_df)

st.session_state["metadata_df"] = metadata_df

if all_spectra_df is None:
    st.stop()

##### Add Display Parameters #####
st.header("Dendrogram Display Options")

with st.expander("Clustering Settings", expanded=False):
    # Add Clustering Method dropdown
    clustering_options = ["average", "single", "complete", "weighted"]
    if st.session_state['distance_measure'] == "euclidean":
        clustering_options += ['ward', 'median', 'centroid']
    st.session_state["clustering_method"] = st.selectbox("Clustering Method", clustering_options, index=0)
    # Add coloring threshold slider
    st.slider("Color code clusters based on a dendrogram cut-height", 0.0, 1.0, step=0.05, key='coloring_threshold',
            help="All branches that split below the specified height will have the same color. Ex: If this value is set to 0.7, all branches\
                  splitting below 0.7 will have the same color. Above 0.7, all branches will be blue.")

# Only include columns that contain any non-empty values
def _col_has_data(col_series):
    non_na = col_series.dropna()
    if non_na.empty:
        return False
    if non_na.dtype == object:
        return non_na.astype(str).str.strip().ne('').any()
    return non_na.notna().any()

with st.expander("Incorporate Metadata into Analysis", expanded=False):
    st.checkbox("Use GenBank IDs to Enrich Metadata", key="enrich_with_genbank", value=False,
                help="If checked, the metadata will be enriched with GenBank accession numbers.")

    # Add Metadata dropdown for scatter plots
    if metadata_df is None:
        # If there is no metadata, then we will disable the dropdown
        st.session_state["metadata_scatter"] = st.multiselect("Select metadata categories that will be displayed as a scatter plot alongside the dendrogram",
                                                               ["No Metadata Available"], default="No Metadata Available", disabled=True, max_selections=5)
    else:
        columns_available = [c for c in metadata_df.columns if _col_has_data(metadata_df[c])]
        # Remove forbidden columns
        columns_available =[x for x in columns_available if x.lower().strip() not in ['filename', 'scan/coordinate', 'genbank accession', 'ncbi taxid', 'ms collected by', 'isolate collected by', 'sample collected by', 'pi', '16s sequence', 'small molecule file name', 'blank filename']]
        st.session_state["metadata_scatter"]  = st.multiselect("Select metadata categories that will be displayed as a scatter plot alongside the dendrogram", sorted(columns_available), default=[], max_selections=5)

    # Add Metadata dropdown for text
    if metadata_df is None:
        # If there is no metadata, then we will disable the dropdown
        st.session_state["metadata_label"] = st.selectbox("Select a metadata category that will be displayed as text next to the strain ID", ["No Metadata Available"], disabled=True)
    else:
        # Enrich metadata with genebank accession
        if st.session_state["enrich_with_genbank"]:
            metadata_df = enrich_genbank_metadata(metadata_df)
        
        filtered_cols = [c for c in metadata_df.columns if _col_has_data(metadata_df[c])]
        columns_available = ["None"] + filtered_cols
        # Remove filename and scan from the metadata
        columns_available =[x for x in columns_available if x.lower().strip() not in ['filename', 'scan/coordinate', 'small molecule file name', 'blank filename']]
        st.session_state["metadata_label"] = st.selectbox("Select a metadata category that will be displayed as text next to the strain ID", sorted(columns_available))

with st.expander("Knowledgebase Search Results", expanded=False):
    if db_search_results is None:
        # Write a message saying there are no db search results
        text = "No knowledgebase search results found for this task."
        st.write(f":grey[{text}]")

    else:
        # Add DB distance threshold slider
        st.session_state["db_distance_threshold"] = st.slider("Maximum Knowledgebase Distance Threshold", 0.0, 1.0, st.session_state["db_distance_threshold"], 0.05,
                                                                  help="To insert known KB strains into the dendrogram, adjust the distance threshold. Note: 0.00 = identical spectra.")
        workflow_thresh = float(st.session_state['workflow_params'].get("database_search_threshold", 1.0))
        if st.session_state["db_distance_threshold"] > workflow_thresh:
            st.warning(f"The knowledgebase distance threshold is greater than the workflow threshold of {workflow_thresh:.2f}")
        # Create a box for the maximum number of knowledgebase results shown
        st.session_state["max_db_results"] = st.number_input("Maximum Number of Knowledgebase Results Shown", min_value=-1, max_value=None, value=st.session_state["max_db_results"], help="The maximum number of unique database isolates shown, highest distance is prefered. Enter -1 to show all database results.")  
        # Create a 'select all' box for the db taxonomy filter
        # Add DB Search Result dropdown
        db_search_columns_for_selection = ['None'] + [x for x in db_search_columns.copy() if x != 'db_strain_name']
        st.session_state["db_search_result_label"] = st.selectbox("Select a metadata category that will be displayed as text alongside the DB strain ID", db_search_columns_for_selection)
        

        
        col1, col2 = st.columns([0.84, 0.16])
        with col1:
            st.write("Filter results by taxonomic category")
        
        with col2:
            st.checkbox("Show All", value=True, key="select_all_db_taxonomies")
        if st.session_state["select_all_db_taxonomies"] is True:
            st.session_state["db_taxonomy_filter"] = db_taxonomies
            # Add disabled multiselect to make this less jarring
            st.multiselect("Select Displayed Knowledgebase Taxonomies", db_taxonomies, disabled=True, label_visibility="collapsed", placeholder ="Select the taxonomies to display in the dendrogram")
        else:
            # Add multiselect with update button
            st.session_state["db_taxonomy_filter"] = st.multiselect("Customize which KB strains are displayed within the dendrogram by taxonomy", db_taxonomies, label_visibility="collapsed", placeholder ="Select the taxonomies to display in the dendrogram")

with st.expander("General", expanded=False):
    # Add a selectbox that hides isolates
    col1, col2 = st.columns([0.84, 0.16])
    with col1:
        st.write("Select isolates that you would like to remove from the dendrogram")
    with col2:
        st.checkbox("Hide All", value=False, key="hide_all_isolates")
        
    if st.session_state["hide_all_isolates"] is True:
        # st.info("Hiding all isolates is currently disabled. Please select isolates manually.")
        if True:
            st.session_state["hidden_isolates"] = all_spectra_df["filename"].unique()
        # Add disabled multiselect to make this less jarring
        st.multiselect("Select isolates that you would like to remove from the dendrogram", all_spectra_df["filename"].unique(), disabled=True, label_visibility="collapsed", placeholder="Select isolates to hide from the dendrogram")
    else:
        # Add multiselect with update button
        st.session_state["hidden_isolates"] = st.multiselect("Select isolates that you would like to remove from the dendrogram", all_spectra_df["filename"].unique(), label_visibility="collapsed", placeholder="Select isolates to hide from the dendrogram")

    # Add option to add a cutoff line
    st.session_state["cutoff"] = st.number_input("Add a vertical dendrogram cut-height", min_value=0.0, max_value=1.0, value=None, help="Add a vertical line to the dendrogram at the specified distance.")
    # Add option to show annotations
    st.session_state["show_annotations"] = st.checkbox("Display Dendrogram Distances", value=bool(st.session_state["show_annotations"]), help="The values listed represent dendrogram distance. \
                                                                                                        To obtain similarity scores, use the 'Knowledgebase Search Summary' tab within the workflow output.")

st.session_state['query_only_spectra_df'] = all_spectra_df

# Process the db search results (it's done in this order to allow for db_search parameters)
all_spectra_df, db_distance_dict = integrate_database_search_results(
    all_spectra_df,
    db_search_results,
    query_query_distance_table,
    db_db_distance_table,
    st.session_state
    )
st.session_state['db_distance_dict'] = db_distance_dict
st.session_state['spectra_df'] = all_spectra_df

# Remove selected ones from all_spectra_df (believe it or not, we want to remove this after integrating the database search results. This will allow users to hide the queries)
all_spectra_df = all_spectra_df[~all_spectra_df["filename"].isin(st.session_state["hidden_isolates"])]
if len(all_spectra_df) == 0:
    # If there are no spectra to display, then we will stop the script
    st.error("There are no protein spectra to display on the dendrogram. Please select different options and check protein spectra exist for this task.")
    st.stop()


def get_svg_download_link(fig: go.Figure) -> str:
    """
    URL encode the svg image and create a download link for the svg file.

    Parameters:
    - fig (plotly.graph_objs.Figure): The figure to download.

    Returns:
    - link (str): The download link for the svg file.
    """
    # Encode the svg image
    svg_string = plotly.io.to_image(fig, format="svg")
    # Base64 encode the SVG string
    encoded_svg = base64.b64encode(svg_string).decode('utf-8')

    # Generate the data URL for the download link
    data_url = f"data:image/svg+xml;base64,{encoded_svg}"

    # Print out the HTML code for the download link
    html_link = f'<a href="{data_url}" download="plot.svg">Download SVG</a>'
    return html_link


def get_newick_tree_download_link(linkage_matrix, labels):
    """
    Credit: https://github.com/scipy/scipy/issues/8274
    """
    tree = to_tree(linkage_matrix, False)

    def build_newick(node, newick, parentdist, leaf_names):
        if node.is_leaf(): # This is for SciPy not for ete or skbio so `is_leaf` utility function does not apply
            return f"{leaf_names[node.id]}:{(parentdist - node.dist)/2}{newick}"
        else:
            if len(newick) > 0:
                newick = f"):{(parentdist - node.dist)/2}{newick}"
            else:
                newick = ");"
            newick = build_newick(node.get_left(), newick, node.dist, leaf_names)
            newick = build_newick(node.get_right(), f",{newick}", node.dist, leaf_names)
            newick = f"({newick}"
            return newick
    tree = build_newick(tree, "", tree.dist, labels)

    # Generate the data URL for the download link
    data_url = f"data:text/plain;base64,{base64.b64encode(tree.encode()).decode()}"
    # Print out the HTML code for the download link
    html_link = f'<a href="{data_url}" download="dendrogram.nwk">Download Newick Tree</a>'
    return html_link

def get_taxa_annotation_file(labels, metadata, all_spectra_df):
    """Download a .txt file for usage in ITOL.

    Parameters:
    - labels (list): List of labels for the dendrogram.
    - metadata (pandas.DataFrame): DataFrame containing metadata for the labels.

    Returns:
    - str: An html string containing the download link for the annotation file.
    """
    header = textwrap.dedent("""
    AUTO_TAXONOMY

    SEPARATOR COMMA

    TAXONOMY_TYPE,NCBI

    REPLACE_LEAF_LABELS 0

    #1: taxonomy_name
    #2: original_label (taxonomy_name)
    #3: original_label|taxonomy_name
    #4: taxonomy_name (original_label)
    #5: taxonomy_name|original_label

    LABEL_FORMAT 4

    DATA

    """)
    try:
        if set(['filename', 'NCBI taxid', 'db_search_result']).issubset(all_spectra_df.columns) is False:
            return "NCBI TaxID information not available for this task."
        metadata = all_spectra_df[['filename', 'NCBI taxid', 'db_search_result']].copy(deep=True)

        metadata.dropna(subset=['NCBI taxid'], inplace=True)
        metadata.loc[metadata['db_search_result'], 'filename'] = 'KB Result - ' + metadata.loc[metadata['db_search_result'], 'filename'].astype(str)

        color_assignments = metadata.set_index('filename')['NCBI taxid'].to_dict()
       
        print("color_assignments", color_assignments, flush=True)

        # Data format is "node_id,taxa_id" on each line
        data_lines = []
        for label in labels:
            # print("label", label, flush=True)
            color = color_assignments.get(label)
            if color is not None:
                data_lines.append(f"{label},{color}")
                # print(f"Added line for label {label}: {int(color)}", flush=True)

        output_str = header + "\n" + "\n".join(data_lines)
        # Encode the output string to base64
        encoded_output = base64.b64encode(output_str.encode()).decode()
        # Generate the data URL for the download link
        data_url = f"data:text/plain;base64,{encoded_output}"

        return f'<a href="{data_url}" download="annotation_file.txt">Download Annotation File</a>'

    except Exception as e:
        print("Error creating annotation file:", e, flush=True)
        print(traceback.format_exc(), flush=True)
        return "Unable to Create Annotation File."
    
def get_taxa_coloring_file(labels, metadata, all_spectra_df):
    header = textwrap.dedent("""
    TREE_COLORS
    SEPERATEOR TAB
    DATA
    """)
    # Data format NODE_ID TYPE COLOR LABEL_OR_STYLE SIZE_FACTOR
    # We will only use NODE_ID,label_backgrounds,COLOR

    try:
        cols = ['filename', 'db_genus', 'db_search_result']
        cols += ['genus'] if 'genus' in all_spectra_df.columns else []
        metadata = all_spectra_df[cols].copy(deep=True)
        # Move db_genus to genus if it exists
        if 'genus' in metadata.columns:
            metadata['genus'] = metadata['db_genus'].fillna(metadata['genus'])
        else:
            metadata['genus'] = metadata['db_genus']

        metadata.dropna(subset=['genus'], inplace=True)
        metadata.loc[metadata['db_search_result'], 'filename'] = 'KB Result - ' + metadata.loc[metadata['db_search_result'], 'filename'].astype(str)

        # Convert taxa to colors using pastel categorical %20
        taxa = metadata['genus'].unique()
        # Create a color map for the taxa
        color_map = {}
        for i, taxon in enumerate(taxa):
            # Use plotly to generate a pastel color
            color = plotly.colors.qualitative.Pastel[i % len(plotly.colors.qualitative.Pastel)]
            color_map[taxon] = color

        # Create a dictionary for taxa assignments
        metadata['color'] = metadata['genus'].map(color_map)
        taxa_assignments = metadata.set_index('filename')['color'].to_dict()

        print("taxa_assignments", taxa_assignments, flush=True)

        data_lines = []

        for label in labels:
            color = taxa_assignments.get(label)
            if color is not None:
                data_lines.append(f"{label}\tlabel_background\t{str(color).replace('rgb', 'rgba').replace(')', ', 1.0)')}")  
                print(f"Added line for label {label}: {color}", flush=True)

        output_str = header + "\n" + "\n".join(data_lines)
        # Encode the output string to base64
        encoded_output = base64.b64encode(output_str.encode()).decode()
        # Generate the data URL for the download link
        data_url = f"data:text/plain;base64,{encoded_output}"
        return f'<a href="{data_url}" download="coloring_file.txt">Download Coloring File</a>'
    except Exception as e:
        print("Error creating coloring file:", e, flush=True)
        print(traceback.format_exc(), flush=True)
        return "Unable to Create Coloring File."


# Creating the dendrogram
if st.session_state['distance_measure'] is not None:
    dendro, linkage_matrix, labels = create_dendrogram(     numpy_array,
                                                            all_spectra_df,
                                                            db_distance_dict,
                                                            plotted_metadata=st.session_state["metadata_scatter"],
                                                            db_label_column=st.session_state["db_search_result_label"],
                                                            metadata_df=metadata_df,
                                                            db_search_columns=st.session_state["db_search_result_label"] ,
                                                            cluster_method=st.session_state["clustering_method"],
                                                            coloring_threshold=st.session_state["coloring_threshold"],
                                                            cutoff=st.session_state["cutoff"],
                                                            show_annotations=st.session_state["show_annotations"]
                                                        )
    if dendro is not None:
        # DEBUG
        # Save the dendrogram to a 400 DPI PNG file
        # plotly.io.write_image(dendro, "dendrogram.png", scale=2, width=1200, height=800, format="png")

        config = {
            'toImageButtonOptions': {
            'format': 'png', # one of png, svg, jpeg, webp
            'filename': 'IDBac_Protein_Dendrogram',
            'scale':5 # Multiply title/legend/axis/canvas sizes by this factor
            }
        }

        st.plotly_chart(dendro, use_container_width=True, config=config)
        # Display number of strains
        st.write(f"Number of Strains: {len(labels)}")
        # Add option to download as svg
        st.markdown(get_svg_download_link(dendro), unsafe_allow_html=True, help="Currently, this feature only officially supports the dendrogram _without_ metadata.")
        # Add option to download as ete tree
        st.markdown(get_newick_tree_download_link(linkage_matrix, labels), unsafe_allow_html=True)
        st.markdown(get_taxa_annotation_file(labels, metadata_df, all_spectra_df), unsafe_allow_html=True, help="Download an annotation preset for usage in ITOL. \
                    This file will replace knowledgebase hits with their NCBI taxonomies. \
                    Include a column 'NCBI taxid' in your metadata to use this feature for queries")
        st.markdown(get_taxa_coloring_file(labels, metadata_df, all_spectra_df), unsafe_allow_html=True, help="Download a coloring file for usage in ITOL. \
                    For KB results, genus will be used. Include a column 'genus' in your metadata to use this feature.\
                    To show colors in ITOL after uploading the annotation file, toggle Advanced>Node options>Leaf node symbols. It can be set back to 'Hide' once done.")

# Create a shareable link to this page
st.write("Shareable Link: ")
link = f"https://analysis.idbac.org/?task={st.session_state['task_id']}&\
        metadata_label={st.session_state['metadata_label']}&\
        metadata_scatter={st.session_state['metadata_scatter']}&\
        db_search_result_label={st.session_state['db_search_result_label']}&\
        db_distance_threshold={st.session_state['db_distance_threshold']}&\
        max_db_results={st.session_state['max_db_results']}&\
        clustering_method={st.session_state['clustering_method']}&\
        coloring_threshold={st.session_state['coloring_threshold']}&\
        hide_isolates={','.join(st.session_state['hidden_isolates'])}&\
        cutoff={st.session_state['cutoff']}&\
        show_annotations={st.session_state['show_annotations']}"
st.code(link)

# Add documentation
st.header("Additional Information")
# Add a bulleted list of more information
st.markdown("""
            #### Visualization
            * A flat line between strains along the y-axis of the dendrogram at '0' represents 
            identical spectra between those strains.
            * Cluster Color: All branches below the specified height will have the same color.
            """)