# Import required packages
import random 
import pandas as pd # version 1.4.4
import numpy as np # version 1.21.5
import os
import datetime as dt 
from pytz import FixedOffset # version 2022.1
import io

from scipy.stats import randint, beta, uniform, reciprocal # version 1.9.1
from scipy.signal import find_peaks 
from scipy.integrate import simps 
from scipy.ndimage import median_filter

# Ignore pandas future warnings
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

from sklearn.base import BaseEstimator, ClassifierMixin, TransformerMixin # version 1.2.2
from sklearn.model_selection import (RandomizedSearchCV, GridSearchCV, GroupKFold, ParameterGrid)  
from sklearn.metrics import (balanced_accuracy_score, mean_absolute_error) 
from sklearn.utils.fixes import loguniform 
from sklearn.preprocessing import StandardScaler 
from sklearn.linear_model import (LinearRegression, SGDRegressor)
from sklearn.svm import SVR 
from sklearn.pipeline import Pipeline

from lightgbm import LGBMRegressor # version 3.3.5
import json 
import joblib # version 1.2.0

import multiprocessing # version
n_jobs = int(max(1, multiprocessing.cpu_count())/2)

# set path to data
data_path= # add data path

# Here, we merge TAC_df and ref_df to make train_df.
# First, create a temporary column in ref_df with the time interval upper bound
ref_df['time_upper_bound'] = ref_df['time'] + pd.Timedelta(hours=5)

# Merge the two DataFrames based on participant_id
merged_df = pd.merge(TAC_df, ref_df, on='participant_id', how='right')

# Filter rows within the 5-hour intervals
merged_df = merged_df[(merged_df['datetime'] >= merged_df['time']) & (merged_df['datetime'] < merged_df['time_upper_bound'])]

# Group the DataFrames by participant_id and time, and aggregate datetime and TAC ug/L(air) values in lists
grouped_df = merged_df.groupby(['participant_id', 'time']).agg({'datetime': list, 'TAC ug/L(air)': list}).reset_index()

# Merge the aggregated values back to ref_df
train_df = pd.merge(ref_df, grouped_df, on=['participant_id', 'time'])

# Drop the temporary column in ref_df
train_df = train_df.drop(columns=['time_upper_bound'])

