##### SETUP ######

import pickle

import numpy as np

import pandas as pd

import math

import scipy.interpolate as interpolate

import matplotlib.pyplot as plt

import pylab as P

import numpy as np

import statsmodels.api as sm

from scipy import stats 

## Set up age intervals

## Float instability created errors, so corrected the following trunction 

intervals = np.trunc(np.concatenate((np.array([0.01]),np.arange(.05,.2,.05),np.arange(.4,2,.2), np.arange(2,20,.5)))*100)/100

##################

##### VARIABLES #####

## Minimum number of datapoints for patient to pass filter

min_datapoints = 5

## CSV raw document locaiton

csv_location = "../../data/csv/BMI.csv"

## Individuals with outlier data

## 16262: First two datapoints when ages ~1.8 have heights of ~2 inches

bad_patients = [16262]

#####################

##### FUNCTIONS #####

## Change numeric datatypes to appropriate floats of ints

def change_dtypes(df):

    if "wt" in df.columns:

        df["wt"] = df["wt"].astype(np.float64)

    else:

        df["wt_ounces"] = df["wt_ounces"].astype(np.float64)

    df["age"] = df["age"].astype(np.float64)

    df["bmi"] = df["bmi"].astype(np.float64)

    df["id"] = df["id"].astype(np.int64)

    df["ht"] = df["ht"].astype(np.float64)

    return df

## Fit a OLS curve to the group patient's y as a function of x

def fit(group_patient, header_y, header_x, alpha = 0.01):

    x = group_patient[header_x]

    y = np.log(group_patient[header_y])

    ## Count number of datapoints for this patient

    num_x = group_patient.shape[0]

    ## np.ones necessary to include constant in regression

    X = np.column_stack((x*x*x, x*x, x, np.ones(num_x)))

    ## OLS fit

    model = sm.OLS(y,X)

    results = model.fit()

    ## Option 1: Outliers defined to be either datapoints with large residuals

    infl = results.get_influence()

    resids = infl.resid_studentized_external

    ## Calculate residual cutoff

    resid_cutoff = stats.t.ppf(1.-alpha/2, results.df_resid)

    large_resid = np.abs(resids) > resid_cutoff

    ## Option 2: Outliers defined to be either datapoints with cooks distance greater than 4/N

    #(c, p) = infl.cooks_distance

    #cooks_cutoff = 4.0/group_patient.shape[0]

    #large_cooks = c > cooks_cutoff

    ## Option 3: .. or both

    #outliers_bool = np.logical_or(outliers_bool_norm_resid, outliers_bool_cooks_distance)

    ## Currently choose option 3 for defining outliers

    return large_resid

## Fit a OLS curve to the group patient's y as a function of x

## Has more restrictive critera than fit(), and therefore better to use when

## patient's measurements do not fluctuate as much

def fit2(group_patient, header_y, header_x):

    x = group_patient[header_x]

    y = group_patient[header_y]

    ## Count number of datapoints for this patient

    num_x = group_patient.shape[0]

    ## np.ones necessary to include constant in regression

    X = np.column_stack((np.log(x), np.ones(num_x)))

    ## OLS fit

    model = sm.OLS(y,X)

    fitted = model.fit()

    ## Print summar of fit

    #print fitted.summary()

    ## Option 1: Outliers defined to be either datapoints with residuals greater than 2.5 standard deviations

    outliers_bool_norm_resid = abs(fitted.norm_resid()) > 2.5

    ## Option 2: Outliers defined to be either datapoints with cooks distance greater than 4/(N-K-1)

    #influence = fitted.get_influence()

    #(c, p) = influence.cooks_distance

    #outliers_bool_cooks_distance = c > 4.0/(num_x-2)

    ## Option 3: .. or both

    #outliers_bool = np.logical_or(outliers_bool_norm_resid, outliers_bool_cooks_distance)

    ## Currently choose option 3 for defining outliers

    return outliers_bool_norm_resid

#####################

## Read in document

df = pd.read_csv(csv_location)

