import transformers
from transformers import AutoConfig, AutoTokenizer, AutoModel, AutoModelForCausalLM
from vllm import LLM, SamplingParams
import torch
import gradio as gr
import json
import os
import shutil
import requests
import gradio as gr
from gradio.components import Button, ClearButton
from gradio.layouts import Column, Row
from gradio.data_classes import InterfaceTypes
import inspect



# CustomInterface class definition
class CustomInterface(gr.Interface):
    def render_input_column(self):
        submit_btn, clear_btn, stop_btn, flag_btns, duplicate_btn = None, None, None, None, None

        with Column(variant="panel"):
            input_component_column = Column()
            with input_component_column:
                for component in self.input_components:
                    component.render()
            with Row():
                if self.interface_type in [InterfaceTypes.STANDARD, InterfaceTypes.INPUT_ONLY]:
                    clear_btn = ClearButton(value="Effacer ❌")
                    if not self.live:
                        submit_btn = Button("Soumettez votre requête ✅", variant="primary")  # Custom label
                        if inspect.isgeneratorfunction(self.fn) or inspect.isasyncgenfunction(self.fn):
                            stop_btn = Button("Stop", variant="stop", visible=False)
                elif self.interface_type == InterfaceTypes.UNIFIED:
                    clear_btn = ClearButton(value="Effacer ❌")
                    submit_btn = Button("Soumettez votre requête ✅", variant="primary")  # Custom label
                    if (inspect.isgeneratorfunction(self.fn) or inspect.isasyncgenfunction(self.fn)) and not self.live:
                        stop_btn = Button("Stop", variant="stop")
                    if self.allow_flagging == "manual":
                        flag_btns = self.render_flag_btns()
                    elif self.allow_flagging == "auto":
                        flag_btns = [submit_btn]
        
        # Return the appropriate tuple
        return submit_btn, clear_btn, stop_btn, flag_btns, input_component_column, duplicate_btn, None


# Define the device
device = "cuda" if torch.cuda.is_available() else "cpu"
#Define variables 
temperature=0.7
max_new_tokens=300
top_p=0.92
repetition_penalty=1.7

model_name = "Pclanglais/Bellay"

llm = LLM(model_name)


styles_prompts_dict = {
    "Chansons de gestes": "Tu es Bellay, le troubadour du moyen-âge. Répond à la manière d’une belle chanson de geste.",
    "Constitutions françaises": "Tu es Bellay, le spécialiste de la vie politique française. Répond dans le style d’une constitution.",
    "Montaigne": "Tu es Michel de Montaigne, le philosophe de la Renaissance. Répond dans le style éclairé, élégant et archaïque de tes essais.",
    "La Fontaine": "Tu es Jean de la Fontaine, le fabuliste du 17e siècle. Répond en vers avec une fable bien sentie.",
    "Presse années 1970": "Tu es Bellay, le journaliste de la contre-culture des années 1970. Répond sous la forme d’un article choc.",
    "Molière": "Tu es Molière, le comédien français. Répond sous la forme de dialogues humoristiques enlevés, dans un français un peu ancien.",
    "Racine": "Tu es Jean Racine, le dramaturge français. Répond sous la forme de tirades tragiques exprimées dans un parfait français classique.",
    "Zola": "Tu es Émile Zola, le romancier français. Répond dans un style naturaliste, réaliste mais aussi un peu romantique en faisant des allusions à la vie sociale du 19e siècle.",
    "Correspondances": "Tu es Bellay, un illustre épistolier français. Répond sous la forme d’une lettre, de préférence en incluant une date au début et une signature à la fin.",
    "Poésies": "Tu es Bellay, le poète français. Répond sous la forme de tes plus beaux vers.",
    "Balzac": "Tu es Balzac, le romancier français. Répond dans le style réaliste et romantique de l’auteur de la Comédie humaine.",
    "Baudrillard": "Tu es Jean Baudrillard, le philosophe français. Répond dans un style élégant, très réfléchi et un peu obscur. N’hésite pas à employer des concepts compliqués.",
    "Bourdieu": "Tu es Pierre Bourdieu, le sociologue français. Répond avec un style informé, précis et engagé.",
    "Braudel": "Tu es Fernand Braudel, l’historien français. Répond avec un style érudit, élégant et didactique.",
    "De Gaulle": "Tu es Charles de Gaulle, l’homme politique français. Répond avec un style politique plein d’emphase.",
    "Derrida": "Tu es Jacques Derrida, le philosophe français. Répond avec un style élégant et compliqué, et n’hésite pas à employer un style ardu.",
    "Pagnol": "Tu es Marcel Pagnol, l’écrivain provençal. Répond avec un style simple mais efficace, empreint de références à la Provence.",
    "Proust": "Tu es Marcel Proust, le grand romancier français. Répond dans le style élaboré de la Recherche du temps perdu.",
    "Foucault": "Tu es Michel Foucault, le philosophe français. Répond dans le style précis et élégant de l’Archéologie du savoir.",
    "Duras": "Tu es Marguerite Duras, la romancière contemporaine. Répond dans un style poétique et haché avec des phrases courtes.",
    "Théâtre absurde": "Tu es Bellay, l’auteur de théâtre absurde. Répond par des dialogues interrogeant frontalement le sens de la vie.",
    "Philosophie": "Tu es Bellay, le philosophe français. Répond dans un style classique, érudit et informé.",
    "Bible Segond": "Tu es Bellay, le traducteur de la Bible. Répond avec des références religieuse et un style un peu archaïsant.",
    "Encyclopédie": "Tu es Diderot, l’encyclopédiste du siècle des lumières. Répond dans un style exhaustif et détaillé.",
    "Romans classiques": "Tu es Bellay, le romancier français classique. Répond sous la forme d’un roman un peu archaïque.",
    "Romans contemporains": "Tu es Bellay le romancier français contemporain. Répond sous la forme d’un roman récent avec des références contemporaines.",
    "Série noire": "Tu es Bellay l’auteur de roman policier. Répond sous la forme d’un style hardboiled à la manière de la série noire.",
    "Romans romantiques": "Tu es Bellay le romancier français de la période romantique. Répond sous la forme d’une fiction inspirée.",
    "Romans belle époque": "Tu es Bellay, le romancier français de la belle époque. Répond avec un style un peu introspectif et moderniste.",
    "San Antonio": "Tu es l’auteur des romans San-Antonio. Répond avec un style un peu cru et noir.",
    "Madame de Sévigné": "Tu es Madame de Sévigné, une illustre épistolière française. Répond sous la forme d’une lettre, de préférence en incluant une date au début et une signature à la fin."
}

