# %%
# Import necessary packages
import numpy as np
import pandas as pd
import torch

# %%
# Read raw data
df_train: pd.DataFrame = pd.read_excel(
    "./datasets/tongji/raw_data/time_series_375_prerpocess_en.xlsx"
)

# %% [markdown]
# Steps:
#
# - fill `patient_id`
# - only reserve y-m-d for `RE_DATE` column
# - merge lab tests of the same (patient_id, date)
# - calculate and save features' statistics information (demographic and lab test data are calculated separately)
# - normalize data
# - feature selection
# - fill missing data (our filling strategy will be described below)
# - combine above data to time series data (one patient one record)
# - export to python pickle file

# %%
# fill `patient_id` rows
df_train["PATIENT_ID"].fillna(method="ffill", inplace=True)

# gender transformation: 1--male, 0--female
df_train["gender"].replace(2, 0, inplace=True)

# only reserve y-m-d for `RE_DATE` and `Discharge time` columns
df_train["RE_DATE"] = df_train["RE_DATE"].dt.strftime("%Y-%m-%d")
df_train["Discharge time"] = df_train["Discharge time"].dt.strftime("%Y-%m-%d")


# %%
df_train = df_train.dropna(
    subset=["PATIENT_ID", "RE_DATE", "Discharge time"], how="any"
)

# %%
# calculate raw data's los interval
df_grouped = df_train.groupby("PATIENT_ID")

los_interval_list = []
los_interval_alive_list = []
los_interval_dead_list = []

for name, group in df_grouped:
    sorted_group = group.sort_values(by=["RE_DATE"], ascending=True)
    # print(sorted_group['outcome'])
    # print('---')
    # print(type(sorted_group))
    intervals = sorted_group["RE_DATE"].tolist()
    outcome = sorted_group["outcome"].tolist()[0]
    cur_visits_len = len(intervals)
    # print(cur_visits_len)
    if cur_visits_len == 1:
        continue
    for i in range(1, len(intervals)):
        los_interval_list.append(
            (pd.to_datetime(intervals[i]) - pd.to_datetime(intervals[i - 1])).days
        )
        if outcome == 0:
            los_interval_alive_list.append(
                (pd.to_datetime(intervals[i]) - pd.to_datetime(intervals[i - 1])).days
            )
        else:
            los_interval_dead_list.append(
                (pd.to_datetime(intervals[i]) - pd.to_datetime(intervals[i - 1])).days
            )

los_interval_list = np.array(los_interval_list)
los_interval_alive_list = np.array(los_interval_alive_list)
los_interval_dead_list = np.array(los_interval_dead_list)

output = {
    "overall": los_interval_list,
    "alive": los_interval_alive_list,
    "dead": los_interval_dead_list,
}
# pd.to_pickle(output, 'raw_tjh_los_interval_list.pkl')


# %%
# we have 2 types of prediction tasks: 1) predict mortality outcome, 2) length of stay

# below are all lab test features
labtest_features_str = """
Hypersensitive cardiac troponinI	hemoglobin	Serum chloride	Prothrombin time	procalcitonin	eosinophils(%)	Interleukin 2 receptor	Alkaline phosphatase	albumin	basophil(%)	Interleukin 10	Total bilirubin	Platelet count	monocytes(%)	antithrombin	Interleukin 8	indirect bilirubin	Red blood cell distribution width 	neutrophils(%)	total protein	Quantification of Treponema pallidum antibodies	Prothrombin activity	HBsAg	mean corpuscular volume	hematocrit	White blood cell count	Tumor necrosis factorα	mean corpuscular hemoglobin concentration	fibrinogen	Interleukin 1β	Urea	lymphocyte count	PH value	Red blood cell count	Eosinophil count	Corrected calcium	Serum potassium	glucose	neutrophils count	Direct bilirubin	Mean platelet volume	ferritin	RBC distribution width SD	Thrombin time	(%)lymphocyte	HCV antibody quantification	D-D dimer	Total cholesterol	aspartate aminotransferase	Uric acid	HCO3-	calcium	Amino-terminal brain natriuretic peptide precursor(NT-proBNP)	Lactate dehydrogenase	platelet large cell ratio 	Interleukin 6	Fibrin degradation products	monocytes count	PLT distribution width	globulin	γ-glutamyl transpeptidase	International standard ratio	basophil count(#)	2019-nCoV nucleic acid detection	mean corpuscular hemoglobin 	Activation of partial thromboplastin time	Hypersensitive c-reactive protein	HIV antibody quantification	serum sodium	thrombocytocrit	ESR	glutamic-pyruvic transaminase	eGFR	creatinine
"""

