import openai
import re
import gradio as gr

# Define OpenAI API key

model_engine = "text-davinci-003"

def GPT_answer(prompt):
    # Generate a response
    completion = openai.Completion.create(
        engine=model_engine,
        prompt=prompt,
        temperature=0.3,
        max_tokens=500,
        top_p=1,
        best_of=1,
        frequency_penalty=0,
        presence_penalty=0
    )
    response = completion.choices[0].text
    return response

def text_preparation(text):
    text = re.sub('[\n+\t+]', ' ', text)
    text = re.sub('\\\m[0-9]+', '', text)
    text = re.sub('"', '', text)
    text = re.sub('\s{2,}', ' ', text)
    return text

def prompt_creation_icd(clean_text):
    prompt = ' What are ICD-10 codes for these diagnoses:\n'+clean_text+'.\n'
    return prompt

def prompt_creation_dia(clean_text):
    prompt = " Extract patient's diagnoses from this text:\n"+clean_text+".\n"
    prompt += "Print each diagnose from the new line."
    return prompt

def get_result(text):
    clean_text = text_preparation(text)
    prompt_dia = prompt_creation_dia(clean_text)
    try:
        response_dia = GPT_answer(prompt_dia)
        #creating ICD codes based on found diagnoses
        diagnoses = response_dia.split('\n')
        diagnoses_cleaned = "\n".join([dia for dia in diagnoses if 'kein' not in dia.lower()])
        prompt_icd = prompt_creation_icd(diagnoses_cleaned)
        response_icd = GPT_answer(prompt_icd)
        answer = 'Found DIAGNOSES:'+diagnoses_cleaned+'\n\nICD-10 CODES:'+response_icd
        return answer
    except openai.error.InvalidRequestError as e:
        # Handle the InvalidRequestError
        return"This model's maximum context length is 4097 tokens error"

print('INFO: starting gradio interface')
default_input_text = """Report: 0121874648 02.01.2019 0121874648 04.01.2019"\m1Therapie: \m002.01.2019 bis 06.01.2019 i.v.-Antibiose mit Ampicillin/Sulbactam\m1  Anamnese:\m0  Die stationäre Übernahme der Patientin erfolgte aus der hiesigen Klinik für Kardiologie/Pneumologie bei akuter Tonsillitis mit positivem Nachweis von A-Streptokokken. Bezüglich der expliziten Anamnese verweisen wir auf den Entlassungsbericht der genannten Klinik vom 03.01.2019.   \m1Befunde:\m0  HNO-Status: Ohren: unauffällig. Nase: unauffällig. Mundrachen: Tonsillen bds. hyperplastisch, gerötet, follikuläre Beläge li. > re. Epipharynx: unauffällig. Oropharynx: unauffällig. Hypopharynx: unauffällig. Larynx: unauffällig.  Paraklinik: CRP vom 03.01.2019: 76,7 mg/l, Leukozyten 8,3 Gpt/l.  CRP vom 5.1.2019: 12.5mg/l.  LZ-RR: s. Anlage.  \m1Verlauf:\m0  Unter o. g. Antibiose sowie symptomatischer Therapie sahen wir einen regelrechten Rückgang der Beschwerden sowie der paraklinischen Infektparameter. Bezüglich der internistisch vordiagnostizierten hypertensiven Entgleisung bitten wir um weitere amb. internistische Kontrollen und ggf. Therapieoptimierung  Wir entließen Frau Kaczmarek am 06.01.2019 in Ihre weitere geschätzte ambulante Betreuung und bitten um Fortführung der Antibiose (z. B. Sultamicillin 375 mg 2 x tgl. per os) für weitere 2 Tage sowie um Kontrolle der Lokalbefunde in Ihrer Sprechstunde.\m0" Mit freundlichen kollegialen Grüßen"""
iface = gr.Interface(
    enable_queue=True,
    title="ICD-10 codes and Diagnoses Extraction",
    description="",
    fn=get_result,
    inputs=[gr.Textbox(label="Input text", value=default_input_text)],
    outputs=gr.outputs.Textbox(label="Found ICD-10 codes and diagnoses"),
)
iface.launch()