import gradio as gr
import pandas as pd

image_path = "data/protein_dendrogram.png"

protein_clusters = pd.read_excel('data/protein_clusters.xlsx')
max_length = 500
protein_clusters['proteins'] = protein_clusters['proteins'].apply(lambda x: x[:max_length] + ('...' if len(x) > max_length else ''))
protein_clusters['protein groups'] = protein_clusters['protein groups'].apply(lambda x: x[:max_length] + ('...' if len(x) > max_length else ''))
protein_clusters['protein features'] = protein_clusters['protein features'].apply(lambda x: x[:max_length] + ('...' if len(x) > max_length else ''))


with gr.Blocks() as demo:
    gr.Markdown("# Protein similarity from BERT point of view")
    gr.Markdown("This app displays protein similarity captured in the model [unikei/bert-base-proteins]("
                "https://huggingface.co/unikei/bert-base-proteins).")

    gr.Image(image_path,
             label="Right click to zoom in new tab.",
             container=True
             )

    gr.Markdown("\n")
    gr.Markdown("Click on the [link](https://huggingface.co/spaces/unikei/proteins-from-bert-point-of-view/blob/main/data/protein_clusters.xlsx) to download the spreadsheet.")
    gr.DataFrame(protein_clusters,
                 interactive=False,
                 wrap=True,
                 column_widths=[5, 5, 30, 30, 30])
    #

demo.launch()