# below are 2 demographic features
demographic_features_str = """
age	gender
"""

labtest_features = [f for f in labtest_features_str.strip().split("\t")]
demographic_features = [f for f in demographic_features_str.strip().split("\t")]
target_features = ["outcome", "LOS"]

# from our observation, `2019-nCoV nucleic acid detection` feature (in lab test) are all -1 value
# so we remove this feature here
labtest_features.remove("2019-nCoV nucleic acid detection")

# %%
# if some values are negative, set it as Null
df_train[df_train[demographic_features + labtest_features] < 0] = np.nan

# %%
# merge lab tests of the same (patient_id, date)
df_train = df_train.groupby(
    ["PATIENT_ID", "RE_DATE", "Discharge time"], dropna=True, as_index=False
).mean()

# %%
# calculate length-of-stay lable
df_train["LOS"] = (
    pd.to_datetime(df_train["Discharge time"]) - pd.to_datetime(df_train["RE_DATE"])
).dt.days

# %%
# if los values are negative, set it as 0
df_train["LOS"] = df_train["LOS"].clip(lower=0)

# %%
# save features' statistics information


def calculate_statistic_info(df, features):
    """all values calculated"""
    statistic_info = {}
    len_df = len(df)
    for _, e in enumerate(features):
        h = {}
        h["count"] = int(df[e].count())
        h["missing"] = str(round(float((100 - df[e].count() * 100 / len_df)), 3)) + "%"
        h["mean"] = float(df[e].mean())
        h["max"] = float(df[e].max())
        h["min"] = float(df[e].min())
        h["median"] = float(df[e].median())
        h["std"] = float(df[e].std())
        statistic_info[e] = h
    return statistic_info


def calculate_middle_part_statistic_info(df, features):
    """calculate 5% ~ 95% percentile data"""
    statistic_info = {}
    len_df = len(df)
    # calculate 5% and 95% percentile of dataframe
    middle_part_df_info = df.quantile([0.05, 0.95])

    for _, e in enumerate(features):
        low_value = middle_part_df_info[e][0.05]
        high_value = middle_part_df_info[e][0.95]
        middle_part_df_element = df.loc[(df[e] >= low_value) & (df[e] <= high_value)][e]
        h = {}
        h["count"] = int(middle_part_df_element.count())
        h["missing"] = (
            str(round(float((100 - middle_part_df_element.count() * 100 / len_df)), 3))
            + "%"
        )
        h["mean"] = float(middle_part_df_element.mean())
        h["max"] = float(middle_part_df_element.max())
        h["min"] = float(middle_part_df_element.min())
        h["median"] = float(middle_part_df_element.median())
        h["std"] = float(middle_part_df_element.std())
        statistic_info[e] = h
    return statistic_info


# labtest_statistic_info = calculate_statistic_info(df_train, labtest_features)


# group by patient_id, then calculate lab test/demographic features' statistics information
groupby_patientid_df = df_train.groupby(
    ["PATIENT_ID"], dropna=True, as_index=False
).mean()


# calculate statistic info (all values calculated)
labtest_patientwise_statistic_info = calculate_statistic_info(
    groupby_patientid_df, labtest_features
)
demographic_statistic_info = calculate_statistic_info(
    groupby_patientid_df, demographic_features
)  # it's also patient-wise

# calculate statistic info (5% ~ 95% only)
demographic_statistic_info_2 = calculate_middle_part_statistic_info(
    groupby_patientid_df, demographic_features
)
labtest_patientwise_statistic_info_2 = calculate_middle_part_statistic_info(
    groupby_patientid_df, labtest_features
)

# take 2 statistics information's union
statistic_info = labtest_patientwise_statistic_info_2 | demographic_statistic_info_2


# %%
# observe features, export to csv file [optional]
to_export_dict = {
    "name": [],
    "missing_rate": [],
    "count": [],
    "mean": [],
    "max": [],
    "min": [],
    "median": [],
    "std": [],
}
for key in statistic_info:
    detail = statistic_info[key]
    to_export_dict["name"].append(key)
    to_export_dict["count"].append(detail["count"])
    to_export_dict["missing_rate"].append(detail["missing"])
    to_export_dict["mean"].append(detail["mean"])
    to_export_dict["max"].append(detail["max"])
    to_export_dict["min"].append(detail["min"])
    to_export_dict["median"].append(detail["median"])
    to_export_dict["std"].append(detail["std"])
to_export_df = pd.DataFrame.from_dict(to_export_dict)
# to_export_df.to_csv('statistic_info.csv')

