import streamlit as st
import pandas as pd
import numpy as np
from math import ceil
from collections import Counter
from string import punctuation
import spacy
from negspacy.negation import Negex
from spacy import displacy
from spacy.lang.en import English
from spacy.matcher import PhraseMatcher
from spacy.tokens import Span
#import en_ner_bc5cdr_md
import re
from streamlit.components.v1 import html

if "load_state" not in st.session_state: 
    st.session_state.load_state = False

if "button_clicked" not in st.session_state: 
    st.session_state.button_clicked = False
    
if "daily_button_clicked" not in st.session_state: 
    st.session_state.daily_button_clicked = False

if "past_button_clicked" not in st.session_state: 
    st.session_state.past_button_clicked = False
    
       
#nlp = en_core_web_lg.load()
nlp = spacy.load("en_ner_bc5cdr_md")

st.set_page_config(page_title ='Patient Inpatient Progression Dashboard', 
                   #page_icon= "Notes",
                   layout='wide')
st.title('Patient Inpatient Progression Dashboard')
st.markdown(
    """
    <style>
    [data-testid="stSidebar"][aria-expanded="true"] > div:first-child {
        width: 400px;
    }
    [data-testid="stSidebar"][aria-expanded="false"] > div:first-child {
        width: 400px;
        margin-left: -230px;
    }
    </style>
    """,
    unsafe_allow_html=True,
)
st.sidebar.markdown('Using transformer model')

## ======== Loading dataset ========
## Loading in Admission Dataset
df = pd.read_csv('shpi25nov.csv')
df.sort_values(by='SUBJECT_ID',ascending = True, inplace=True)

# Loading in Admission chief Complaint and diagnosis
df2 = pd.read_csv('cohort_cc_adm_diag.csv')

# Loading in Dischare History
df3 = pd.read_csv('cohort_past_history_12072022.csv')
df3.sort_values(by='CHARTDATE',ascending = False, inplace=True)

# Loading in Daily Narrative 
df4 = pd.read_csv('24houreventsFulltextwdifference.csv')
df4.sort_values(by=['SUBJECT_ID','HADM_ID','STORETIME'],ascending = True, inplace=True)


# combining both data into one 
df = pd.merge(df, df2, on=['HADM_ID','SUBJECT_ID'])

# Deleting admission chief complaint and diagnosis after combining
del df2

# Remove decimal point from Admission ID
df['HADM_ID'] = df['HADM_ID'].astype(str).apply(lambda x: x.replace('.0',''))
df3['HADM_ID'] = df3['HADM_ID'].astype(str).apply(lambda x: x.replace('.0',''))
df4['HADM_ID'] = df4['HADM_ID'].astype(str).apply(lambda x: x.replace('.0',''))
df3['INDEX_HADM_ID'] = df3['INDEX_HADM_ID'].astype(str).apply(lambda x: x.replace('.0',''))
df3["CHARTDATE_HADM_ID"] = df3["CHARTDATE"].astype(str) +' ('+ df3["HADM_ID"] +')'
df3["DIAGNOSIS"] = df3["DIAGNOSIS"].str.capitalize()
df3["DISCHARGE_LOCATION"] = df3["DISCHARGE_LOCATION"].str.capitalize()

df3["TEXT"] =df3["TEXT"].replace(r'\n','  \n ', regex=True)
df3["TEXT"] =df3["TEXT"].replace(r'#',' ', regex=True) 
df3["BertSummarizer"] =df3["BertSummarizer"].replace(r'#',' ', regex=True) 


#Renaming column
df.rename(columns={'SUBJECT_ID':'Patient_ID',
                  'HADM_ID':'Admission_ID',
                  'hpi_input_text':'Original_Text',
                  'hpi_reference_summary':'Reference_text'}, inplace = True)
df3.rename(columns={'SUBJECT_ID':'Patient_ID',
                   'HADM_ID':'PAST_Admission_ID',
                   'INDEX_HADM_ID':'Admission_ID'}, inplace = True) 

df4.rename(columns={'SUBJECT_ID':'Patient_ID',
                   'HADM_ID':'Admission_ID',
                   'Full_24_Hour_Events':'Full Text'}, inplace = True) 


#Filter selection 
st.sidebar.header("Search for Patient:")

