from typing import List

import gradio as gr
import numpy as np
import pandas as pd

_ORIGINAL_DF = pd.read_csv("./data/benchmark.csv")
_METRICS = ["MCC", "F1", "ACC"]
_AGGREGATION_METHODS = ["mean", "max", "min", "median"]
_TASKS = {
    "histone_marks": [
        "H4",
        "H3",
        "H3K14ac",
        "H3K4me1",
        "H3K4me3",
        "H3K4me2",
        "H3K36me3",
        "H4ac",
        "H3K79me3",
        "H3K9ac",
    ],
    "regulatory_elements": [
        "promoter_no_tata",
        "enhancers",
        "enhancers_types",
        "promoter_all",
        "promoter_tata",
    ],
    "RNA_production": [
        "splice_sites_donors",
        "splice_sites_all",
        "splice_sites_acceptors",
    ],
}

_BIBTEX = """@article{DallaTorre2023TheNT,
  title={The Nucleotide Transformer: Building and Evaluating Robust Foundation Models for Human Genomics},
  author={Hugo Dalla-Torre and Liam Gonzalez and Javier Mendoza Revilla and Nicolas Lopez Carranza and Adam Henryk Grzywaczewski and Francesco Oteri and Christian Dallago and Evan Trop and Hassan Sirelkhatim and Guillaume Richard and Marcin J. Skwark and Karim Beguir and Marie Lopez and Thomas Pierrot},
  journal={bioRxiv},
  year={2023},
  url={https://api.semanticscholar.org/CorpusID:255943445}
}
"""  # noqa
_LAST_UPDATED = "Sept 15, 2023"

banner_url = "./assets/logo.png"
_BANNER = f'<div style="display: flex; justify-content: space-around;"><img src="{banner_url}" alt="Banner" style="width: 40vw; min-width: 300px; max-width: 600px;"> </div>'  # noqa

_INTRODUCTION_TEXT = """The 🤗 Nucleotide Transformer Leaderboard aims to track, rank and evaluate DNA foundational models on a set of curated downstream tasks introduced in the huggingface dataset [nucleotide_transformer_downstream_tasks](https://huggingface.co/datasets/InstaDeepAI/nucleotide_transformer_downstream_tasks) , with a standardized evaluation protocole presented in the "Methods" tab."""  # noqa

_METHODS_TEXT = """We have compared the fine-tuned performance of Nucleotide Transformer models on the 18 downstream tasks with four different pre-trained models: [DNABERT-1](https://academic.oup.com/bioinformatics/article/37/15/2112/6128680), [DNABERT-2](https://arxiv.org/abs/2306.15006), [HyenaDNA](https://arxiv.org/abs/2306.15794) (1kb and 32kb context length) and the [Enformer](https://www.nature.com/articles/s41592-021-01252-x) (which was trained as a supervised model on several genomics tasks). We ported the architecture and trained weights of each model to our code framework and performed parameter-efficient fine-tuning for every model as described above, using the same cross-validation scheme for a fair comparison. All results can be visulaized in an interactive leader-board 2. Only for HyenaDNA we performed full fine-tuning due to the incompatibility of our parameter-efficient fine-tuning approach with the model architecture."""  # noqa


def retrieve_array_from_text(text):
    return np.fromstring(text.replace("[", "").replace("]", ""), dtype=float, sep=",")


def format_number(x):
    return float(f"{x:.3}")


def get_dataset(
    histone_tasks: List[str],
    regulatory_tasks: List[str],
    rna_tasks: List[str],
    target_metric: str = "MCC",
    aggregation_method: str = "mean",
):
    tasks = histone_tasks + regulatory_tasks + rna_tasks

    aggr_fn = getattr(np, aggregation_method)
    scores = _ORIGINAL_DF[target_metric].apply(retrieve_array_from_text).apply(aggr_fn)
    scores = scores.apply(format_number)
    df = _ORIGINAL_DF.drop(columns=_METRICS)
    df["Score"] = scores
    df = df.pivot(index="Model", columns="Dataset", values="Score")
    df = df[tasks]
    df["All Tasks"] = df.agg("mean", axis="columns").apply(format_number)
    columns = list(df.columns.values)
    columns.sort()
    df = df[columns]
    df.reset_index(inplace=True)
    df = df.rename(columns={"index": "Model"})
    df = df.sort_values(by=["All Tasks"], ascending=False)

    leaderboard_table = gr.components.Dataframe.update(
        value=df,
        # datatype=TYPES,
        max_rows=None,
        interactive=False,
        visible=True,
    )
    return leaderboard_table


