import streamlit as st
import pandas as pd
import project_config
import base64


@st.cache_data(show_spinner = 'Loading knowledge graph nodes...')
def load_kg():
    # with st.spinner('Loading knowledge graph...'):
    kg_nodes = pd.read_csv(project_config.DATA_DIR / 'kg_nodes.csv', dtype = {'node_index': int}, low_memory = False)
    return kg_nodes


@st.cache_data(show_spinner = 'Loading knowledge graph edges...')
def load_kg_edges():
    # with st.spinner('Loading knowledge graph...'):
    kg_edges = pd.read_csv(project_config.DATA_DIR / 'kg_edges.csv', dtype = {'edge_index': int, 'x_index': int, 'y_index': int}, low_memory = False)
    return kg_edges


def capitalize_after_slash(s):
    # Split the string by slashes first
    parts = s.split('/')
    # Capitalize each part separately
    capitalized_parts = [part.title() for part in parts]
    # Rejoin the parts with slashes
    capitalized_string = '/'.join(capitalized_parts).replace('_', ' ')
    return capitalized_string


# Define dictionary mapping node types to database URLs
map_dbs = {
    'gene/protein': lambda x: f"https://ncbi.nlm.nih.gov/gene/?term={x}",
    'drug': lambda x: f"https://go.drugbank.com/drugs/{x}",
    'effect/phenotype': lambda x: f"https://hpo.jax.org/app/browse/term/HP:{x.zfill(7)}", # pad with 0s to 7 digits
    'disease': lambda x: x, # MONDO
    # pad with 0s to 7 digits
    'biological_process': lambda x: f"https://amigo.geneontology.org/amigo/term/GO:{x.zfill(7)}", 
    'molecular_function': lambda x: f"https://amigo.geneontology.org/amigo/term/GO:{x.zfill(7)}",
    'cellular_component': lambda x: f"https://amigo.geneontology.org/amigo/term/GO:{x.zfill(7)}",
    'exposure': lambda x: f"https://ctdbase.org/detail.go?type=chem&acc={x}",
    'pathway': lambda x: f"https://reactome.org/content/detail/{x}",
    'anatomy': lambda x: x,
}

# Define dictionary mapping node types to database names
map_db_names = {
    'gene/protein': 'NCBI',
    'drug': 'DrugBank',
    'effect/phenotype': 'HPO',
    'disease': 'MONDO',
    'biological_process': 'GO: BP',
    'molecular_function': 'GO: MF',
    'cellular_component': 'GO: CC',
    'exposure': 'CTD',
    'pathway': 'Reactome',
    'anatomy': 'UBERON',
}