--- a
+++ b/tests/tests.py
@@ -0,0 +1,472 @@
+import unittest
+import random
+import sys, os
+import re
+import nltk
+import json
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'allennlp')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'summarization', 'pubmed_summarization')))
+# print(sys.path)
+from ehrkit import ehrkit
+from getpass import getpass
+
+try: 
+    from config import USERNAME, PASSWORD
+except:
+    print("Please put your username and password in config.py")
+    USERNAME = input('DB_username?')
+    PASSWORD = getpass('DB_password?')
+
+
+DOC_ID = 1354526 # Temporary!!!
+
+
+# Number of documents in NOTEEVENTS.
+NUM_DOCS = 2083180
+
+# Number of patients in PATIENTS.
+NUM_PATIENTS = 46520
+
+# Number of diagnoses in DIAGNOSES_ICD.
+NUM_DIAGS = 823933
+
+
+def select_ehr(ehrdb, requires_long=False, recursing=False):
+    if recursing:
+        doc_id = ''
+    else:
+        #doc_id = input("MIMIC Document ID [press Enter for random]: ")
+        
+        doc_id = ''
+    
+    if doc_id == '':
+        # Picks random document
+        ehrdb.cur.execute("SELECT ROW_ID FROM mimic.NOTEEVENTS ORDER BY RAND() LIMIT 1")
+        doc_id = ehrdb.cur.fetchall()[0][0]
+        text = ehrdb.get_document(int(doc_id))
+        if len(text.split()) > 200 or not requires_long:
+            return doc_id, text
+        else:
+            return select_ehr(ehrdb, requires_long, True)
+    else:
+        # Get inputted document
+        try:
+            text = ehrdb.get_document(int(doc_id))
+            return doc_id, text
+        except:
+            message = 'Error: There is no document with ID \'' + doc_id + '\' in mimic.NOTEEVENTS'
+            sys.exit(message)
+
+
+def get_nb_dir(ending, SUMM_DIR):
+    # Gets path of Naive Bayes model trained on most examples
+    dir_nums = []
+    for dir in os.listdir(SUMM_DIR):
+        if os.path.isdir(os.path.join(SUMM_DIR, dir)) and dir.endswith('_exs_' + ending):
+            if os.path.exists(os.path.join(SUMM_DIR, dir, 'nb')):  
+                try:
+                    dir_nums.append(int(dir.split('_')[0]))
+                except:
+                    continue
+    if len(dir_nums) > 0:
+        best_dir_name = str(max(dir_nums)) + '_exs_' + ending
+        return best_dir_name
+    else:
+        return None
+
+def show_summary(doc_id, text, summary, model_name):
+    #x = input('Show full EHR (DOC ID %s)? [DEFAULT=Yes]' % doc_id)
+    x = ''
+    if x.lower() in ['y', 'yes', '']:
+        print('\n\n' + '-'*30 + 'Full EHR' + '-'*30)
+        print(text + '\n')
+        print('-'*80 + '\n\n')
+
+    print('-'*30 + 'Predicted Summary ' + model_name + '-'*30)
+    print(summary)
+    print('-'*80 + '\n\n')
+
+
+class tests(unittest.TestCase):
+    def setUp(self):
+        self.ehrdb = ehrkit.start_session(USERNAME, PASSWORD)
+        self.ehrdb.get_patients(3)
+
+
+''' Runs tests 1.1-1.4 '''
+class t1(tests):
+    def test1_1_count_patients(self):
+        kit_count = self.ehrdb.count_patients()
+        print("Patient count: ", kit_count)
+
+        self.ehrdb.cur.execute("SELECT COUNT(*) FROM mimic.PATIENTS")
+        raw = self.ehrdb.cur.fetchall()
+        test_count = int(raw[0][0])
+
+        self.assertEqual(test_count, kit_count)
+
+    # def test1_2_count_docs(self):
+    #     # Fails! count_docs returns 1573339, but mimic.NOTEEVENTS has 2083180 documents. 
+    #     # TO DO: Fix whatever is wrong here
+    #     kit_count = self.ehrdb.count_docs(['NOTEEVENTS'])
+    #     print("Document count: ", kit_count)
+
+    #     self.ehrdb.cur.execute("SELECT COUNT(*) FROM mimic.NOTEEVENTS")
+    #     raw = self.ehrdb.cur.fetchall()
+    #     test_count = int(raw[0][0])
+
+    #     self.assertEqual(test_count, kit_count)
+
+    def test1_3_note_info(self):
+        self.ehrdb.get_note_events()
+        print('output format: SUBJECT_ID, ROW_ID, NoteEvent length')
+        lens = [(patient.id, note[0], len(note[1])) for patient in self.ehrdb.patients.values() for note in patient.note_events]
+        print(lens)
+
+        # placeholder, this output cannot be checked easily
+        self.assertEqual(1, 1)
+
+    def test1_4_longest_note(self):
+        # Gets longest note among the patient notes queued by get_note_events()
+        self.ehrdb.get_note_events()
+        pid, rowid, doclen = self.ehrdb.longest_NE()
+        print('patient id is:', pid, 'NoteEvent id is:', rowid, 'length: ', doclen)
+
+        # placeholder, this output cannot be checked easily
+        self.assertEqual(1, 1)
+
+
+class t2(tests):
+    def test2_1_print_note(self):
+        ### There are 2083180 patient records in NOTEEVENTS. ###
+        record_id = random.randint(1, NUM_DOCS + 1)
+        kit_rec = self.ehrdb.get_document(record_id)
+        print("Document with ID %d\n: " % record_id, kit_rec)
+
+        self.ehrdb.cur.execute("select TEXT from mimic.NOTEEVENTS where ROW_ID = %d" % record_id)
+        test_rec = self.ehrdb.cur.fetchall()[0][0]
+
+        self.assertEqual(kit_rec, test_rec)
+
+    def test2_2_patient_info(self):
+        ### There are records from 46520 unique patients in MIMIC. ###
+        patient_id = random.randint(1, NUM_PATIENTS + 1)
+        kit_ids = self.ehrdb.get_all_patient_document_ids(patient_id)
+        print('Document IDs related to Patient %d: ' % patient_id, kit_ids)
+        print("Number of docs related to Patient %d: " % patient_id, len(kit_ids))
+
+        self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where SUBJECT_ID = %d" % patient_id)
+        raw = self.ehrdb.cur.fetchall()
+        test_ids = ehrkit.flatten(raw)
+
+        self.assertEqual(kit_ids, test_ids)
+
+    #@unittest.skipIf("t2.test2_3" not in sys.argv, "Test 2_3 must be run explicitly due to runtime.")
+    def test2_3_doc_ids(self):
+        kit_ids = self.ehrdb.list_all_document_ids()
+
+        self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS")
+        raw = self.ehrdb.cur.fetchall()
+        test_ids = ehrkit.flatten(raw)
+        print('test_ids')
+        print(test_ids[:30])
+        print('...')
+
+        self.assertEqual(kit_ids, test_ids)
+
+    def test2_4_patient_ids(self):
+        kit_ids = self.ehrdb.list_all_patient_ids()
+
+        self.ehrdb.cur.execute("select SUBJECT_ID from mimic.PATIENTS")
+        raw = self.ehrdb.cur.fetchall()
+        test_ids = ehrkit.flatten(raw)
+        print(test_ids[:30])
+        print('...')
+
+        self.assertEqual(kit_ids, test_ids)
+
+    #@unittest.skipIf("t2.test2_5" not in sys.argv, "Test 2_5 must be run explicitly due to runtime.")
+    def test2_5_docs_on_date(self):
+        ### Select random date from a date in the database. 
+        ### Dates are shifted to future but preserve time, weekday, and seasonality.
+        random_id = random.randint(1, NUM_DOCS + 1)
+        self.ehrdb.cur.execute("select CHARTDATE from mimic.NOTEEVENTS where ROW_ID = %d" % random_id)
+        date = self.ehrdb.cur.fetchall()[0][0]
+
+        kit_ids = self.ehrdb.get_documents_d(date)
+
+        self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where CHARTDATE = \"%s\"" % date)
+        raw = self.ehrdb.cur.fetchall()
+        test_ids = ehrkit.flatten(raw)
+        print(f"Selected date: {date}")
+        print(f"Test ids {test_ids[:30]} ...")
+
+        self.assertEqual(kit_ids, test_ids)
+
+
+class t3(tests):
+    def test3_1_extract_abbreviations(self):
+        # Defines abbreviation as a string of capitalized letters
+        random_id = random.randint(1, NUM_DOCS + 1)
+        print("Collecting abbreviations for document %d..." % random_id)
+        kit_abbs = self.ehrdb.get_abbreviations(random_id)
+
+        sents = self.ehrdb.get_document_sents(random_id)
+        test_abbs = set()
+        for sent in sents:
+            for word in ehrkit.word_tokenize(sent):
+                print(word)
+                pattern = r'[A-Z]{2}'  # Only selects words in ALL CAPS
+                if re.match(pattern, word):
+                    test_abbs.add(word)
+
+        print(kit_abbs)
+
+        self.assertEqual(kit_abbs, list(test_abbs))
+
+    #@unittest.skipIf("t3.test3_2" not in sys.argv, "Test 3_2 must be run explicitly due to runtime.")
+    def test3_2_docs_with_query(self):
+        query = "meningitis"
+        print('Printing a list of all document ids including query like ', query)
+        kit_ids = self.ehrdb.get_documents_q(query)
+        print(kit_ids[:30])  # Extremely long list of DOC_IDs
+        print("...")
+
+        query = "%"+query+"%"
+        self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where TEXT like \'%s\'" % query)
+        raw = self.ehrdb.cur.fetchall()
+        test_ids = ehrkit.flatten(raw)
+
+        self.assertEqual(kit_ids, test_ids)
+
+    #@unittest.skipIf("t3.test3_3" not in sys.argv, "Test 3_3 must be run explicitly due to runtime. Also, this is essentially a duplicate of task 3.2.")
+    def test3_3_query_docs(self):
+        ### Task 3.3 is the same as task 3.2 with a different query. ###
+        query = "Service: SURGERY"
+        print('Printing a list of all document ids including query like ', query)
+        kit_ids = self.ehrdb.get_documents_q(query)
+        print(kit_ids[:30])  # Extremely long list of DOC_IDs
+        print("...")
+
+        query = "%"+query+"%"
+        self.ehrdb.cur.execute("select ROW_ID from mimic.NOTEEVENTS where TEXT like \'%s\'" % query)
+        raw = self.ehrdb.cur.fetchall()
+        test_ids = ehrkit.flatten(raw)
+
+        self.assertEqual(kit_ids, test_ids)
+
+    def test3_4_doc_sentences(self):
+        doc_id = random.randint(1, NUM_DOCS + 1)
+        print('Kit function printing a numbered list of all sentences in document %d' % doc_id)
+        # MIMIC EHRs are very messy and sentence tokenizaton often doesn't work
+        kit_doc = self.ehrdb.get_document_sents(doc_id)
+        ehrkit.numbered_print(kit_doc)
+
+        self.ehrdb.cur.execute("select TEXT from mimic.NOTEEVENTS where ROW_ID = %d " % doc_id)
+        raw = self.ehrdb.cur.fetchall()
+        test_doc = ehrkit.sent_tokenize(raw[0][0])
+        print(test_doc)
+
+        self.assertEqual(kit_doc, test_doc)
+
+    #@unittest.skipIf("t3.test3_7" not in sys.argv, "Test 3_7 must be run explicitly due to runtime.")
+    def test3_7_medications(self):
+        kit_meds = self.ehrdb.count_all_prescriptions()
+
+        test_meds = {}
+        self.ehrdb.cur.execute("select DRUG from mimic.PRESCRIPTIONS")
+        raw = self.ehrdb.cur.fetchall()
+        meds_list = ehrkit.flatten(raw)
+        for med in meds_list:
+            if med in test_meds:
+                test_meds[med] += 1
+            else:
+                test_meds[med] = 1
+
+        print(meds_list[:30])
+        print("...")
+
+        self.assertEqual(kit_meds, test_meds)
+
+
+class t4(tests):
+    @unittest.skip("Task 4.1 is not ready to be tested yet.")
+    def test4_1(self):
+        d = self.ehrdb.get_documents_icd9()
+        print(d)
+        self.assertIsNotNone(d['code'])
+
+    @unittest.skip("Task 4.4 is not ready to be tested yet.")
+    def test4_4(self):
+        pass
+
+
+class t5(tests):
+    @unittest.skipIf("t5.test5_1" not in sys.argv, "Test 5_1 must be run explicitly due to runtime.")
+    def test5_1_extract_phrases(self):
+        doc_id = random.randint(1, NUM_DOCS + 1)
+        kit_phrases = self.ehrdb.extract_phrases(doc_id)
+
+        print("Testing task 5.1\n Check phrases manually: ", kit_phrases)
+
+        self.assertIsNotNone(kit_phrases)
+
+    def test5_4_count_gender(self):
+        gender = random.choice(['M', 'F'])
+        kit_count = self.ehrdb.count_gender(gender)
+
+        self.ehrdb.cur.execute('SELECT COUNT(*) FROM mimic.PATIENTS WHERE GENDER = \'%s\'' % gender)
+        raw = self.ehrdb.cur.fetchall()
+        test_count = raw[0][0]
+        print('Gender:', gender, '\tCount:', str(test_count))
+
+        self.assertEqual(kit_count, test_count)
+
+
+class t6(tests):
+    @unittest.skipIf("t6.test6_1_sentiment_classification" not in sys.argv, "Test 6_1 must be run explicitly due to verbosity.")
+    def test6_1_sentiment_classification(self):
+        import loader
+
+        doc_id, text = select_ehr(self.ehrdb)
+
+        x = input('GloVe or RoBERTa predictor [g=GloVe, r=RoBERTa]? ')
+        if x == 'g':
+            glove_predictor = loader.load_glove()
+            probs = glove_predictor.predict(text)['probs']
+        elif x == 'r':
+            roberta_predictor = loader.load_roberta()
+            try:
+                probs = roberta_predictor.predict(text)['probs']
+            except:
+                print('Document too long for RoBERTa model. Using GLoVe instead.')
+                glove_predictor = loader.load_glove()
+                probs = glove_predictor.predict(text)['probs']
+        else:
+            sys.exit('Error: Must input \'g\' or  \'r\'')
+
+        classification = 'Positive' if probs[0] >= 0.5 else 'Negative'
+        print("Document ID: ", doc_id, "\tPredicted Sentiment: ", classification)
+
+    @unittest.skipIf("t6.test6_2_ner" not in sys.argv, "Test 6_2 must be run explicitly due to verbosity.")
+    def test6_2_ner(self):
+        import loader
+
+        doc_id, text = select_ehr(self.ehrdb)
+
+        if os.path.exists("../allennlp/elmo-ner/whole_model.pt"):
+            predictor = loader.load_ner()
+        else:
+            predictor = loader.download_ner()
+
+        text = self.ehrdb.get_document(int(doc_id))
+        pred = predictor.predict(text)
+        # pred = predictor.predict("John likes and Bill hates ice cream")
+        print_results = input("Prediction complete. Print results? (y/n): ")
+        if print_results == "y":
+            print("Document ID: ", doc_id, "  Results: ", pred['tags'])
+    
+    @unittest.skipIf("t6.test6_3_tokenize" not in sys.argv, "Test 6_3 must be run explicitly due to runtime.")
+    def test6_3_tokenize(self):
+        import torch
+        from transformers import BertTokenizer#, BertModel, BertForMaskedLM
+
+        doc_id, text = select_ehr(self.ehrdb)
+        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        bert_tokenized_text = tokenizer.tokenize(text)
+        print('\n' + '-'*20 + 'text' + '-'*20)
+        print(text)
+        print('\n' + '-'*20 + 'Tokenized text from Huggingface BERT Tokenizer' + '-'*20)    
+        print(bert_tokenized_text)
+
+
+        # library function
+        ehr_bert_tokenized_text = self.ehrdb.get_bert_tokenize(doc_id)
+        self.assertEqual(bert_tokenized_text, ehr_bert_tokenized_text)
+
+
+class t7(tests):
+    # Summarization algorithms
+    #@unittest.skipIf("t7.test7_1_naive_bayes" not in sys.argv, "Test 7_1 must be run explicitly due to verbosity.")
+    def test7_1_naive_bayes(self):
+        from pubmed_naive_bayes import classify_nb
+        from get_pubmed_nb_data import build_vecs
+        from sklearn.naive_bayes import GaussianNB
+
+        doc_id, text = select_ehr(self.ehrdb)
+        #body_type = input('Use Naive Bayes model trained from whole body sections or just their body introductions?\n\t'\
+        #                '[w=whole body, j=just intro, DEFAULT=just intro]: ')
+        
+        body_type = 'j'
+
+        if body_type == 'w':
+            ending = 'body'
+        elif body_type in ['j', '']:
+            ending = 'intro'
+        else:
+            sys.exit('Error: Must input \'w\' or \'j.\'')
+        SUMM_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'summarization', 'pubmed_summarization'))
+        best_dir_name = get_nb_dir(ending, SUMM_DIR)
+        if not best_dir_name:
+            message = 'No Naive Bayes models of this type have been fit. '\
+                        'Would you like to do so now?\n\t[DEFAULT=Yes] '
+            #response = input(message)
+            
+            response = 'y'
+
+            if response.lower() in ['y', 'yes', '']:
+                command = 'python ' + os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'summarization', 'pubmed_summarization', 'pubmed_naive_bayes.py'))
+                os.system(command)
+                best_dir_name = get_nb_dir(ending)
+            if response.lower() not in ['y', 'yes', ''] or not best_dir_name:
+                sys.exit('Exiting.')
+
+        # Fits model to data        
+        NB_DIR = os.path.join(SUMM_DIR, best_dir_name, 'nb')
+        with open(os.path.join(NB_DIR, 'feature_vecs.json'), 'r') as f:
+            data = json.load(f)
+        xtrain, ytrain = data['train_features'], data['train_outputs']
+        gnb = GaussianNB()
+        gnb.fit(xtrain, ytrain)
+
+        # Evaluates on model
+        tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
+        feature_vecs, _ = build_vecs(text, None, tokenizer)
+        PCT_SUM = 0.3
+        preds = classify_nb(feature_vecs, PCT_SUM, gnb)
+        sents = tokenizer.tokenize(text)
+        summary = ''
+        for i in range(len(preds)):
+            if preds[i] == 1:
+                summary += sents[i]
+
+        show_summary(doc_id, text, summary, 'Naive Bayes')
+        
+    #@unittest.skipIf("t7.test7_2_distilbart_summary" not in sys.argv, "Test 7_2 must be run explicitly due to runtime.")
+    def test7_2_distilbart_summary(self):
+        # Distilbart for summarization. Trained on CNN/ Daily Mail (~4x longer summaries than XSum)
+        doc_id, text = select_ehr(self.ehrdb, requires_long=True)
+        model_name = 'sshleifer/distilbart-cnn-12-6'
+        summary = self.ehrdb.summarize_huggingface(text, model_name)
+
+        show_summary(doc_id, text, summary, model_name)
+        print('Number of Words in Full EHR: %d' % len(text.split()))
+        print('Number of Words in %s Summary: %d' % (model_name, len(summary.split())))
+
+    #@unittest.skipIf("t7.test7_3_t5_summary" not in sys.argv, "Test 7_3 must be run explicitly due to runtime.")
+    def test7_3_t5_summary(self):
+        # T5 for summarization. Trained on CNN/ Daily Mail
+        doc_id, text = select_ehr(self.ehrdb, requires_long=True)
+        model_name = 't5-small'
+        summary = self.ehrdb.summarize_huggingface(text, model_name)
+
+        show_summary(doc_id, text, summary, model_name)
+        print('Number of Words in Full EHR: %d' % len(text.split()))
+        print('Number of Words in %s Summary: %d' % (model_name, len(summary.split())))
+
+def testing():
+    unittest.main()
+
+if __name__ == '__main__':
+    unittest.main()