import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Charger le modèle et le tokenizer
model_name = "MaziyarPanahi/BioMistral-7B-GGUF"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model = model.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(inputs['input_ids'], max_length=150, num_return_sequences=1, no_repeat_ngram_size=2)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def add_message(history, message):
    if message["text"] is not None:
        history.append((message["text"], None))
    for x in message["files"]:
        history.append(((x,), None))
    return history, gr.MultimodalTextbox(value=None, interactive=False)

def bot(history):
    if history and history[-1][0]:
        history[-1] = (history[-1][0], generate_response(history[-1][0]))
    return history

def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)

# Création de l'interface Gradio avec le fond animé et des boutons stylisés
with gr.Blocks(css="""

.gradio-container {
     background: url('https://st4.depositphotos.com/8211188/25405/v/450/depositphotos_254059962-stock-illustration-abstract-medical-background-with-flat.jpg')50% 50% no-repeat;
        background-size: cover;
               }
    .chatbox-container {
        max-width: 80%;
        margin: 20px auto;
        padding: 20px 20px;
        background-color: rgb(39 150 160);
        border-radius: 12px;
        box-shadow: 0 0 20px rgba(0, 0, 0, 0.1);
        display: flex;
        flex-direction: row;
        align-items: stretch;
    }

    .chatbox {
        flex: 1;
        overflow-y: auto;
        padding: 10px;
        display: flex;
        flex-direction: column;
        justify-content: flex-end;
        border-bottom: 1px solid #e39d05;
        height: 400px; 
    }
    
    .chat-input-container {
        display: flex;
        flex-direction: column;
        padding: 10px;
        border-top: 1px solid #e39d05;
        width: 100%;
    }
    
    .chat-input {
    color:blue;
        margin-bottom: 10px;
        flex: 1;
        border-radius: 5px;
        border: 1px solid #e39d05;  
        padding: 10px;
        font-size: 16px;
    }
    
    .button-container {
        color:#e39d05;
        display: flex;
        flex-direction: row;
        justify-content: flex-start;
    }
    
    .button {
        background-color: #e39d05; 
        color: black;
        border: none;
        border-radius: 20px;
        padding: 8px 16px;
        margin: 12px;
        cursor: pointer;
        font-size: 20px;
        display: flex;
        align-items: center;
        justify-content: center;
        transition: background-color 0.3s, box-shadow 0.3s;
    }
    
    .button:hover {
        background-color: #fa0a0a;
        box-shadow: 0 4px 8px #e39d05;
    }

    
    
    
    .titre h1 {
        font-family: 'Centaur', serif;
        font-size: 4em;
        margin: 0;
         color: #ad2727;
         text-align: center;
    }
    
    .titre p {
        font-size: 3em;
        font-family: 'Centaur', serif;
        margin-top: 10px;
        
            color: rgb(9 129 118);
         text-align: center;
               font-weight: bold;
    }
    .titre img{
               display: block;
            margin-left: auto;
            margin-right: auto;
            width: 20%;
               }
    
""") as demo:
    
    with gr.Row(elem_classes="titre"):
        
        gr.Markdown("<h1>Diagnostique médicale</h1><p class='description'>Bienvenue ! Entrez vos symptômes ou questions pour des conseils médicaux rapides</p><img src='https://imageio.forbes.com/specials-images/imageserve/64b54b7467fcc06271e9bcff/Chatbot-in-a-medical-cap--a-pen-and-a-notebook-in-his-hands-asks-how-he-can-help-/960x0.jpg?height=592&width=711&fit=bounds'>")
        

    with gr.Column(scale=1, elem_classes="chatbox-container"):
        with gr.Row():
            chatbot = gr.Chatbot(
                elem_id="chatbot",
                bubble_full_width=False,
                scale=1,
                elem_classes="chatbox"
            )
        
        with gr.Row(elem_classes="chat-input-container"):
            chat_input = gr.MultimodalTextbox(interactive=True,
                                              file_count="multiple",
                                              placeholder="Entrez un message ou téléchargez un fichier...", 
                                              show_label=False,
                                              elem_classes="chat-input")
            
            # Container for buttons below the input
            with gr.Row(elem_classes="button-container"):
                clear_button = gr.Button("Effacer", elem_classes="button")
                stop_button = gr.Button("Arrêter", elem_classes="button")
                generate_button = gr.Button("Générer", elem_classes="button")

                # Configuration des interactions
                chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
                bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
                bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

                clear_button.click(lambda: ([], gr.MultimodalTextbox(value=None, interactive=True)), None, [chatbot, chat_input])
                stop_button.click(lambda: "Arrêter cliqué", None, None)
                generate_button.click(lambda: "Générer cliqué", None, None)
                
    save_button = gr.Button("Sauvegarder", elem_classes="button")
    save_button.click(fn=lambda history: open("discussion_history.txt", "w").write(str(history)), inputs=chatbot, outputs=None)

# Lancement de l'interface
demo.launch()