from __future__ import absolute_import
from __future__ import print_function

import csv
import numpy as np
import os
import gzip
import pandas as pd
from tqdm import tqdm
from data_utils.preprocessing.mimic3benchmark.util import dataframe_from_csv

def read_patients_table(mimic3_path):
	pats = dataframe_from_csv(os.path.join(mimic3_path, 'PATIENTS.csv'))
	pats = pats[['SUBJECT_ID', 'GENDER', 'DOB', 'DOD']]
	pats.DOB = pd.to_datetime(pats.DOB)
	pats.DOD = pd.to_datetime(pats.DOD)
	return pats


def read_admissions_table(mimic3_path):
	admits = dataframe_from_csv(os.path.join(mimic3_path, 'ADMISSIONS.csv'))
	admits = admits[['SUBJECT_ID', 'HADM_ID', 'ADMITTIME', 'DISCHTIME', 'DEATHTIME', 'ETHNICITY', 'DIAGNOSIS']]
	admits.ADMITTIME = pd.to_datetime(admits.ADMITTIME)
	admits.DISCHTIME = pd.to_datetime(admits.DISCHTIME)
	admits.DEATHTIME = pd.to_datetime(admits.DEATHTIME)
	return admits


def read_icustays_table(mimic3_path):
	stays = dataframe_from_csv(os.path.join(mimic3_path, 'ICUSTAYS.csv'))
	stays.INTIME = pd.to_datetime(stays.INTIME)
	stays.OUTTIME = pd.to_datetime(stays.OUTTIME)
	return stays


def read_icd_diagnoses_table(mimic3_path):
	codes = dataframe_from_csv(os.path.join(mimic3_path, 'D_ICD_DIAGNOSES.csv'))
	codes = codes[['ICD9_CODE', 'SHORT_TITLE', 'LONG_TITLE']]
	diagnoses = dataframe_from_csv(os.path.join(mimic3_path, 'DIAGNOSES_ICD.csv'))
	diagnoses = diagnoses.merge(codes, how='inner', left_on='ICD9_CODE', right_on='ICD9_CODE')
	diagnoses[['SUBJECT_ID', 'HADM_ID', 'SEQ_NUM']] = diagnoses[['SUBJECT_ID', 'HADM_ID', 'SEQ_NUM']].astype(int)
	return diagnoses


def read_events_table_by_row(mimic3_path, table):
	nb_rows = {'chartevents': 330712484, 'labevents': 27854056, 'outputevents': 4349219}
	with gzip.open(os.path.join(mimic3_path, table.upper() + '.csv.gz'), 'rt') as f:
		#reader = csv.DictReader(open(os.path.join(mimic3_path, table.upper() + '.csv'), 'r'))
		reader = csv.DictReader(f)
		for i, row in enumerate(reader):
			if 'ICUSTAY_ID' not in row:
				row['ICUSTAY_ID'] = ''
			yield row, i, nb_rows[table.lower()]


def count_icd_codes(diagnoses, output_path=None):
	codes = diagnoses[['ICD9_CODE', 'SHORT_TITLE', 'LONG_TITLE']].drop_duplicates().set_index('ICD9_CODE')
	codes['COUNT'] = diagnoses.groupby('ICD9_CODE')['ICUSTAY_ID'].count()
	codes.COUNT = codes.COUNT.fillna(0).astype(int)
	codes = codes[codes.COUNT > 0]
	if output_path:
		codes.to_csv(output_path, index_label='ICD9_CODE')
	return codes.sort_values('COUNT', ascending=False).reset_index()