class DrinkDetector(BaseEstimator, ClassifierMixin, TransformerMixin):
    """A scikit-learn compatible estimator that conducts procedures I and II of the model.
    
    Parameters
    ----------
        size: median_filter kernel size
        window: rolling average window size
        distance: peak minimum distance in peak detection algorithm
        prominence: peak minimum prominence
        wlen: window length
        width: peak minimum width
        
        Please see the following website for more details on the parameters: 
        https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.median_filter.html
        https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.rolling.html
        https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.find_peaks.html
    """
    
    def __init__(self, size=39, window=206, distance=265, prominence=8.0,
                 wlen=2202, width=36):
        self.size = size
        self.window = window
        self.distance = distance
        self.prominence = prominence
        self.wlen = wlen
        self.width = width
        
    def fit(self, X, y=None):
        """
        Parameters
        ----------
        X : TAC data collected with Skyn. X is a numpy array with three 
        columns "participant_id","datetime","TAC ug/L(air)".
        
        y: EMA app/self-reported drinking event start times and number of drinkgs. y is a numpy array
        with columns for "participant_id","time","drinking_timestamp","drinking_event","ema_n_drinks"
        """
        return self
        
    def predict(self, X):
        X=pd.DataFrame(X,columns=["participant_id","datetime","TAC ug/L(air)"])

        # Make X in long format with one row per participan TAC timestamp 
        data = [{'participant_id': row.participant_id, 'datetime': dt, 'TAC ug/L(air)': tac}
                for row in X.itertuples() for dt, tac in zip(row.datetime, row._3)]
        X = pd.DataFrame(data)
        X.index=X.datetime
        X = X[["participant_id","TAC ug/L(air)"]]
        
        y_pred = pd.DataFrame(columns=["participant_id","index_test","right_bases","peak_maximum","peak_auc"]) 
        for participant_id in X.iloc[:,0].unique():
            # Get data for each participant.
            p_data = X.query('participant_id == @participant_id').copy() # participant level TAC data
            
            # PROCEDURE I (SIGNAL FILTERING)
            ## Median filter
            p_data["medfilt"] = median_filter(p_data["TAC ug/L(air)"], size=self.size)

            ## Moving average on median filter
            p_data["medfilt_rolling"] = p_data["medfilt"].rolling(window=self.window).mean()   
            ## Remove missing values created in procedure I
            p_data = p_data.loc[p_data["medfilt_rolling"].notna()]

            # PROCEDURE II (PEAK DETECTION)
            peaks, properties = find_peaks(p_data["medfilt_rolling"].values, 
                                           distance = self.distance, 
                                           prominence = (self.prominence, None),
                                           wlen = self.wlen,
                                           width = (self.width,None))
            
            ## store peak properties: Left bases are the index test, model detected drinking event start times
            index_test = p_data["medfilt_rolling"].index[properties["left_bases"]]
            ## get right bases of detected peak too
            rbs = p_data["medfilt_rolling"].index[properties["right_bases"]]
            max_peak = p_data.loc[p_data.index.isin(p_data["medfilt_rolling"].index[peaks]),"medfilt_rolling"].values          
            dict_df = {'participant_id': participant_id, 'index_test':index_test,'right_bases':rbs, 'peak_maximum':max_peak} 
            peaks_props = pd.DataFrame(dict_df) 
            ## AUC: we calculate area under the peak curve and add it to peaks_props
            peaks_props["peak_auc"] = [simps(dx=1, y=p_data[it:rb].medfilt_rolling)
                                       for it, rb in zip(peaks_props.index_test, peaks_props.right_bases)]
            y_pred = pd.concat([y_pred, peaks_props])
        
        y_pred = y_pred.values
        return y_pred
    
    def main_analysis(self,X, y=None):
        y=pd.DataFrame(y,columns=["participant_id","time","drinking_timestamp","drinking_event","ema_n_drinks"]) 
        y.participant_id=y.participant_id.astype("int64")
        y.drinking_event=y.drinking_event.astype("int64")
        y.ema_n_drinks=y.ema_n_drinks.astype(float)
        y.sort_values(by="drinking_timestamp",inplace=True) # EMA app recorded drinking start time
        
        
        y_pred = self.predict(X)
        y_pred = pd.DataFrame(data=y_pred, columns=["participant_id","index_test","right_bases","peak_maximum","peak_auc"])
        y_pred['index_test'] = pd.to_datetime(y_pred['index_test'], 
                                              utc=True).dt.tz_convert(FixedOffset(-240))

        y_pred["peak_id"]=np.arange(0,y_pred.shape[0]) # make peak id in y_pred
        y_pred.sort_values(by="index_test",inplace=True)
        y_pred.participant_id=y_pred.participant_id.astype("int64")
        
        
        # make a df for true pos. and false neg.
        TP_FN=pd.merge_asof(left=y[y.drinking_event ==1 ], right=y_pred,
                            by="participant_id", left_on="drinking_timestamp", right_on="index_test",
                            allow_exact_matches=True, 
                            direction="nearest", 
                            tolerance=pd.Timedelta("5h")) 
        
        TP_FN["time_difference"]=abs(TP_FN.drinking_timestamp-TP_FN.index_test)
        # for the duplicates, keep the ones with smaller time_difference values
        # To do that, first find duplicate peak_ids with larger time difference
        code_nan=TP_FN.loc[(TP_FN.peak_id.notna())&(TP_FN.duplicated(subset=["peak_id"],keep=False))
                      ].groupby(by="index_test").max()
        
        if code_nan.shape[0] > 0:
            print(code_nan.shape[0],"duplicates were generated while calculating true pos. false neg.",end='\r')
            # Set them to nan:
            TP_FN.loc[(TP_FN.peak_id.isin(code_nan.peak_id))&
                      (TP_FN.time_difference.isin(code_nan.time_difference)),
                      ['index_test','peak_id']]=np.nan
        # drop time_difference
        TP_FN.drop(columns=["time_difference"],inplace=True) 
        
        # Next, Make a df for true neg and concat it to TP_FN
        TN=y.loc[y.drinking_timestamp.isna()]
        TP_FN_TN=pd.concat([TP_FN,TN])
        
        # Next, Make a df for false pos. (these are detected peaks not in TP_FN)
        FP=y_pred.loc[~y_pred.peak_id.isin(TP_FN.peak_id)]
        FP=FP.rename(columns={"index_test":"false_p"})

        # merge FP to TP_FN_TN: If more than 1 FP in a 5-hour interval, just one of them is counted
        TP_FN_TN.sort_values(by="time",inplace=True)
        FP.sort_values(by="false_p",inplace=True)
        TP_FN_TN_FP=pd.merge_asof(left=TP_FN_TN, right=FP, suffixes=('_TP', '_FP'), 
                                  by="participant_id", left_on="time", right_on="false_p",
                                  allow_exact_matches=True, direction="forward",
                                  tolerance=pd.Timedelta("5h"))
        
        # Make a true_label and pred_label for scoring
        TP_FN_TN_FP.rename(columns={"drinking_event":"true_label"},inplace=True) 
        TP = TP_FN_TN_FP.query('index_test.notnull() &'
                               'drinking_timestamp.notnull()').index
        FN = TP_FN_TN_FP.query('index_test.isnull() &'
                               'drinking_timestamp.notnull()').index
        FP = TP_FN_TN_FP.query('false_p.notnull() &'
                               'index_test.isnull()').index
        TN = TP_FN_TN_FP.query('index_test.isnull() &'
                               'drinking_timestamp.isnull() &'
                               'false_p.isnull()').index
        
        TP_FN_TN_FP.loc[TP,"pred_label"]=1 # True pos       
        TP_FN_TN_FP.loc[FN,"pred_label"]=0 # False neg
        # if for a drinking interval, there are more than one peak, 
        # it is counted as one True positive
        TP_FN_TN_FP.loc[FP,"pred_label"]=1 # False pos
        TP_FN_TN_FP.loc[TN,"pred_label"]=0 # True neg.

        # Compute the AUC score based on the binary labels and return it
        score = balanced_accuracy_score(TP_FN_TN_FP['true_label'], TP_FN_TN_FP['pred_label'])
        
        return score, TP_FN_TN_FP
    
    def score(self,X, y=None):
        score, _ = self.main_analysis(X, y)
        return score
    
    def outputs(self,X, y=None):
        _, outputs = self.main_analysis(X, y)
        outputs = outputs[["participant_id","true_label","pred_label",
                           "ema_n_drinks","peak_maximum_TP","peak_auc_TP"]]
        return outputs

# set up X,y, and group
X = np.array(train_df[["participant_id","datetime","TAC ug/L(air)"]])
y = np.array(train_df[["participant_id","time","drinking_timestamp","drinking_event","ema_n_drinks"]])
groups = np.array(train_df["participant_id"])

# Set X and y 
X = np.array(DrinkDetector_outputs[["peak_maximum_TP","peak_auc_TP"]])
y = np.array(DrinkDetector_outputs["ema_n_drinks"])
groups = np.array(DrinkDetector_outputs["participant_id"])

# Update pipeline based on results from random grid search
pipeline = Pipeline([
            ('scaler', StandardScaler()),
            ('SVR', SVR(C = 0.672, epsilon = 1.065, gamma = 'scale', kernel = 'sigmoid',
                        shrinking = False))
])