## Drop columns not needed

df = df.drop(["Height (ft&inc)","Body Mass Index Percentile","Age on 5/1/2012","BMI","Patient Ethnicity","Patient Race"], axis=1)

df.rename(columns={"Patient ID": "id", "Gender": "gender", "Race/Ethnicity": "race_ethnicity", "Patient Age At Encounter": "age", \

                   "Body Mass Index": "bmi", "Weight (Ounces)":"wt_ounces", "Height (Inches)":"ht"}, inplace=True)

## Convert numerical columns to floats (and remove #VALUE! if necessary)

cond = df["ht"] == "#VALUE!"

df["ht"][cond] = "NaN"

df = change_dtypes(df)

## Fill in missing values: need at least one of bmi, wt, ht to resolve third

selector = np.isnan(df["bmi"]) & ~np.isnan(df["wt_ounces"]) & ~np.isnan(df["ht"])

missing_series = df[selector]["wt_ounces"]/16/(np.power(df[selector]["ht"],2))*703

df["bmi"] = pd.concat([df["bmi"].dropna(), missing_series.dropna()]).reindex_like(df)

selector = ~np.isnan(df["bmi"]) & np.isnan(df["wt_ounces"]) & ~np.isnan(df["ht"])

missing_series = df[selector]["bmi"]*np.power(df[selector]["ht"],2)/703*16

df["wt"] = pd.concat([df["wt_ounces"].dropna(), missing_series.dropna()]).reindex_like(df)

selector = ~np.isnan(df["bmi"]) & ~np.isnan(df["wt_ounces"]) & np.isnan(df["ht"])

missing_series = np.power(df[selector]["wt_ounces"]/16/(df[selector]["bmi"])*703,0.5)

df["ht"] = pd.concat([df["ht"].dropna(), missing_series.dropna()]).reindex_like(df)

## Remove rows that can't be resolved for bmi, wt, or ht

selector = ~np.isnan(df["bmi"]) & ~np.isnan(df["wt_ounces"]) & ~np.isnan(df["ht"]) & df["gender"].notnull()

df = df[selector]

## Convert from ounces to pounds

df["wt"] = df["wt_ounces"]/16

df = df.drop("wt_ounces", axis=1)

## Remove bad patients

for patient_id in bad_patients:

    df[df["id"] == patient_id] = np.nan

df = df.dropna()

## Remove datapoints where ht = 0

df = df[df["ht"] != 0]

## Sort columns by header id and reindex

df = df.sort(["id", "age"], ascending=[1,1])

df = df.reindex()

## Group by patient id

grouped_patients = df.groupby("id")

df_resampled_list = []

df_filtered_list = []

df_filtered_out_list = []

df_outliers_list = []

## Iterate through each patient in the dataframe

## Note: using pandas to append is O(n^2), whereas appending list of dictionaries is O(n)

## Takes ~24 mins

for name_patient, group_patient in grouped_patients:

    ## To grab patients, type:

    #grouped_patients.get_group(np.int64(5035))

    ## Skip patients that have less than minimum number of datapoints

    if group_patient.shape[0] < min_datapoints:

        group_patient.apply(lambda x: df_filtered_out_list.append(x.to_dict()), axis = 1)

        #df_filtered_out = pd.DataFrame.append(df_filtered_out, group_patient)

        continue

    ## Change datapoints with age 0 to age 0.01 (in order to take its log for the regression)

    group_patient["age"][group_patient["age"] == np.float64(0)] = np.float64(0.01)

    ## Conduct regression of ht vs age and wt vs age to determine outliers

    ## Use fit2() if patient only has data from 15 yrs old and on, as

    ## patient's measurements do not change much

    if group_patient["age"].min() > 15:

        outliers_bool_ht = fit2(group_patient, "ht", "age")

        outliers_bool_wt = fit2(group_patient, "wt", "age")

    ## Otherwise use fit()

    else:

        outliers_bool_ht = fit(group_patient, "ht", "age")

        outliers_bool_wt = fit(group_patient, "wt", "age")

    outliers_bool = np.logical_or(outliers_bool_ht, outliers_bool_wt)

    ## Remove outlier datapoints and save removed datapoints to df_outliers

    if any(outliers_bool):

        group_patient[outliers_bool].apply(lambda x: df_outliers_list.append(x.to_dict()), axis = 1)

    group_patient = group_patient[np.logical_not(outliers_bool)]

    ## Check again after removing outliers

    ## Skip patients that have less than minimum number of datapoints

    if group_patient.shape[0] < min_datapoints:

        group_patient.apply(lambda x: df_filtered_out_list.append(x.to_dict()), axis = 1)

        continue

    ## Interpolate ht and wt linearly and resample at defined age intervals

