import os

import csv

import shutil

import sqlite3

from datetime import datetime

from process_mimic_db.utils import *

from process_mimic_db.process_tables import *

# Specify the path to the downloaded MIMIC III data

data_dir = '/home/nas1_userA/junwoopark/datasets/mimiciii/1.4/'

# Path to the generated mimic.db. No need to update.

out_dir = 'mimic_db'

# Generate five tables and the database with all admissions

# if os.path.exists(out_dir):

#     shutil.rmtree(out_dir)

# os.mkdir(out_dir)

conn = sqlite3.connect(os.path.join(out_dir, 'mimic_all.db'))

build_demographic_table(data_dir, out_dir, conn)

build_diagnoses_table(data_dir, out_dir, conn)

build_procedures_table(data_dir, out_dir, conn)

build_prescriptions_table(data_dir, out_dir, conn)

build_lab_table(data_dir, out_dir, conn)

'''

1. We did not emumerate all possible questions about MIMIC III.

MIMICSQL data is generated based on the patient information 

related to 100 randomly selected admissions.

2. The following codes are used for sampling the admissions 

from the large database. 

3. The parameter 'random_state=0' in line 41 will provide you 

the same set of sampled admissions and the same database as we used.

'''

print('Begin sampling ...')

# DEMOGRAPHIC

print('Processing DEMOGRAPHIC')

conn = sqlite3.connect(os.path.join(out_dir, 'mimic.db'))

data_demo = pandas.read_csv(os.path.join(out_dir, "DEMOGRAPHIC.csv"))

data_demo_sample = data_demo.sample(100, random_state=0)

for k, v in data_demo_sample.dtypes.items():

    data_demo_sample[k] = data_demo_sample[k].apply(lambda x: x.lower() if type(x) == str else x)

data_demo_sample.to_sql('DEMOGRAPHIC', conn, if_exists='replace', index=False)

sampled_id = data_demo_sample['HADM_ID'].values

# DIAGNOSES

print('Processing DIAGNOSES')

data_input = pandas.read_csv(os.path.join(out_dir, "DIAGNOSES.csv"))

data_filter = []

cnt = 0

for itm in sampled_id:

    msg = 'HADM_ID=='+str(itm)

    data_filter.append(data_input.query(msg))

    cnt += 1

    show_progress(cnt, len(sampled_id))

data_out = pandas.concat(data_filter, ignore_index=True)

for k, v in data_out.dtypes.items():

    data_out[k] = data_out[k].apply(lambda x: x.lower() if type(x) == str else x)

data_out.to_sql('DIAGNOSES', conn, if_exists='replace', index=False)

# PROCEDURES

print('Processing PROCEDURES')

data_input = pandas.read_csv(os.path.join(out_dir, "PROCEDURES.csv"))

data_filter = []

cnt = 0

for itm in sampled_id:

    msg = 'HADM_ID=='+str(itm)

    data_filter.append(data_input.query(msg))

    cnt += 1

    show_progress(cnt, len(sampled_id))

data_out = pandas.concat(data_filter, ignore_index=True)

for k, v in data_out.dtypes.items():

    data_out[k] = data_out[k].apply(lambda x: x.lower() if type(x) == str else x)

data_out.to_sql('PROCEDURES', conn, if_exists='replace', index=False)

# PRESCRIPTIONS

print('Processing PRESCRIPTIONS')

data_input = pandas.read_csv(os.path.join(out_dir, "PRESCRIPTIONS.csv"))

data_filter = []

cnt = 0

for itm in sampled_id:

    msg = 'HADM_ID=='+str(itm)

    data_filter.append(data_input.query(msg))

    cnt += 1

    show_progress(cnt, len(sampled_id))

data_out = pandas.concat(data_filter, ignore_index=True)

for k, v in data_out.dtypes.items():

    data_out[k] = data_out[k].apply(lambda x: x.lower() if type(x) == str else x)

data_out.to_sql('PRESCRIPTIONS', conn, if_exists='replace', index=False)

# LAB

print('Processing LAB')

data_input = pandas.read_csv(os.path.join(out_dir, "LAB.csv"))

data_filter = []

cnt = 0

for itm in sampled_id:

    msg = 'HADM_ID=='+str(itm)

    data_filter.append(data_input.query(msg))

    cnt += 1

    show_progress(cnt, len(sampled_id))

data_out = pandas.concat(data_filter, ignore_index=True)

for k, v in data_out.dtypes.items():

    data_out[k] = data_out[k].apply(lambda x: x.lower() if type(x) == str else x)

data_out.to_sql('LAB', conn, if_exists='replace', index=False)

print('Done!')