def remove_icustays_with_transfers(stays):
	stays = stays[(stays.FIRST_WARDID == stays.LAST_WARDID) & (stays.FIRST_CAREUNIT == stays.LAST_CAREUNIT)]
	return stays[['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'LAST_CAREUNIT', 'DBSOURCE', 'INTIME', 'OUTTIME', 'LOS']]


def merge_on_subject(table1, table2):
	return table1.merge(table2, how='inner', left_on=['SUBJECT_ID'], right_on=['SUBJECT_ID'])


def merge_on_subject_admission(table1, table2):
	return table1.merge(table2, how='inner', left_on=['SUBJECT_ID', 'HADM_ID'], right_on=['SUBJECT_ID', 'HADM_ID'])


def add_age_to_icustays(stays):
	#NOTE: https://mimic.mit.edu/docs/iii/about/time/ - <1900 means age >= 89 on the first admit
	diff = stays.INTIME.map(lambda x: pd.to_datetime(x).date()) - stays.DOB.map(lambda x: pd.to_datetime(x).date())
	stays['AGE'] = diff.map(lambda x: x.days/365)
	stays.loc[stays.AGE < 0, 'AGE'] = 90
	stays.loc[stays.AGE > 200, 'AGE'] = 91
	return stays


def add_inhospital_mortality_to_icustays(stays):
	mortality = stays.DOD.notnull() & ((stays.ADMITTIME <= stays.DOD) & (stays.DISCHTIME >= stays.DOD))
	mortality = mortality | (stays.DEATHTIME.notnull() & ((stays.ADMITTIME <= stays.DEATHTIME) & (stays.DISCHTIME >= stays.DEATHTIME)))
	stays['MORTALITY'] = mortality.astype(int)
	stays['MORTALITY_INHOSPITAL'] = stays['MORTALITY']
	return stays


def add_inunit_mortality_to_icustays(stays):
	mortality = stays.DOD.notnull() & ((stays.INTIME <= stays.DOD) & (stays.OUTTIME >= stays.DOD))
	mortality = mortality | (stays.DEATHTIME.notnull() & ((stays.INTIME <= stays.DEATHTIME) & (stays.OUTTIME >= stays.DEATHTIME)))
	stays['MORTALITY_INUNIT'] = mortality.astype(int)
	return stays


def filter_admissions_on_nb_icustays(stays, min_nb_stays=1, max_nb_stays=1):
	to_keep = stays.groupby('HADM_ID').count()[['ICUSTAY_ID']].reset_index()
	to_keep = to_keep[(to_keep.ICUSTAY_ID >= min_nb_stays) & (to_keep.ICUSTAY_ID <= max_nb_stays)][['HADM_ID']]
	stays = stays.merge(to_keep, how='inner', left_on='HADM_ID', right_on='HADM_ID')
	return stays


def filter_icustays_on_age(stays, min_age=18, max_age=np.inf):
	stays = stays[(stays.AGE >= min_age) & (stays.AGE <= max_age)]
	return stays


def filter_diagnoses_on_stays(diagnoses, stays):
	return diagnoses.merge(stays[['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID']].drop_duplicates(), how='inner',
						   left_on=['SUBJECT_ID', 'HADM_ID'], right_on=['SUBJECT_ID', 'HADM_ID'])


def break_up_stays_by_subject(stays, output_path, subjects=None):
	subjects = stays.SUBJECT_ID.unique() if subjects is None else subjects
	nb_subjects = subjects.shape[0]
	for subject_id in tqdm(subjects, total=nb_subjects, desc='Breaking up stays by subjects'):
		dn = os.path.join(output_path, str(subject_id))
		try:
			os.makedirs(dn)
		except:
			pass

		stays[stays.SUBJECT_ID == subject_id].sort_values(by='INTIME').to_csv(os.path.join(dn, 'stays.csv'),
																			  index=False)


def break_up_diagnoses_by_subject(diagnoses, output_path, subjects=None):
	subjects = diagnoses.SUBJECT_ID.unique() if subjects is None else subjects
	nb_subjects = subjects.shape[0]
	for subject_id in tqdm(subjects, total=nb_subjects, desc='Breaking up diagnoses by subjects'):
		dn = os.path.join(output_path, str(subject_id))
		try:
			os.makedirs(dn)
		except:
			pass

		diagnoses[diagnoses.SUBJECT_ID == subject_id].sort_values(by=['ICUSTAY_ID', 'SEQ_NUM'])\
													 .to_csv(os.path.join(dn, 'diagnoses.csv'), index=False)


def read_events_table_and_break_up_by_subject(mimic3_path, table, output_path,
											  items_to_keep=None, subjects_to_keep=None):
	obs_header = ['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'CHARTTIME', 'ITEMID', 'VALUE', 'VALUEUOM']
	if items_to_keep is not None:
		items_to_keep = set([str(s) for s in items_to_keep])
	if subjects_to_keep is not None:
		subjects_to_keep = set([str(s) for s in subjects_to_keep])

	class DataStats(object):
		def __init__(self):
			self.curr_subject_id = ''
			self.curr_obs = []

	data_stats = DataStats()

	def write_current_observations():
		dn = os.path.join(output_path, str(data_stats.curr_subject_id))
		try:
			os.makedirs(dn)
		except:
			pass
		fn = os.path.join(dn, 'events.csv')
		if not os.path.exists(fn) or not os.path.isfile(fn):
			f = open(fn, 'w')
			f.write(','.join(obs_header) + '\n')
			f.close()
		w = csv.DictWriter(open(fn, 'a'), fieldnames=obs_header, quoting=csv.QUOTE_MINIMAL)
		w.writerows(data_stats.curr_obs)
		data_stats.curr_obs = []

	nb_rows_dict = {'chartevents': 330712484, 'labevents': 27854056, 'outputevents': 4349219}
	nb_rows = nb_rows_dict[table.lower()]

	for row, row_no, _ in tqdm(read_events_table_by_row(mimic3_path, table), total=nb_rows,
														desc='Processing {} table'.format(table)):

		if (subjects_to_keep is not None) and (row['SUBJECT_ID'] not in subjects_to_keep):
			continue
		if (items_to_keep is not None) and (row['ITEMID'] not in items_to_keep):
			continue

		row_out = {'SUBJECT_ID': row['SUBJECT_ID'],
				   'HADM_ID': row['HADM_ID'],
				   'ICUSTAY_ID': '' if 'ICUSTAY_ID' not in row else row['ICUSTAY_ID'],
				   'CHARTTIME': row['CHARTTIME'],
				   'ITEMID': row['ITEMID'],
				   'VALUE': row['VALUE'],
				   'VALUEUOM': row['VALUEUOM']}
		if data_stats.curr_subject_id != '' and data_stats.curr_subject_id != row['SUBJECT_ID']:
			write_current_observations()
		data_stats.curr_obs.append(row_out)
		data_stats.curr_subject_id = row['SUBJECT_ID']

	if data_stats.curr_subject_id != '':
		write_current_observations()