{
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.6.5 64-bit ('base': conda)",
   "metadata": {
    "interpreter": {
     "hash": "5f55f32a86243f5615d447fc6161221c78f375b55f8f0019de541d78503bbfd0"
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../../')\n",
    "# put mimic3 related files in MIMIC_3_DIR\n",
    "from constant import MIMIC_3_DIR\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from icd_50 import codes_50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "2083180it [01:10, 29445.58it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import csv\n",
    "\n",
    "note_dict = {}\n",
    "others_dict = {}\n",
    "\n",
    "notes_file = '%s/NOTEEVENTS.csv' % (MIMIC_3_DIR)\n",
    "with open(notes_file, 'r') as csvfile:\n",
    "    notereader = csv.reader(csvfile)\n",
    "    next(notereader)\n",
    "    i = 0\n",
    "    for line in tqdm(notereader):\n",
    "        subject_id = int(line[1])\n",
    "        hadm_id = str(line[2])\n",
    "        category = str(line[6])\n",
    "        note = line[10]\n",
    "        if (subject_id, hadm_id) not in note_dict:\n",
    "            note_dict[(subject_id, hadm_id)] = []\n",
    "            others_dict[(subject_id, hadm_id)] = []\n",
    "        if category == \"Discharge summary\":\n",
    "            note_dict[(subject_id, hadm_id)].append(note)\n",
    "        else:\n",
    "            others_dict[(subject_id, hadm_id)].append(category + \": \" + note)\n",
    "        i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "52726\n"
     ]
    }
   ],
   "source": [
    "df_codes = pd.read_csv('%s/ALL_CODES_filtered.csv' % MIMIC_3_DIR, index_col=None)\n",
    "print(len(set(df_codes['HADM_ID'])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_note_to_paragraph(note):\n",
    "    result = {}\n",
    "    now_paragraph = 'other'\n",
    "    for line in note.split('\\n'):\n",
    "        if line.find(\":\") >= 0:\n",
    "            now_paragraph = line[0:line.find(\":\")]\n",
    "            text = line[line.find(\":\") + 1:]\n",
    "            if not now_paragraph in result:\n",
    "                result[now_paragraph] = []\n",
    "            result[now_paragraph].append(text.strip())\n",
    "        elif not line.strip():\n",
    "            now_paragraph = 'other'\n",
    "        else:\n",
    "            if not now_paragraph in result:\n",
    "                result[now_paragraph] = []\n",
    "            result[now_paragraph].append(line.strip())\n",
    "    for paragraph in result:\n",
    "        result[paragraph] = \" \".join(result[paragraph])\n",
    "    return result\n",
    "\n",
    "# split_note_to_paragraph(text_list[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# p_count = {}\n",
    "# for idx, text in tqdm(enumerate(text_list)):\n",
    "#     result = split_note_to_paragraph(text)\n",
    "#     for key in result.keys():\n",
    "#         if not key in p_count:\n",
    "#             p_count[key] = 0\n",
    "#         p_count[key] += 1\n",
    "# print(len(p_count))\n",
    "# p_count_sort = {k: v for k, v in sorted(p_count.items(), key=lambda item: -item[1])}\n",
    "# list(p_count_sort.items())[0:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 848692/848692 [00:38<00:00, 21865.04it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "\n",
    "icd_dict = {}\n",
    "for i in trange(df_codes.shape[0]):\n",
    "    subject_id = df_codes['SUBJECT_ID'][i]\n",
    "    hadm_id = str(df_codes['HADM_ID'][i])\n",
    "    code = df_codes['ICD9_CODE'][i]\n",
    "    if (subject_id, hadm_id) not in icd_dict:\n",
    "        icd_dict[(subject_id, hadm_id)] = []\n",
    "    icd_dict[(subject_id, hadm_id)].append(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in trange(df_notes.shape[0]):\n",
    "#     subject_id = df_notes['SUBJECT_ID'][i]\n",
    "#     hadm_id = df_notes['HADM_ID'][i]\n",
    "#     text = df_notes['TEXT'][i]\n",
    "#     if (subject_id, hadm_id) in note_dict:\n",
    "#         note_dict[(subject_id, hadm_id)].append(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|██████████| 52726/52726 [00:06<00:00, 8739.29it/s]\n",
      "47723 1631 3372\n",
      "100%|██████████| 52726/52726 [00:02<00:00, 23364.72it/s]\n",
      "8066 1573 1729\n"
     ]
    }
   ],
   "source": [
    "for base_name in ['mimic3', 'mimic3-50']:\n",
    "    train_name = '%s_train.json' % (base_name)\n",
    "    dev_name = '%s_dev.json' % (base_name)\n",
    "    test_name = '%s_test.json' % (base_name)\n",
    "\n",
    "    hadm_ids = {}\n",
    "\n",
    "    #read in train, dev, test splits\n",
    "    for splt in ['train', 'dev', 'test']:\n",
    "        hadm_ids[splt] = set()\n",
    "        if base_name == \"mimic3\":\n",
    "            base = \"full\"\n",
    "        if base_name == \"mimic3-50\":\n",
    "            base = \"50\"\n",
    "        with open('%s/%s_%s_hadm_ids.csv' % (MIMIC_3_DIR, splt, base), 'r') as f:\n",
    "            for line in f:\n",
    "                hadm_ids[splt].add(line.rstrip())\n",
    "\n",
    "    train_list = []\n",
    "    dev_list = []\n",
    "    test_list = []\n",
    "\n",
    "    for key in tqdm(icd_dict.keys()):\n",
    "        subject_id = str(int(key[0]))\n",
    "        hadm_id = str(int(key[1]))\n",
    "        icd = \";\".join([str(c) for c in icd_dict[key]])\n",
    "        if base_name == \"mimic-50\":\n",
    "            filtered_codes = set(icd_dict[key]).intersection(set(codes_50))\n",
    "            if len(filtered_codes) > 0:\n",
    "                icd = \";\".join([str(c) for c in filtered_codes])\n",
    "            else:\n",
    "                continue\n",
    "\n",
    "        text = \"\\t\".join([str(c) for c in note_dict[key]])\n",
    "        additional_text = \"\\t\".join([str(c) for c in others_dict[key]])\n",
    "        row = {'subject_id':subject_id,\n",
    "               'hadm_id':hadm_id,\n",
    "               'LABELS':icd,\n",
    "               'TEXT':text,\n",
    "               'Addition':additional_text}\n",
    "\n",
    "        if hadm_id in hadm_ids['train']:\n",
    "            train_list.append(row)\n",
    "        elif hadm_id in hadm_ids['dev']:\n",
    "            dev_list.append(row)\n",
    "        elif hadm_id in hadm_ids['test']:\n",
    "            test_list.append(row)\n",
    "        else:\n",
    "            #print(key)\n",
    "            pass\n",
    "\n",
    "    print(len(train_list), len(dev_list), len(test_list))\n",
    "    import json\n",
    "    with open(train_name, \"w\") as f:\n",
    "        json.dump(train_list, f, indent=4)\n",
    "    with open(dev_name, \"w\") as f:\n",
    "        json.dump(dev_list, f, indent=4)\n",
    "    with open(test_name, \"w\") as f:\n",
    "        json.dump(test_list, f, indent=4)"
   ]
  }
 ]
}