##    f_ht = interpolate.interp1d(group_patient["age"], np.log(group_patient["ht"]), kind='cubic', bounds_error=False)

##    f_wt = interpolate.interp1d(group_patient["age"], np.log(group_patient["wt"]), kind='cubic', bounds_error=False)

##    df_resampled_patient = pd.DataFrame(intervals, columns=["age"])

##    df_resampled_patient["ht"] = np.power(math.e, df_resampled_patient["age"].apply(f_ht))

##    df_resampled_patient["wt"] = np.power(math.e, df_resampled_patient["age"].apply(f_wt))

    f_ht = interpolate.interp1d(group_patient["age"], group_patient["ht"], kind='linear', bounds_error=False)

    f_wt = interpolate.interp1d(group_patient["age"], group_patient["wt"], kind='linear', bounds_error=False)

    df_resampled_patient = pd.DataFrame(intervals, columns=["age"])

    df_resampled_patient["ht"] = df_resampled_patient["age"].apply(f_ht)

    df_resampled_patient["wt"] = df_resampled_patient["age"].apply(f_wt)

    ## Drop NAs, which occur because no extrapolation is conducted

    df_resampled_patient = df_resampled_patient.dropna()

    ## Recalculate BMIs

    df_resampled_patient["bmi"] = df_resampled_patient["wt"]/(np.power(df_resampled_patient["ht"],2))*703

    ## Save patient attributes

    df_resampled_patient["id"] = name_patient

    df_resampled_patient["gender"] = group_patient["gender"].max()

    df_resampled_patient["race_ethnicity"] = group_patient["race_ethnicity"].max()

    ## Save filtered and resampled data

    df_resampled_patient.apply(lambda x: df_resampled_list.append(x.to_dict()), axis = 1)

    group_patient.apply(lambda x: df_filtered_list.append(x.to_dict()), axis = 1)

## Convert list of dictionaries to dataframes

df_resampled = change_dtypes(pd.DataFrame(df_resampled_list).dropna())

df_filtered = change_dtypes(pd.DataFrame(df_filtered_list).dropna())

df_filtered_out = change_dtypes(pd.DataFrame(df_filtered_out_list).dropna())

df_outliers = change_dtypes(pd.DataFrame(df_outliers_list).dropna())

## Print dataframes to CSVs

df_resampled.to_csv("../../data/csv/BMI_resampled.csv", index_label=False, index=False)

df_filtered.to_csv("../../data/csv/BMI_filtered.csv", index_label=False, index=False)

df_filtered_out.to_csv("../../data/csv/BMI_filtered_out.csv", index_label=False, index=False)

df_outliers.to_csv("../../data/csv/BMI_outliers.csv", index_label=False, index=False)

## Dump files with pickle using the highest protocol available

output = open('../../data/pkl/BMI_resampled.pkl', 'wb')

pickle.dump(df_resampled, output, -1)

output.close()

output = open('../../data/pkl/BMI_filtered.pkl', 'wb')

pickle.dump(df_filtered, output, -1)

output.close()

output = open('../../data/pkl/BMI_filtered_out.pkl', 'wb')

pickle.dump(df_filtered_out, output, -1)

output.close()

output = open('../../data/pkl/BMI_outliers.pkl', 'wb')

pickle.dump(df_outliers, output, -1)

output.close()