# ===== Initial filter for patient and admission id =====
patientid = df['Patient_ID'].unique()
patient = st.sidebar.selectbox('Select Patient ID:', patientid)    #Filter Patient
admissionid = df['Admission_ID'].loc[df['Patient_ID'] == patient]  #Filter available Admission id for patient
HospitalAdmission = st.sidebar.selectbox(' ', admissionid)   
pastHistoryEpDate = df3['CHARTDATE_HADM_ID'].loc[(df3['Patient_ID'] == patient) & (df3['Admission_ID']== HospitalAdmission)] 
countOfAdmission = len(pastHistoryEpDate)


# List of Model available
model = st.sidebar.selectbox('Select Model', ('BertSummarizer','BertGPT2','t5seq2eq','t5','gensim','pysummarizer'))

# ===== to display selected patient and admission id on main page
col3,col4 = st.columns(2) 
patientid = col3.write(f"Patient ID:  {patient} ")
admissionid =col4.write(f"Admission ID:  {HospitalAdmission} ")

runtext = ''
inputNote  ='Input note here:'
# Query out relevant Clinical notes
original_text =  df.query(
    "Patient_ID  == @patient & Admission_ID == @HospitalAdmission"
)

original_text2 = original_text['Original_Text'].values
AdmissionChiefCom = original_text['Admission_Chief_Complaint'].values
diagnosis =original_text['DIAGNOSIS'].values
reference_text = original_text['Reference_text'].values

#dailyNoteChange = df4['_24_Hour_Events'].loc[(df4['Patient_ID'] == patient) & (df3['Admission_ID']== HospitalAdmission)] 
dailyNoteChange =df4[['STORETIME','_24_Hour_Events']].loc[(df4['Patient_ID'] == patient) & (df4['Admission_ID']==HospitalAdmission) & df4['_24_Hour_Events'].notnull()]
                        
dailyNoteChange.rename(columns={'STORETIME':'Time of Record',
                   '_24_Hour_Events':'Note Changes'}, inplace = True) 
dailyNote = df4['Full Text'].loc[(df4['Patient_ID'] == patient) & (df4['Admission_ID']==HospitalAdmission)]
dailyNote = dailyNote.unique()

    
##========= Buttons to the 5 tabs ======== Temp disabled Discharge Plan and Social Notes 
##col1, col2, col3, col4, col5 = st.columns([1,1,1,1,1]) -- to uncomment and comment below line to include discharge plan and social notes 
col1, col2, col5 = st.columns([1,1,1])
col6, col7 =st.columns([2,2])
with st.container():
    with col1:
        btnAdmission = st.button("🏥 Admission")
        inputNote = "Input Admission Note"
 
    with col2:
        btnDailyNarrative = st.button('📆Daily Narrative')
#    with col3:what
#        btnDischargePlan = st.button('🗒️Discharge Plan')
#        if btnDischargePlan:
#            inputNote = "Input Discharge Plan"
#    with col4:
#        btnSocialNotes = st.button('📝Social Notes')
#        if btnSocialNotes: 
#            inputNote = "Input Social Note"
    with col5:
        btnPastHistory = st.button('📇Past History (6 Mths)')



##======================== Start of NER Tagging ========================
                 
#lemmatizing the notes to capture all forms of negation(e.g., deny: denies, denying)
def lemmatize(note, nlp):
    doc = nlp(note)
    lemNote = [wd.lemma_ for wd in doc]
    return " ".join(lemNote)

#function to modify options for displacy NER visualization
def get_entity_options():
    entities = ["DISEASE", "CHEMICAL", "NEG_ENTITY"]
    colors = {'DISEASE': 'pink', 'CHEMICAL': 'orange', "NEG_ENTITY":'white'}
    options = {"ents": entities, "colors": colors}    
    return options

#adding a new pipeline component to identify negation
def neg_model():
    nlp.add_pipe('sentencizer')
    nlp.add_pipe(
    "negex",
    config={
        "chunk_prefix": ["no"],
    },
    last=True)
    return nlp

def negation_handling(note, neg_model):
    results = []
    nlp = neg_model() 
    note = note.split(".") #sentence tokenizing based on delimeter 
    note = [n.strip() for n in note] #removing extra spaces at the begining and end of sentence
    for t in note:
        doc = nlp(t)
        for e in doc.ents:
            rs = str(e._.negex)
            if rs == "True": 
                results.append(e.text)
    return results

#function to identify span objects of matched negative phrases from text
def match(nlp,terms,label):
    patterns = [nlp.make_doc(text) for text in terms]
    matcher = PhraseMatcher(nlp.vocab)
    matcher.add(label, None, *patterns)
    return matcher

