import gradio as gr
import uuid
import asyncio

from substra_launcher import launch_substra_space
from huggingface_hub import HfApi

hf_api = HfApi()


theme = gr.themes.Default(primary_hue="blue").set(
    background_fill_primary="#F9F2EA",
    block_background_fill="#FFFFFF",
)


async def launch_experiment(hospital_a, hospital_b):
    experiment_id = str(uuid.uuid4())

    asyncio.create_task(launch_substra_space(
        hf_api=hf_api,
        repo_id=experiment_id,
        hospital_a=hospital_a,
        hospital_b=hospital_b,
    ))

    url = f"https://hf.space/owkin/trainer-{experiment_id}"

    return (
        gr.Button.update(interactive=False),
        gr.Markdown.update(
            visible=True,
            value=f"Your experiment is available at [hf.space/owkin/trainer-{experiment_id}]({url})! - If the image does not build in under a minute, please refresh and try again"
        )
        
    )


demo = gr.Blocks(theme=theme, css="""\
@font-face {
    font-family: "Didact Gothic";
    src: url('https://huggingface.co/datasets/NimaBoscarino/assets/resolve/main/substra/DidactGothic-Regular.ttf') format('truetype');
}

@font-face {
    font-family: "Inter";
    src: url('https://huggingface.co/datasets/NimaBoscarino/assets/resolve/main/substra/Inter-Regular.ttf') format('truetype');
}

h1 {
    font-family: "Didact Gothic";
    font-size: 40px !important;
}

p {
    font-family: "Inter";
}

.gradio-container {
    min-width: 100% !important;
}

.margin-top {
    margin-top: 20px;
}

.white {
    background-color: white;
}

.column {
    border-radius: 20px;
    padding: 30px;
}

.blue {
    background-image: url("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/substra-banner.png");
    background-size: cover;
}

.blue p {
    color: white !important;
}

.blue strong {
    color: white !important;
}

.info-box {
    background: transparent !important;
    border-radius: 20px !important;
    border-color: white !important;
    border-width: 4px !important;
    padding: 20px !important;
}
""")

with demo:
    gr.HTML("""
    <img src="https://raw.githubusercontent.com/substra/substra/main/Substra-logo-colour.svg" style="height: 2em;" />
    """)
    gr.Markdown("# Federated  Learning with Substra")
    with gr.Row():
        with gr.Column(scale=1, elem_classes=["blue", "column"]):
            gr.Markdown("Here you can run a **quick simulation of Federated Learning**.")
            gr.Markdown("Check out the accompanying [blog post](https://huggingface.co/blog/owkin-substra/) to learn more.")
            with gr.Box(elem_classes=["info-box"]):
                gr.Markdown("""\
                This space is an introduction to federated learning. \
                We will create new spaces soon where you will be able to control the models, datasets and \
                federation strategies.\
                """)
        with gr.Column(scale=3, elem_classes=["white", "column"]):
            gr.Markdown("""\
            Data scientists doing medical research often face a shortage of high quality and diverse data to \
            effectively train models. This challenge can be overcome by securely allowing training on protected \
            data through Federated Learning. [Substra](https://docs.substra.org/) is a Python based Federated \
            Learning software that enables researchers to easily train ML models on remote data regardless of the \
            ML library they are using or the data type they are working with.
            """)
            gr.Markdown("### Here we show an example of image data located in **two different hospitals**.")
            gr.Markdown("""\
            By playing with the distribution of data in the two simulated hospitals, you'll be able to compare how \
            the federated models compare with models trained on single datasets. The data used is from the \
            Camelyon17 dataset, a commonly used benchmark in the medical world that comes from \
            [this challenge](https://camelyon17.grand-challenge.org/). The sample below shows normal cells on the \
            left compared with cancer cells on the right.
            """)
            gr.HTML("""
            <img
            src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/substra-tumor.png"
            style="height: 300px; margin: auto;"
            />
            """)
            gr.Markdown("""\
            A problem often faced by researchers is that datasets lack the necessary amount of positive samples \
            (samples containing cancer tissues) that are needed to reliably classify cancer. In this interface you \
            can use the slider to control the percentage of negative and positive samples in each hospital. \
            Setting this slider to minimum will mean there are 0 positive samples, whereas 50 would mean that \
            half the dataset contains slides with positive tumor samples.\
            """)

            with gr.Row(elem_classes=["margin-top"]):
                hospital_a_slider = gr.Slider(
                    label="Percentage of positive samples in Hospital A",
                    value=80,
                )
                hospital_b_slider = gr.Slider(
                    label="Percentage of positive samples in Hospital B",
                    value=20,
                )
            launch_experiment_button = gr.Button(value="Launch Experiment 🚀")
            visit_experiment_text = gr.Markdown(visible=False)

    launch_experiment_button.click(
        fn=launch_experiment,
        inputs=[hospital_a_slider, hospital_b_slider],
        outputs=[launch_experiment_button, visit_experiment_text]
    )

demo.launch()