# %%
# normalize data
def normalize_data(df, features, statistic_info):

    df_features = df[features]
    df_features = df_features.apply(
        lambda x: (x - statistic_info[x.name]["mean"])
        / (statistic_info[x.name]["std"] + 1e-12)
    )
    df = pd.concat(
        [df[["PATIENT_ID", "gender", "RE_DATE", "outcome", "LOS"]], df_features], axis=1
    )
    return df


df_train = normalize_data(
    df_train, ["age"] + labtest_features, statistic_info
)  # gender don't need to be normalized

# %%
# filter outliers
def filter_data(df, features, bar=3):
    for f in features:
        df[f] = df[f].mask(df[f].abs().gt(bar))
    return df


df_train = filter_data(df_train, demographic_features + labtest_features, bar=3)

# %%
# drop rows if all labtest_features are recorded nan
df_train = df_train.dropna(subset=labtest_features, how="all")

# %%
# Calculate data statistics after preprocessing steps (before imputation)

# Step 1: reverse z-score normalization operation
df_reverse = df_train
# reverse normalize data
def reverse_normalize_data(df, features, statistic_info):
    df_features = df[features]
    df_features = df_features.apply(
        lambda x: x * (statistic_info[x.name]["std"] + 1e-12)
        + statistic_info[x.name]["mean"]
    )
    df = pd.concat(
        [df[["PATIENT_ID", "gender", "RE_DATE", "outcome", "LOS"]], df_features], axis=1
    )
    return df


df_reverse = reverse_normalize_data(
    df_reverse, ["age"] + labtest_features, statistic_info
)  # gender don't need to be normalized

statistics = {}

for f in demographic_features + labtest_features:
    statistics[f] = {}


def calculate_quantile_statistic_info(df, features, case):
    """all values calculated"""
    for _, e in enumerate(features):
        # print(e, lo, mi, hi)
        if e == "gender":
            unique, count = np.unique(df[e], return_counts=True)
            data_count = dict(zip(unique, count))  # key = 1 male, 0 female
            print(data_count)
            male_percentage = (
                data_count[1.0] * 100 / (data_count[1.0] + data_count[0.0])
            )
            statistics[e][case] = f"{male_percentage:.2f}% Male"
            print(statistics[e][case])
        else:
            lo = round(np.nanpercentile(df[e], 25), 2)
            mi = round(np.nanpercentile(df[e], 50), 2)
            hi = round(np.nanpercentile(df[e], 75), 2)
            statistics[e][case] = f"{mi:.2f} [{lo:.2f}, {hi:.2f}]"


def calculate_missing_rate(df, features, case="missing_rate"):
    for _, e in enumerate(features):
        missing_rate = round(float(df[e].isnull().sum() * 100 / df[e].shape[0]), 2)
        statistics[e][case] = f"{missing_rate:.2f}%"


tmp_groupby_pid = df_reverse.groupby(["PATIENT_ID"], dropna=True, as_index=False).mean()

calculate_quantile_statistic_info(tmp_groupby_pid, demographic_features, "overall")
calculate_quantile_statistic_info(
    tmp_groupby_pid[tmp_groupby_pid["outcome"] == 0], demographic_features, "alive"
)
calculate_quantile_statistic_info(
    tmp_groupby_pid[tmp_groupby_pid["outcome"] == 1], demographic_features, "dead"
)

calculate_quantile_statistic_info(df_reverse, labtest_features, "overall")
calculate_quantile_statistic_info(
    df_reverse[df_reverse["outcome"] == 0], labtest_features, "alive"
)
calculate_quantile_statistic_info(
    df_reverse[df_reverse["outcome"] == 1], labtest_features, "dead"
)

calculate_missing_rate(
    df_reverse, demographic_features + labtest_features, "missing_rate"
)

export_quantile_statistics = {
    "Characteristics": [],
    "Overall": [],
    "Alive": [],
    "Dead": [],
    "Missing Rate": [],
}
for f in demographic_features + labtest_features:
    export_quantile_statistics["Characteristics"].append(f)
    export_quantile_statistics["Overall"].append(statistics[f]["overall"])
    export_quantile_statistics["Alive"].append(statistics[f]["alive"])
    export_quantile_statistics["Dead"].append(statistics[f]["dead"])
    export_quantile_statistics["Missing Rate"].append(statistics[f]["missing_rate"])

# pd.DataFrame.from_dict(export_quantile_statistics).to_csv('statistics.csv')

# %%
def calculate_data_existing_length(data):
    res = 0
    for i in data:
        if not pd.isna(i):
            res += 1
    return res