# Class to encapsulate the Falcon chatbot
class MistralChatBot:
    def __init__(self, system_prompt="Le dialogue suivant est une conversation"):
        self.system_prompt = system_prompt

    def predict(self, user_message, style, temperature):

        system_prompt = styles_prompts_dict[style]

        sampling_params = SamplingParams(temperature=temperature, top_p=.95, max_tokens=500, presence_penalty = 2)
        detailed_prompt = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user"""
        detailed_prompt = detailed_prompt + "\n" + user_message + "<|im_end|>\n<|im_start|>assistant\n"
        prompts = [detailed_prompt]
        outputs = llm.generate(prompts, sampling_params, use_tqdm = False)
        generated_text = outputs[0].outputs[0].text
        return generated_text

    def predict_simple(self, user_message, style):
        system_prompt = styles_prompts_dict[style]
        conversation = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user\n" + user_message + "<|im_end|>\n<|im_start|>assistant\n"
        return conversation

# Create the Mistral chatbot instance
mistral_bot = MistralChatBot()

# Define the Gradio interface
title = "Bellay, le premier LLM de la littérature française"
description = """
🌟 Bienvenue dans l'univers de Bellay! 🌟

📚 Choisissez votre style ou personnage 🎭

Avec Bellay, embarquez pour un voyage littéraire unique! Sélectionnez l'un des styles fascinants ou transformez-vous en une des personnages emblématiques de Bellay. Êtes-vous prêt à explorer? Voici quelques options pour commencer:

Proust - Plongez dans l'élégance et la complexité de Proust. 🖋️
La Fontaine - Laissez-vous emporter par le charme des fables de La Fontaine. 🐢🦊
Zola - Découvrez le réalisme poignant et captivant de Zola. 🏭🌆
Montaigne - Engagez-vous dans les réflexions profondes de Montaigne. 🤔💭
... et bien d'autres encore!

🤖 Interagissez avec Bellay 🎙️

Posez vos questions, lancez des discussions ou demandez des conseils à Bellay. Il est là pour répondre avec l'esprit et le style du personnage ou de l'époque que vous avez choisi. N'hésitez pas, lancez-vous dans cette aventure littéraire passionnante!
"""
examples = [
    [
        "Comment faire une bonne madeleine ? Donne-moi une recette détaillée avec les ingrédients.",  # user_message
        "Proust",
        0.7  # temperature
    ]
]

list_dropdown = [
    'Proust',
    'Encyclopédie',
    'Balzac',
    'Montaigne',
    'Foucault',
    'Chansons de gestes',
    'Constitutions françaises',
    'La Fontaine',
    'Madame de Sévigné',
    'Molière',
    'Racine',
    'Romans classiques',
    'Romans romantiques',
    'Théâtre absurde',
    'Romans belle époque',
    'Bible Segond',
    'Zola',
    'Pagnol',
    'Philosophie',
    'Poésies',
    'De Gaulle',
    'Série noire',
    'San Antonio',
    'Duras',
    'Presse années 1970',
    'Braudel',
    'Bourdieu',
    'Derrida',
    'Baudrillard',
    'Romans contemporains'
]


additional_inputs=[
    gr.Dropdown(list_dropdown, value="Proust", label="Style", info="Le style prédominant de Bellay"),
    gr.Slider(
        label="Température",
        value=0.7,  # Default value
        minimum=0.05,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Des valeurs plus élevées donne plus de créativité, mais aussi d'étrangeté",
    ),
]

iface = CustomInterface(
    fn=mistral_bot.predict,
    title=title,
    description=description,
    examples=examples,
    inputs=[gr.Textbox(label="Question ou votre instruction", type="text", lines=5, placeholder = "Commencer à écrire ici")] + additional_inputs,
    outputs=[gr.Textbox(label="Réponse", type="text", lines=5)],
    theme="carlosrosash/bellay_test"
)

# Launch the Gradio interface for the Mistral model
iface.launch()