#replacing the labels for identified negative entities
def overwrite_ent_lbl(matcher, doc):
    matches = matcher(doc)
    seen_tokens = set()
    new_entities = []
    entities = doc.ents
    for match_id, start, end in matches:
        if start not in seen_tokens and end - 1 not in seen_tokens:
            new_entities.append(Span(doc, start, end, label=match_id))
            entities = [e for e in entities if not (e.start < end and e.end > start)]
            seen_tokens.update(range(start, end))
    doc.ents = tuple(entities) + tuple(new_entities)
    return doc

#deduplicate repeated entities 
def dedupe(items):
    seen = set()
    for item in items:
        item = str(item).strip()
        if item not in seen:
            yield item
            seen.add(item)



##======================== End of NER Tagging ========================

def run_model(input_text):    
    if model == "BertSummarizer":
        output = original_text['BertSummarizer2s'].values
        st.write('Summary')

    elif model == "BertGPT2":
        output = original_text['BertGPT2'].values
        st.write('Summary')
        
      
    elif model == "t5seq2eq": 
        output = original_text['t5seq2eq'].values
        st.write('Summary')
         
    elif model == "t5":
        output = original_text['t5'].values
        st.write('Summary')
        
    elif model == "gensim": 
        output = original_text['gensim'].values
        st.write('Summary')
        
    elif model == "pysummarizer": 
        output = original_text['pysummarizer'].values
        st.write('Summary')


    
    st.success(output)

##========= on Past History Tab  =========  
if btnAdmission or st.session_state["button_clicked"]:
    st.session_state["daily_button_clicked"] =  False
    st.session_state["past_button_clicked"] =  False
    st.session_state["button_clicked"] =  True


    runtext =st.text_area(inputNote, str(original_text2)[1:-1], height=300)
    
if btnPastHistory or st.session_state["past_button_clicked"]:
    st.session_state["button_clicked"] =  False
    st.session_state["daily_button_clicked"] =  False
    st.session_state["past_button_clicked"] =  True

    with st.container(): 
        with col6: 

            st.markdown('**No. of admission past 6 months:**')
            st.markdown(countOfAdmission)
       
        with col7:
            #st.date_input('Select Admission Date') # To replace with a dropdown filter instead 
            #st.selectbox('Past Episodes',pastHistoryEp)
            pastHistory = st.selectbox('Select Past History Admission', pastHistoryEpDate, format_func=lambda x: 'Select an option' if x == '' else x)        
    
    historyAdmission =  df3.query(
                "Patient_ID  == @patient & CHARTDATE_HADM_ID == @pastHistory"
                )
    
   
    if historyAdmission.shape[0] == 0:
        runtext = "No past episodes" 
    else: 
        #runtext = historyAdmission['hospital_course_processed'].values[0]
        runtext = historyAdmission['hospital_course_processed'].values[0]
    
if btnDailyNarrative or st.session_state["daily_button_clicked"]:
    st.session_state["button_clicked"] =  False
    st.session_state["past_button_clicked"] =  False
    st.session_state["daily_button_clicked"] =  True


    


lem_clinical_note= lemmatize(runtext, nlp)
#creating a doc object using BC5CDR model
doc = nlp(lem_clinical_note)
options = get_entity_options()

#list of negative concepts from clinical note identified by negspacy
results0 = negation_handling(lem_clinical_note, neg_model)

matcher = match(nlp, results0,"NEG_ENTITY")

#doc0: new doc object with added "NEG_ENTITY label"
doc0 = overwrite_ent_lbl(matcher,doc)

#visualizing identified Named Entities in clinical input text 
ent_html = displacy.render(doc0, style='ent', options=options)


    
col1, col2 = st.columns([1,1])