# elements in data are sorted in time ascending order
def fill_missing_value(data, to_fill_value=0):
    data_len = len(data)
    data_exist_len = calculate_data_existing_length(data)
    if data_len == data_exist_len:
        return data
    elif data_exist_len == 0:
        # data = [to_fill_value for _ in range(data_len)]
        for i in range(data_len):
            data[i] = to_fill_value
        return data
    if pd.isna(data[0]):
        # find the first non-nan value's position
        not_na_pos = 0
        for i in range(data_len):
            if not pd.isna(data[i]):
                not_na_pos = i
                break
        # fill element before the first non-nan value with median
        for i in range(not_na_pos):
            data[i] = to_fill_value
    # fill element after the first non-nan value
    for i in range(1, data_len):
        if pd.isna(data[i]):
            data[i] = data[i - 1]
    return data


# %%
# fill missing data using our strategy and convert to time series records
grouped = df_train.groupby("PATIENT_ID")

all_x_demographic = []
all_x_labtest = []
all_y = []
all_missing_mask = []

for name, group in grouped:
    sorted_group = group.sort_values(by=["RE_DATE"], ascending=True)
    patient_demographic = []
    patient_labtest = []
    patient_y = []

    for f in demographic_features + labtest_features:
        to_fill_value = (statistic_info[f]["median"] - statistic_info[f]["mean"]) / (
            statistic_info[f]["std"] + 1e-12
        )
        # take median patient as the default to-fill missing value
        # print(sorted_group[f].values)
        fill_missing_value(sorted_group[f].values, to_fill_value)
        # print(sorted_group[f].values)
        # print('-----------')
    all_missing_mask.append(
        (
            np.isfinite(
                sorted_group[demographic_features + labtest_features].to_numpy()
            )
        ).astype(int)
    )

    for _, v in sorted_group.iterrows():
        patient_y.append([v["outcome"], v["LOS"]])
        demo = []
        lab = []
        for f in demographic_features:
            demo.append(v[f])
        for f in labtest_features:
            lab.append(v[f])
        patient_labtest.append(lab)
        patient_demographic.append(demo)
    all_y.append(patient_y)
    all_x_demographic.append(patient_demographic[-1])
    all_x_labtest.append(patient_labtest)

# all_x_demographic (2 dim, record each patients' demographic features)
# all_x_labtest (3 dim, record each patients' lab test features)
# all_y (3 dim, patients' outcome/los of all visits)

# %%
all_x_labtest = np.array(all_x_labtest, dtype=object)
x_lab_length = [len(_) for _ in all_x_labtest]
x_lab_length = torch.tensor(x_lab_length, dtype=torch.int)
max_length = int(x_lab_length.max())
all_x_labtest = [torch.tensor(_) for _ in all_x_labtest]
# pad lab test sequence to the same shape
all_x_labtest = torch.nn.utils.rnn.pad_sequence((all_x_labtest), batch_first=True)

all_x_demographic = torch.tensor(all_x_demographic)
batch_size, demo_dim = all_x_demographic.shape
# repeat demographic tensor
all_x_demographic = torch.reshape(
    all_x_demographic.repeat(1, max_length), (batch_size, max_length, demo_dim)
)
# demographic tensor concat with lab test tensor
all_x = torch.cat((all_x_demographic, all_x_labtest), 2)

all_y = np.array(all_y, dtype=object)
all_y = [torch.Tensor(_) for _ in all_y]
# pad [outcome/los] sequence as well
all_y = torch.nn.utils.rnn.pad_sequence((all_y), batch_first=True)

all_missing_mask = np.array(all_missing_mask, dtype=object)
all_missing_mask = [torch.tensor(_) for _ in all_missing_mask]
all_missing_mask = torch.nn.utils.rnn.pad_sequence((all_missing_mask), batch_first=True)

# %%
# save pickle format dataset (export torch tensor)
pd.to_pickle(all_x, f"./datasets/tongji/processed_data/x.pkl")
pd.to_pickle(all_y, f"./datasets/tongji/processed_data/y.pkl")
pd.to_pickle(x_lab_length, f"./datasets/tongji/processed_data/visits_length.pkl")
pd.to_pickle(all_missing_mask, f"./datasets/tongji/processed_data/missing_mask.pkl")

# %%
# Calculate patients' outcome statistics (patients-wise)
outcome_list = []
y_outcome = all_y[:, :, 0]
indices = torch.arange(len(x_lab_length), dtype=torch.int64)
for i in indices:
    outcome_list.append(y_outcome[i][0].item())
