# Simple APP for specialty pharmacy 

# Import packages
import numpy as np 
import os 

import gradio as gr
from transformers import pipeline

#Import LLMs 
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI

# Prompt template 
from langchain import PromptTemplate

# Chains
from langchain.chains import LLMChain


# Import "secret" OPENAI_API_KEY
os.environ["OPENAI_API_KEY"] 

# Import GPT-4
llm_gpt = ChatOpenAI(model='gpt-4-0613',temperature=0.)

# ======================================================
# Set up an ASR pipeline using facebook's wav2vec2
p = pipeline("automatic-speech-recognition", chunk_length_s=40)

# =======================================================
# LLM Chains

# Dialogue chain
template_diag = """
You are an AI assistant with medical language understanding.

The input is a dialogue between a specialty pharmacist and patient: {input}

To give you context, the dialogue will have to do about symptoms, side effects, medications etc
of a rare disease, most probably multiple sclerosis.

You have a couple of tasks:

- First: If there are some non-sensical words, convert them to the most probable real word,
taking into account that this is a pharmaxist, so most of them should describe medical conditions
or symptoms, most probably about multiple sclerosis.
If a medication is mentioned, do your best to find which is that, if any. Correct any mispellings
Capitalize the names of the medications.

- Second: Convert the text into a dialogue of the form:

[Pat]: 
[PRx]:

Where [PRx]: Pharmacist, [Pat]: Patient

Use your judgement to distinguish between the two roles and who said what.
Output only this dialogue.

Output:
"""


prompt_diag = PromptTemplate(template=template_diag, input_variables=["input"])
chain_diag = LLMChain(llm=llm_gpt, prompt=prompt_diag, verbose=False)


# ==============================================
template_struct = """
You are an AI assistant with medical language understanding.

The input is a dialogue between a specialty pharmacist and patient: {input}

To give you context, the dialogue will have to do about symptoms, side effects, medications etc
of a rare disease, most probably multiple sclerosis.

Some words may not be clearly spelled, because they come from an automatic 
audio to text transcript. 

Your have a few tasks:

- First task: If there are some non-sensical words, convert them to the most probable real word,
taking into account that this is a dialogue about a medical condition, probably multiple sclerosis

- Second task: extract information from this dialogue

Specifically the following:

- A brief summary of the dialogue, highlighting the chief complaint 
- The main disease mentioned by the patient
- Medications mentioned by the patient
- Side effets mentioned by the patient

The output should have the form of a json file with those four keys: (Summary, Disease, Medications, Side_Effects)

Do not hallucinate and do not make up information that is not included in the original file.

Output:
"""

# SOAP notes 
prompt_struct = PromptTemplate(template=template_struct, input_variables=["input"])
chain_struct = LLMChain(llm=llm_gpt, prompt=prompt_struct, verbose=False)


# Transcription function
def transcribe(audio):
    #text = fake_audio
    text = p(audio)["text"]
    output_1 = eval(chain_struct.run(text))
    output_2 = chain_diag.run(text)
    summa = output_1['Summary']
    disease = output_1['Disease']
    meds = output_1['Medications']
    sides = output_1['Side_Effects']
    return summa, disease, meds, sides, output_2  
        
    
# 
with gr.Blocks(title="AI specialty scriber",theme=gr.themes.Soft()) as demo:
    
    with gr.Row():
        image_wag = gr.Image(value="Walgreens_AI.png", width=10, show_label=False,show_download_button=False, scale=1)
    gr.Markdown("## <center> Walgreens AI-powered specialty pharmacy tool </center>")
    #gr.Markdown("**<center>"+scriber_description+"</center>**")
    gr.Markdown("<center> ________________________________________________________________________ </center>")
    
    # ====================================================
    # Dictation tool
    gr.Markdown("**Record Patient Interaction**")
    audio = gr.Audio(label='Your recording here',source="microphone", type="filepath",container=True)
    audio_submit_btn = gr.Button(value="Submit Recording", variant="primary")
    
    # Clinical notess and transcript
    with gr.Tab("Extracted Information"):
        with gr.Row():
            summary = gr.Textbox(label='Summary',lines=3,interactive=True) 
            disease = gr.Textbox(label='Disease mentioned',lines=3,interactive=True)
        with gr.Row():
            medications = gr.Textbox(label='Medications mentioned',lines=3,interactive=True)
            sides = gr.Textbox(label='Side Effects mentioned',lines=3,interactive=True)
        
    with gr.Tab("Original Transcript"):
        dialogue = gr.Textbox(label='Full conversation transcript',lines=10)
        
    # ===============================================
    # Submit and clear tool    
    audio_submit_btn.click(transcribe, inputs = audio, outputs=[summary,disease,medications,sides,dialogue])
    audio_clear_btn = gr.ClearButton([audio,summary,disease,medications,sides,dialogue])
    
    
demo.launch()