import datasets
import streamlit as st
from st_aggrid import AgGrid, GridOptionsBuilder, JsCode

st.set_page_config(layout='wide')

# parse out gene_ids from URL query args to it's possible to link to this page
query_params = st.query_params
if "gene_id" in query_params.keys():
    gene_id = query_params["gene_id"]
else:
    gene_id = "CNAG_04365"


st.markdown("""
# CryptoCEN Top50 co-expressed partners

**CryptoCEN** is a co-expression network for *Cryptococcus neoformans* built on 1,524 RNA-seq runs across 34 studies.
A pair of genes are said to be co-expressed when their expression is correlated across different conditions and
is often a marker for genes to be involved in similar processes. 

To Cite:
    O'Meara MJ, Rapala JR, Nichols CB, Alexandre C, Billmyre RB, Steenwyk JL, A Alspaugh JA, O'Meara TR
    CryptoCEN: A Co-Expression Network for Cryptococcus neoformans reveals novel proteins involved in DNA damage repair.
    PLoS Genet 20(2): e1011158. (2024) https://doi.org/10.1371/journal.pgen.1011158

* Code available at https://github.com/maomlab/CalCEN/tree/master/vignettes/CryptoCEN
* Full network and dataset: https://huggingface.co/datasets/maomlab/CryptoCEN

## Look up top-coexpressed partners:
Put in the ``CNAG_#####`` gene_id for a gene and expand the table to get the top 50 co-expressed genes.
``coexp_score`` ranges between ``[0-1]``, where ``1`` is the best and greater than ``0.85`` can be considered significant.
""")

top_coexp_hits = datasets.load_dataset(
    path = "maomlab/CryptoCEN",
    data_files = {"top_coexp_hits": "top_coexp_hits.tsv"})
top_coexp_hits = top_coexp_hits["top_coexp_hits"].to_pandas()

col1, col2, col3 = st.columns(spec = [0.3, 0.2, 0.5])
with col1:
    gene_id = st.text_input(
        label = "Gene ID",
        value = f"{gene_id}",
        max_chars = 10,
        help = "CNAG Gene ID e.g. CNAG_04365")

top_coexp_hits = top_coexp_hits[
    top_coexp_hits.gene_id_1 == gene_id]
top_coexp_hits = top_coexp_hits[[
    'gene_id_1', 'gene_symbol_1', 'description_1',
    'gene_id_2', 'gene_symbol_2', 'description_2',
    'coexp_score', 'blastp_EValue']]
top_coexp_hits["scatter_link"] = top_coexp_hits.apply(
    lambda row: f"https://huggingface.co/spaces/maomlab/CryptoCEN-ExpressionScatter?gene_id_1={row.gene_id_1}&gene_id_2={row.gene_id_2}",
    axis = 1)

top_coexp_hits.reset_index()

with col2:
    st.text('') # help alignment with input box
    st.download_button(
        label = "Download data as TSV",
        data = top_coexp_hits.to_csv(sep ='\t').encode('utf-8'),
        file_name = f"top_coexp_hits_{gene_id}.tsv",
        mime = "text/csv")

# third column is padding only


######### Table #########

fungidb_link_renderer = JsCode("""
        class UrlCellRenderer {
          init(params) {
            this.eGui = document.createElement('a');
            this.eGui.innerText = params.value;
            this.eGui.setAttribute('href', "https://fungidb.org/fungidb/app/record/gene/" + params.value);
            this.eGui.setAttribute('style', "text-decoration:none");
            this.eGui.setAttribute('target', "_blank");
          }
          getGui() {
            return this.eGui;
          }
        }
""")

coexp_scatter_link_renderer = JsCode("""
        class UrlCellRenderer {
          init(params) {
            this.eGui = document.createElement('a');
            this.eGui.innerText = 'plot';
            this.eGui.setAttribute('href', params.value);
            this.eGui.setAttribute('style', "text-decoration:none");
            this.eGui.setAttribute('target', "_blank");
          }
          getGui() {
            return this.eGui;
          }
        }
""")



grid_option_builder = GridOptionsBuilder()
grid_option_builder.configure_auto_height()
grid_option_builder.configure_default_column(
    filterable=False,
    groupable=False,
    editable=False,
    wrapText=True,
    autoHeight=True)
grid_option_builder.configure_column("gene_id_1", header_name="GeneID 1", pinned="left", cellRenderer=fungidb_link_renderer, width=550)
grid_option_builder.configure_column("gene_symbol_1", header_name="Gene 1", pinned="left", width=500)
grid_option_builder.configure_column("description_1", header_name="Description 1", width=1600)
grid_option_builder.configure_column("gene_id_2", header_name="GeneID 2", pinned="left", cellRenderer=fungidb_link_renderer, width=550)
grid_option_builder.configure_column("gene_symbol_2", header_name="Gene 2", pinned="left", width=500)
grid_option_builder.configure_column("description_2", header_name="Description 2", width=1600)
grid_option_builder.configure_column(
    "coexp_score",
    header_name="Coexp Score",
    type=["numericColumn", "customNumericFormat"],
    precision=3,
    width=600)
grid_option_builder.configure_column(
    "blast_EValue",
    header_name="Blast E-value",
    type=["numericColumn", "customNumericFormat"],
    precision=3,
    width=600)
grid_option_builder.configure_column(
    "scatter_link",
    header_name="Expression Scatterplot",
    cellRenderer=coexp_scatter_link_renderer,
    width=550)

grid_option_builder.configure_selection(selection_mode=False, use_checkbox=False)

AgGrid(
    data = top_coexp_hits,
    gridOptions = grid_option_builder.build(),
    fit_columns_on_grid_load=True,
    theme="streamlit",
    allow_unsafe_jscode=True)