outcome_list = np.array(outcome_list)
print(len(outcome_list))
unique, count = np.unique(outcome_list, return_counts=True)
data_count = dict(zip(unique, count))
print(data_count)

# %%
# Calculate patients' outcome statistics (records-wise)
outcome_records_list = []
y_outcome = all_y[:, :, 0]
indices = torch.arange(len(x_lab_length), dtype=torch.int64)
for i in indices:
    outcome_records_list.extend(y_outcome[i][0 : x_lab_length[i]].tolist())
outcome_records_list = np.array(outcome_records_list)
print(len(outcome_records_list))
unique, count = np.unique(outcome_records_list, return_counts=True)
data_count = dict(zip(unique, count))
print(data_count)

# %%
# Calculate patients' mean los and 95% percentile los
los_list = []
y_los = all_y[:, :, 1]
indices = torch.arange(len(x_lab_length), dtype=torch.int64)
for i in indices:
    # los_list.extend(y_los[i][: x_lab_length[i].long()].tolist())
    los_list.append(y_los[i][0].item())
los_list = np.array(los_list)
print(los_list.mean() * 0.5)
print(np.median(los_list) * 0.5)
print(np.percentile(los_list, 95))

print("median:", np.median(los_list))
print("Q1:", np.percentile(los_list, 25))
print("Q3:", np.percentile(los_list, 75))

# %%
los_alive_list = np.array(
    [los_list[i] for i in range(len(los_list)) if outcome_list[i] == 0]
)
los_dead_list = np.array(
    [los_list[i] for i in range(len(los_list)) if outcome_list[i] == 1]
)
print(len(los_alive_list))
print(len(los_dead_list))

print("[Alive]")
print("median:", np.median(los_alive_list))
print("Q1:", np.percentile(los_alive_list, 25))
print("Q3:", np.percentile(los_alive_list, 75))

print("[Dead]")
print("median:", np.median(los_dead_list))
print("Q1:", np.percentile(los_dead_list, 25))
print("Q3:", np.percentile(los_dead_list, 75))

# %%
tjh_los_statistics = {
    "overall": los_list,
    "alive": los_alive_list,
    "dead": los_dead_list,
}
# pd.to_pickle(tjh_los_statistics, 'tjh_los_statistics.pkl')

# %%
# calculate visits length Median [Q1, Q3]
visits_list = np.array(x_lab_length)
visits_alive_list = np.array(
    [x_lab_length[i] for i in range(len(x_lab_length)) if outcome_list[i] == 0]
)
visits_dead_list = np.array(
    [x_lab_length[i] for i in range(len(x_lab_length)) if outcome_list[i] == 1]
)
print(len(visits_alive_list))
print(len(visits_dead_list))

print("[Total]")
print("median:", np.median(visits_list))
print("Q1:", np.percentile(visits_list, 25))
print("Q3:", np.percentile(visits_list, 75))

print("[Alive]")
print("median:", np.median(visits_alive_list))
print("Q1:", np.percentile(visits_alive_list, 25))
print("Q3:", np.percentile(visits_alive_list, 75))

print("[Dead]")
print("median:", np.median(visits_dead_list))
print("Q1:", np.percentile(visits_dead_list, 25))
print("Q3:", np.percentile(visits_dead_list, 75))

# %%
# Length-of-stay interval (overall/alive/dead)
los_interval_list = []
los_interval_alive_list = []
los_interval_dead_list = []

y_los = all_y[:, :, 1]
indices = torch.arange(len(x_lab_length), dtype=torch.int64)
for i in indices:
    cur_visits_len = x_lab_length[i]
    if cur_visits_len == 1:
        continue
    for j in range(1, cur_visits_len):
        los_interval_list.append(y_los[i][j - 1] - y_los[i][j])
        if outcome_list[i] == 0:
            los_interval_alive_list.append(y_los[i][j - 1] - y_los[i][j])
        else:
            los_interval_dead_list.append(y_los[i][j - 1] - y_los[i][j])

los_interval_list = np.array(los_interval_list)
los_interval_alive_list = np.array(los_interval_alive_list)
los_interval_dead_list = np.array(los_interval_dead_list)

output = {
    "overall": los_interval_list,
    "alive": los_interval_alive_list,
    "dead": los_interval_dead_list,
}
# pd.to_pickle(output, 'tjh_los_interval_list.pkl')

# %%
len(los_interval_list), len(los_interval_alive_list), len(los_interval_dead_list)

# %%
def check_nan(x):
    if np.isnan(np.sum(x.cpu().numpy())):
        print("some values from input are nan")
    else:
        print("no nan")