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]
]

# Função para gerar respostas
def gerar(inputs):
    # Coloque aqui a lógica para processar os inputs e gerar respostas
    pass

# Tutorial para orientar os usuários
tutorial = """
Bem-vindo ao Medchat! Esta ferramenta foi projetada para ajudá-lo a analisar casos médicos e fornecer recomendações relevantes. Aqui estão algumas dicas para começar:

1. Escolha um exemplo de caso médico na lista de exemplos.
2. Leia o caso cuidadosamente e identifique as informações relevantes.
3. Use essas informações para formular uma análise ou recomendação adequada.
4. Se necessário, forneça detalhes adicionais ou faça perguntas para obter mais informações.
5. Após revisar o caso e suas recomendações, sinta-se à vontade para copiar o texto ou compartilhá-lo.

Estamos aqui para ajudar! Se você tiver alguma dúvida, não hesite em perguntar.
"""

# Interface do Chat com tutorial
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,
    introduction=tutorial,
    concurrency_limit=20,
).launch(show_api=False)