#to not show summary and references text for Past History and Daily Narrative
if btnAdmission or st.session_state["button_clicked"]:
        st.session_state["daily_button_clicked"] =  False
        st.session_state["past_button_clicked"] =  False
        st.session_state["button_clicked"] =  True
        with col1:
            st.button('Summarize')
            run_model(runtext)
            #sentences=runtext.split('.')
            st.text_area('Reference text', str(reference_text), height=150)
        with col2:
            st.button('NER')
            # ===== Adding the Disease/Chemical into a list =====    
            problem_entities = list(dedupe([t for t in doc0.ents if t.label_ == 'DISEASE']))
            medication_entities = list(dedupe([t for t in doc0.ents if t.label_ == 'CHEMICAL']))
            st.markdown('**CHIEF COMPLAINT:**')
            st.write(str(AdmissionChiefCom)[1:-1])
            st.markdown('**ADMISSION DIAGNOSIS:**')
            st.markdown(str(diagnosis)[1:-1].capitalize())
            st.markdown('**PROBLEM/ISSUE**')
            #st.markdown(problem_entities)
            st.markdown(f'<p style="background-color:PINK;color:#080808;font-size:16px;">{str(problem_entities)[1:-1]}</p>', unsafe_allow_html=True)
            #genEntities(trans_df, 'DISEASE')
            st.markdown('**MEDICATION**')
            st.markdown(f'<p style="background-color:orange;color:#080808;font-size:16px;">{str(medication_entities)[1:-1]}</p>', unsafe_allow_html=True)
            #genEntities(trans_df, 'CHEMICAL')
            #st.table(trans_df)   
            st.markdown('**NER**')
            with st.expander("See NER Details"):
                st.markdown(ent_html, unsafe_allow_html=True)
                
                
elif btnDailyNarrative or st.session_state["daily_button_clicked"]:
    st.session_state["daily_button_clicked"] =  True
    st.session_state["past_button_clicked"] =  False
    st.session_state["button_clicked"] =  False
   
    with st.container():
        st.markdown('Daily Progress Note (24 hour event only):') 
        st.markdown(str(dailyNote)[1:-1])
 
    
    with st.container():
        styler = dailyNoteChange.style.hide_index()
        st.write(styler.to_html(), unsafe_allow_html=True)
    st.markdown(f'<p style="color:#828080;font-size:12px;">*Current prototype displays only a single section within the daily progress note, could also potentially include all sections within each progress note and allow user to select the section changes the user wants to look at</p>', unsafe_allow_html=True)

#else:
elif btnPastHistory or st.session_state["past_button_clicked"]:
    st.session_state["past_button_clicked"] =  True
    st.session_state["button_clicked"] =  False
    st.session_state["daily_button_clicked"] =  False
# ===== Adding the Disease/Chemical into a list =====    
    problem_entities = list(dedupe([t for t in doc0.ents if t.label_ == 'DISEASE']))
    medication_entities = list(dedupe([t for t in doc0.ents if t.label_ == 'CHEMICAL']))
    if historyAdmission.shape[0] == 0:
        st.markdown('Admission Date: NA')
        st.markdown('Date of Discharge: NA') 
        st.markdown('Days from current admission: NA')
    else: 
        st.markdown('Admission Date: ' + historyAdmission['ADMITTIME'].values[0])
        st.markdown('Date of Discharge: ' + historyAdmission['DISCHTIME'].values[0])
        st.markdown('Days from current admission: ' + str(historyAdmission['days_from_index'].values[0]) +' days')

    #st.markdown('Summary: ')
    st.markdown(f'<p style="color:#080808;font-size:16px;"><b>Summary: </b></p>', unsafe_allow_html=True)
    
              
    if model == "BertSummarizer":
        if historyAdmission.shape[0] == 0:
            st.markdown('NA')
        else: 
            st.markdown(str(historyAdmission['BertSummarizer'].values[0]))
    elif model == "t5seq2eq":
        if historyAdmission.shape[0] == 0:
            st.markdown('NA')
        else: 
            st.markdown(str(historyAdmission['t5seq2eq'].values[0]))     
    st.markdown(f'<p style="color:#080808;font-size:16px;"><b>Diagnosis: </b></p>', unsafe_allow_html=True)
    
    if historyAdmission.shape[0] == 0:
        st.markdown('NA')
    else:
        st.markdown(str(historyAdmission['Diagnosis_Description'].values[0]))
        st.markdown('**PROBLEM/ISSUE**')
        st.markdown(f'<p style="background-color:PINK;color:#080808;font-size:16px;">{str(problem_entities)[1:-1]}</p>', unsafe_allow_html=True)
        st.markdown('**MEDICATION**')
        st.markdown(f'<p style="background-color:orange;color:#080808;font-size:16px;">{str(medication_entities)[1:-1]}</p>', unsafe_allow_html=True)        
        st.markdown('Discharge Disposition: ' + str(historyAdmission['DISCHARGE_LOCATION'].values[0]))
        with st.expander('Full Discharge Summary'):
                #st.write("line 1  \n  line 2  \n line 3")
                fulldischargesummary = historyAdmission['TEXT'].values[0]
                st.write(fulldischargesummary)