def get_bar_plot(
    histone_tasks: List[str],
    regulatory_tasks: List[str],
    rna_tasks: List[str],
    target_metric: str = "MCC",
    aggregation_method: str = "mean",
):
    tasks = histone_tasks + regulatory_tasks + rna_tasks

    aggr_fn = getattr(np, aggregation_method)
    scores = _ORIGINAL_DF[target_metric].apply(retrieve_array_from_text).apply(aggr_fn)
    scores = scores.apply(format_number)
    df = _ORIGINAL_DF.drop(columns=_METRICS)
    df["Score"] = scores / len(tasks)
    df = df.query(f"Dataset == {tasks}")

    bar_plot = gr.BarPlot.update(
        df,
        x="Model",
        y="Score",
        color="Dataset",
        width=500,
        x_label_angle=-45,
        x_title="Model",
        y_title="Score",
        color_legend_title="Downstream Task",
    )
    return bar_plot


with gr.Blocks() as demo:
    with gr.Row():
        gr.Image(banner_url, height=160, scale=1)
        gr.Markdown(_INTRODUCTION_TEXT, elem_classes="markdown-text")
        # gr.Textbox(_INTRODUCTION_TEXT, scale=5)

    with gr.Row():
        metric_choice = gr.Dropdown(
            choices=_METRICS,
            value="MCC",
            label="Metric displayed.",
        )
        aggr_choice = gr.Dropdown(
            choices=_AGGREGATION_METHODS,
            value="mean",
            label="Aggregation used over 10-folds.",
        )

    with gr.Row():
        regulatory_tasks = gr.CheckboxGroup(
            choices=_TASKS["regulatory_elements"],
            value=_TASKS["regulatory_elements"],
            label="Regulatory Elements Downstream Tasks.",
            info="Human data.",
            scale=3,
        )
        rna_tasks = gr.CheckboxGroup(
            choices=_TASKS["RNA_production"],
            value=_TASKS["RNA_production"],
            label="RNA Production Downstream Tasks.",
            info="Human data.",
            scale=3,
        )
        histone_tasks = gr.CheckboxGroup(
            choices=_TASKS["histone_marks"],
            label="Histone Modification Downstream Tasks.",
            info="Yeast data.",
            scale=4,
        )

    with gr.Tabs(elem_classes="tab-buttons") as tabs:
        with gr.TabItem("🏅 Leaderboard", elem_id="od-benchmark-tab-table", id=0):
            dataframe = gr.components.Dataframe(
                elem_id="leaderboard-table",
            )

        with gr.TabItem("📈 Graph", elem_id="od-benchmark-tab-table", id=2):
            bar_plot = gr.BarPlot(
                elem_id="leaderboard-bar-plot",
            )

        with gr.TabItem("ℹ️ Methods", elem_id="od-benchmark-tab-table", id=1):
            gr.Markdown(_METHODS_TEXT, elem_classes="markdown-text")

    gr.Markdown(f"Last updated on **{_LAST_UPDATED}**", elem_classes="markdown-text")

    with gr.Row():
        with gr.Accordion("📙 Citation", open=False):
            gr.Textbox(
                value=_BIBTEX,
                lines=7,
                label="Copy the BibTeX snippet to cite this source",
                elem_id="citation-button",
            ).style(show_copy_button=True)

    histone_tasks.change(
        get_dataset,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=dataframe,
    )
    regulatory_tasks.change(
        get_dataset,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=dataframe,
    )
    rna_tasks.change(
        get_dataset,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=dataframe,
    )
    metric_choice.change(
        get_dataset,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=dataframe,
    )
    aggr_choice.change(
        get_dataset,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=dataframe,
    )
    demo.load(
        fn=get_dataset,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=dataframe,
    )

    histone_tasks.change(
        get_bar_plot,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=bar_plot,
    )
    regulatory_tasks.change(
        get_bar_plot,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=bar_plot,
    )
    rna_tasks.change(
        get_bar_plot,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=bar_plot,
    )
    metric_choice.change(
        get_bar_plot,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=bar_plot,
    )
    aggr_choice.change(
        get_bar_plot,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=bar_plot,
    )
    demo.load(
        fn=get_bar_plot,
        inputs=[histone_tasks, regulatory_tasks, rna_tasks, metric_choice, aggr_choice],
        outputs=bar_plot,
    )

demo.launch()