# Import Transformers & Gradio
from datasets import load_dataset
from transformers import pipeline
import gradio as gr

# Load clinical note data

# Print the columns in this dataset and a sample entry.
from datasets import load_dataset
ds = load_dataset("AGBonnet/augmented-clinical-notes")

# Loads the JSON column and indicates whether the JSON was valid. Some records contain invalid JSON formatting.
def load_json_with_status(raw):
    import json
    try:
        parsed = json.loads(raw)
        return True, parsed
    except:
        return False, raw

# View samples of Clinical Notes

train = ds["train"]

# Sample up to 25 of today's patients since we never have more than that per day.
sampleLengthTrain = min(25, len(ds['train']))
print("length=", sampleLengthTrain)
patientSample = {}  # Dictionary keyed by patient ID

for i in range(sampleLengthTrain):

  # A JSON summary contains multiple sections describing the patient's problem and the prescribed tests.
  patientJsonSummaryStr = ds['train'][i]['summary']
  patientIdx = ds['train'][i]['idx']
  patientNote = ds['train'][i]['note']
  success, parsedJsonSummary = load_json_with_status(patientJsonSummaryStr);

  # Only if the JSON parsing succeeded do we load this record.
  if success:
    motivation = parsedJsonSummary['visit motivation']
    discharge = parsedJsonSummary['discharge']
    treatments = parsedJsonSummary['treatments']
    patientSample[patientIdx] = {'note': patientNote, 'visit motivation': motivation, 'summary': parsedJsonSummary, 'discharge': discharge, 'treatments': treatments}
    #(DEBUG ONLY) #print("idx:", ds['train'][i]['idx'], "summary:", repr(str(patientJsonSummaryStr)))

  # otherwise, skip the record with invalid JSON.
print("Done. patientSample=", len(patientSample))

# Create functions to use from the UI to show content in dropdowns and textboxes.

def get_motivation_from_patient_id(patientId):
    return patientSample.get(patientId)['visit motivation']

def get_note_from_patient_id(patientId):
    return patientSample.get(patientId)['note']

def get_treatment_plan_from_patient_id(patientId):
    print(patientSample.get(patientId))
    summary = patientSample.get(patientId)['summary']
    return ("Treatments: " + str(summary.get('treatments', 'N/A')) +
            " Discharge notes: " + str(summary.get('discharge', 'N/A')))


# Assist medical staff by summarizing treatment plan from their notes.

def json_section_to_text(json_obj):
    parts = []

    # Simple mapping that includes only non-empty fields.
    if json_obj.get("visit motivation"):
        motiviation = "Visit motivation: " + json_obj["visit motivation"]
        parts.append(motiviation)
        print(motiviation)
    if json_obj.get("patient information"):
        patientInfo = "Patient information: " + json_obj["patient information"]
        parts.append(patientInfo)
        print(patientInfo)

    return "\n\n".join(parts)

# Create a text generation pipeline to generate patient instructions.

from transformers import pipeline
generator = pipeline("text-generation", model="microsoft/MediPhi-Clinical", max_new_tokens=256, do_sample=True)

def summarize_notes(patient_id):
    treatment_text = get_treatment_plan_from_patient_id(patient_id)
    prompt = f"Based on this treatment plan, summarize key care instructions for the patient in 2-3 sentences: {treatment_text}"

    result = generator(prompt, return_full_text=False)
    return result[0]["generated_text"].strip()

# Build a UI using Gradio to list a sampling of patient's from Today's Office Visits.
#   - Once selected, display the Visit Motivation and the Clinical Summary
#   - A Generate

import gradio as gr

# Load patient IDs for the dropdown. The first item will be a blank so that the user must select a patient ID.
choices = [""] + list(patientSample.keys())

with gr.Blocks(title="Patient Instructions generated from Clinical Notes", theme="soft") as demo:
  gr.Markdown("# Patient Treatment Plan Summarizer");
  gr.Markdown("### Choose a patient to view their summary, and then generate concise patient instructions.")

  patientDropdown = gr.Dropdown(choices, label="Choose a Patient ID")
  motivationTextbox = gr.Textbox(lines=4, label="Visit Motivation")
  summaryTextbox = gr.Textbox(lines=4, label="Clinical summary")
  patientDropdown.change(get_motivation_from_patient_id, inputs=patientDropdown, outputs=motivationTextbox)
  patientDropdown.change(get_note_from_patient_id, inputs=patientDropdown, outputs=summaryTextbox)

  # Button to generate summarized patient instructions and display them.
  patientInstructions = gr.Textbox(lines=5, label="Simplified Next-Steps")
  generateBtn = gr.Button("Get simplified Next Steps for Patient")
  generateBtn.click(summarize_notes, inputs=patientDropdown, outputs=patientInstructions)


demo.queue().launch(share=True)