import pandas as pd
import streamlit as st

from utils_casemaker import CaseMaker, format_casemaker_data

st.set_page_config(layout="wide")

st.title("Juni Health Patient Casemaker")

casemaker = CaseMaker("terms.json")

uploaded_file = st.file_uploader("Choose a file")

if uploaded_file is not None:
    # Can be used wherever a "file-like" object is accepted:
    df = pd.read_csv(uploaded_file)
    reports = format_casemaker_data(
        df=df,
        patient_id_column="patient_id",
        date_column="report_id",
        text_column="text",
        )
    
    patient_options = {
        f"Patient {patient_id} ({len(reports[patient_id])} reports)": patient_id 
        for patient_id in reports.keys()
    }
    selected_patient_string = st.radio(
        "Select a Patient ID",
        list(patient_options.keys()),
        key = "patient_select_button"
        )
    
    if st.button("Generate Case", key = "task_begin_button"):
        selected_patient_id = patient_options[selected_patient_string]
        summary_by_organ = casemaker.parse_records(reports[selected_patient_id])
        summary_by_organ = casemaker.format_reports(summary_by_organ)

        # Display the report
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("Original")
            display_table = pd.DataFrame.from_records([item.dict() for item in reports[selected_patient_id]])
            display_table = display_table[["date", "text"]]
            display_table["text"] = display_table["text"].apply(lambda s: casemaker.remove_header_names(s))
            display_table = display_table.rename(columns={"date": "ID/Date", "text": "Report"})
            st.table(display_table)
                
        with col2:
            st.subheader("With Casemaker")
            for chosen_organ in summary_by_organ.keys():
                if summary_by_organ[chosen_organ]:
                    st.header(chosen_organ.capitalize())
                    st.write(summary_by_organ[chosen_organ])