"""

Gradio related stuff

"""

from enum import Enum

from typing import Optional

import gradio as gr

from src import utils

class InjuryType(str, Enum):

    NO_INJURY = "NO_INJURY"

    LIVER = "LIVER"

    KIDNEY = "KIDNEY"

    SPLEEN = "SPLEEN"

    EXTRAVASATION = "EXTRAVASATION"

    BOWEL = "BOWEL"

INJURY_PROMPTS = {

    InjuryType.LIVER: "[INST] What are some first aid for liver internal injury [/INST]",

    InjuryType.KIDNEY: "[INST] What are some first aid for kidney internal injury [/INST]",

    InjuryType.SPLEEN: "[INST] What are some first aid for spleen internal injury [/INST]",

    InjuryType.EXTRAVASATION: "[INST] What are some first aid for extravasation internal injury [/INST]",

    InjuryType.BOWEL: "[INST] What are some first aid for bowel internal injury [/INST]",

    InjuryType.NO_INJURY: "[INST] What are some general medical fitness tips [/INST]",

}

# Define function to be executed when "Run" button is clicked

def interface_handler(choice: Optional[InjuryType]) -> str:

    if choice is None:

        return ""

    input_prompt = INJURY_PROMPTS[choice]

    return utils.query(inputs=input_prompt)

interface_description = """

<center>

<img src="/static/images/image-robot.jpg" alt="image of a robot" width="350px"/>

</center>

<strong>This application is based on the [RSNA 2023 Kaggle Competition](),

and attempts to provide first aid information to patients with abdominal injuries.</strong>

It does this by prompting the <em>[mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1)</em>

text generation model via the [HuggingFace Inference API](https://huggingface.co/inference-api).

This model is not medically accurate, and this entire application has been created for illustrative purposes.

Do enjoy!

"""

# Gradio interface configuration

iface = gr.Interface(

    fn=interface_handler,

    allow_flagging="never",

    title="Ask eMLily",

    description=interface_description,

    inputs=[

        gr.Radio(

            label="Get some first aid advice!",

            info="What is the patient's diagnosis?",

            choices=[

                ("Liver", InjuryType.LIVER),

                ("Bowel", InjuryType.BOWEL),

                ("Extravasation", InjuryType.EXTRAVASATION),

                ("Kidney", InjuryType.KIDNEY),

                ("Spleen", InjuryType.SPLEEN),

                (

                    "No Internal Injury. I'd just like general medical advice",

                    InjuryType.NO_INJURY,

                ),

            ],

        ),

    ],

    outputs=gr.Textbox("Model output ..."),

    live=True,

    theme="light",

)