from huggingface_hub import InferenceClient
import gradio as gr

client = InferenceClient(
    "mistralai/Mixtral-8x7B-Instruct-v0.1"
)

# Função para formatar a prompt de entrada com o histórico de diálogo
def formatar_prompt(mensagem, historico):
    prompt = "<s>"
    for prompt_usuario, resposta_bot in historico:
        prompt += f"[INST] {prompt_usuario} [/INST]"
        prompt += f" {resposta_bot}</s> "
    prompt += f"[INST] {mensagem} [/INST]"
    return prompt

# Função para gerar resposta do modelo
def gerar(
    prompt, historico, prompt_sistema, temperatura=0.9, max_tokens_novos=256, top_p=0.95, penalidade_repeticao=1.0,
):
    temperatura = float(temperatura)
    if temperatura < 1e-2:
        temperatura = 1e-2
    top_p = float(top_p)

    kwargs_geracao = dict(
        temperature=temperatura,
        max_new_tokens=max_tokens_novos,
        top_p=top_p,
        repetition_penalty=penalidade_repeticao,
        do_sample=True,
        seed=42,
    )

    prompt_formatada = formatar_prompt(f"{prompt_sistema}, {prompt}", historico)
    stream = client.text_generation(prompt_formatada, **kwargs_geracao, stream=True, details=True, return_full_text=False)
    output = ""

    for resposta in stream:
        output += resposta.token.text
        yield output
    return output

# Inputs adicionais para o modelo
inputs_adicionais=[
    gr.Textbox(
        label="Prompt do Sistema",
        max_lines=1,
        interactive=True,
    ),
    gr.Slider(
        label="Temperatura",
        value=0.9,
        minimum=0.0,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Valores mais altos produzem saídas mais diversas",
    ),
    gr.Slider(
        label="Máximo de Novos Tokens",
        value=256,
        minimum=0,
        maximum=1048,
        step=64,
        interactive=True,
        info="O número máximo de novos tokens",
    ),
    gr.Slider(
        label="Top-p (amostragem de núcleo)",
        value=0.90,
        minimum=0.0,
        maximum=1,
        step=0.05,
        interactive=True,
        info="Valores mais altos amostram mais tokens de baixa probabilidade",
    ),
    gr.Slider(
        label="Penalidade de Repetição",
        value=1.2,
        minimum=1.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Penalize tokens repetidos",
    )
]

# Exemplos de prompts
exemplos=[
    ["Analise o caso de um paciente de 55 anos com histórico familiar de doença cardíaca. Ele apresenta um colesterol total de 260 mg/dL, HDL de 40 mg/dL, LDL de 180 mg/dL e triglicerídeos de 200 mg/dL. Interpretar os resultados e sugerir intervenções dietéticas e farmacológicas.", None, None, None, None, None],
    ["Examine os resultados de um ecocardiograma em um paciente de 70 anos que se queixa de dispneia aos esforços. O exame revela uma fração de ejeção de 40%, espessamento do septo interventricular e dilatação atrial esquerda. Propor um plano de tratamento e acompanhamento para insuficiência cardíaca.", None, None, None, None, None],
    ["Avalie os achados de um teste de esforço em esteira em um paciente de 45 anos com dor torácica atípica. Ele atingiu 10 METs, mas desenvolveu depressão do segmento ST de 2 mm no pós-exercício. Recomendar próximos passos diagnósticos e terapêuticos.", None, None, None, None, None],
    ["Interpretar os resultados de um Holter de 24 horas em um paciente de 60 anos com síncope recorrente. O registro mostra episódios de pausas sinusais de até 5 segundos e arritmia sinusal respiratória. Formular um plano de monitoramento e tratamento para bradicardia sintomática.", None, None, None, None, None],
    ["Analise os resultados de um cateterismo cardíaco em um paciente de 65 anos com angina instável. A angiografia revela lesões significativas em duas artérias coronárias principais, com uma fração de fluxo fracionada (FFR) de 0,68. Recomendar opções de revascularização e manejo da doença coronariana.", None, None, None, None, None]
]

# Tutorial em HTML
tutorial_html = """
<h2>Guia de Uso</h2>
<h3>1. Introdução</h3>
<p>Bem-vindo ao Medchat, sua ferramenta de automação inteligente de saúde! Este guia foi projetado para ajudá-lo a utilizar esta interface de maneira eficaz.</p>
<h3>2. Como usar</h3>
<p>Siga estes passos simples para aproveitar ao máximo o Medchat:</p>
<ol>
  <li>Insira sua consulta médica na caixa de texto e pressione "Enviar" para receber uma resposta do chatbot.</li>
  <li>Utilize os exemplos pré-definidos clicando neles para carregar automaticamente na caixa de texto.</li>
  <li>Explore as respostas do chatbot e siga as recomendações fornecidas.</li>
</ol>
<h3>3. Recursos Adicionais</h3>
<p>Além disso, você pode aproveitar os recursos adicionais, como botões de compartilhamento e cópia, para compartilhar ou salvar suas consultas.</p>
<h3>4. Ajuda Adicional</h3>
<p>Se precisar de assistência adicional, clique no botão de ajuda ou entre em contato conosco através dos canais de suporte fornecidos.</p>
<button onclick="fecharTutorial()">Fechar</button>
<script>
  function fecharTutorial() {
    document.getElementById("tutorial").style.display = "none";
  }
</script>
"""

# Interface do Chat
gr.ChatInterface(
    fn=gerar,
    chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
    additional_inputs=inputs_adicionais,
    title="Medchat | Automação Inteligente de Saúde",
    examples=exemplos,
    concurrency_limit=20,
    tutorial_html=tutorial_html
).launch(show_api=False)