'''Inspired by several scripts at:
https://www.kaggle.com/c/facebook-v-predicting-check-ins/scripts
Special thanks to Sandro for starting the madness. :-)
https://www.kaggle.com/svpons/facebook-v-predicting-check-ins/grid-plus-classifier

KNN accelerated V3.
Numpy argsort replaced with faster JIT compiled function for finding 
top three predictions.

'''

import numpy as np
import pandas as pd
from sklearn.naive_bayes import GaussianNB
import time
from datetime import timedelta
import gc
from numba import jit

# Found at: https://www.kaggle.com/rshekhar2/facebook-v-predicting-check-ins/xgboost-cv-example-with-small-bug
def mapkprecision(truthvalues, predictions):
    '''
    This is a faster implementation of MAP@k valid for numpy arrays.
    It is only valid when there is one single truth value. 

    m ~ number of observations
    k ~ MAP at k -- in this case k should equal 3

    truthvalues.shape = (m,) 
    predictions.shape = (m, k)
    '''
    z = (predictions == truthvalues[:, None]).astype(np.float32)
    weights = 1./(np.arange(predictions.shape[1], dtype=np.float32) + 1.)
    z = z * weights[None, :]
    return np.mean(np.sum(z, axis=1))

def load_data(data_name):
    types = {'row_id': np.dtype(np.int32),
         'x': np.dtype(float),
         'y' : np.dtype(float),
         'accuracy': np.dtype(np.int16),
         'place_id': np.int64,
         'time': np.dtype(np.int32)}
    df = pd.read_csv(data_name, dtype=types, na_filter=False)
    return df

def process_one_cell(cell_train, cell_test, fw, th, n_neighbors):
    
    # Remove infrequent places
    places, idx, counts = np.unique(cell_train[:, -1], return_inverse=True, 
                                    return_counts=True)
    count_per_row = counts[idx]
    cell_train = cell_train[count_per_row >= th]

    # Store row_ids for test
    row_ids = cell_test[:, -1].flatten().astype(np.int32)
    cell_test = cell_test[:, :-1]
    
    # Get predictions
    y_pred, classes = get_preds(cell_train, cell_test, n_neighbors)
    
    # Get top three predictions
    y_pred_labels = top_three_preds(y_pred)
    pred_labels = classes[y_pred_labels]
    cell_pred = np.column_stack((row_ids, pred_labels)).astype(np.int64) 
    
    return cell_pred
    
def get_preds(cell_train, cell_test, n_neighbors):
    # Preparing data
    y = cell_train[:, -1].flatten().astype(np.int64)
    X = cell_train[:, :-1]
    
    #Applying the classifier
    clf = GaussianNB()
    clf.fit(X, y)
    y_pred = clf.predict_proba(cell_test)
    return y_pred, clf.classes_
    
# Faster than argsort at getting top three predictions
@jit(nopython=True)
def top_three_preds(preds):
    place_count = preds.shape[0]
    preds_count = preds.shape[1]
    pred_labels = np.zeros((place_count, 3), dtype=np.int32)
    for i in range(place_count):
        first_place_score = 0
        second_place_score = 0
        third_place_score = 0
        for j in range(preds_count):
            this_pred = preds[i,j]
            if (this_pred > 0) and (this_pred > third_place_score):
                if this_pred > second_place_score:
                    pred_labels[i, 2] = pred_labels[i, 1]
                    third_place_score = second_place_score
                    if this_pred > first_place_score:
                        pred_labels[i, 1] = pred_labels[i, 0]
                        second_place_score = first_place_score
                        first_place_score = this_pred
                        pred_labels[i, 0] = j
                    else:
                        second_place_score = this_pred
                        pred_labels[i, 1] = j
                else:
                    third_place_score = this_pred
                    pred_labels[i, 2] = j
    return pred_labels
    
# Precompute the trig values
def time_trig(max_time):
    time_array = np.linspace(0, 2*np.pi, max_time)
    sin_values = np.sin(time_array)
    cos_values = np.cos(time_array)
    return (sin_values, cos_values)
    
# Generate a dictionary of the time limits so it doesn't have to be 
# recalculated each loop
def create_time_dict(t_cuts, time_mod, time_weight, time_aug):
    
    t_slice = 24 / t_cuts
    time_dict = dict()
    trig_array = time_trig(time_mod)
    for t in range(t_cuts):
        
        t_min = int(t * t_slice * 12)
        t_max = int((t + 1) * t_slice * 12 - 1)
        sin_t_start = trig_array[0][t_min] * time_weight
        sin_t_stop = trig_array[0][t_max] * time_weight
        cos_t_start = trig_array[1][t_min] * time_weight
        cos_t_stop = trig_array[1][t_max] * time_weight
        sin_t_min = min((sin_t_start, sin_t_stop))
        sin_t_max = max((sin_t_start, sin_t_stop))
        cos_t_min = min((cos_t_start, cos_t_stop))
        cos_t_max = max((cos_t_start, cos_t_stop))
        time_dict[t] = [sin_t_min, sin_t_max, cos_t_min, cos_t_max]

        t_min = int((t * t_slice - time_aug) * 12)%time_mod
        t_max = int(((t + 1) * t_slice + time_aug)* 12 - 1)%time_mod
        sin_t_start = trig_array[0][t_min] * time_weight
        sin_t_stop = trig_array[0][t_max] * time_weight
        cos_t_start = trig_array[1][t_min] * time_weight
        cos_t_stop = trig_array[1][t_max] * time_weight
        sin_t_min = min((sin_t_start, sin_t_stop, sin_t_min))
        sin_t_max = max((sin_t_start, sin_t_stop, sin_t_max))
        cos_t_min = min((cos_t_start, cos_t_stop, cos_t_min))
        cos_t_max = max((cos_t_start, cos_t_stop, cos_t_max))
        time_dict[t] += [sin_t_min, sin_t_max, cos_t_min, cos_t_max]
        
    return time_dict

@jit
def apply_mask(data, feature, mask_min, mask_max):
    mask = (data[:, feature] >= mask_min)
    mask = mask & (data[:, feature] < mask_max)      
    return data[mask]    

def process_grid(train, test, x_cuts, y_cuts, t_cuts,
                 x_border_aug, y_border_aug, time_aug, fw, th, n_neighbors):
    preds_list = []
    x_slice = train[:, 0].max() / x_cuts
    y_slice = train[:, 1].max() / y_cuts
    time_mod = 288
    time_weight = fw[2]
    time_dict = create_time_dict(t_cuts, time_mod, time_weight, time_aug)

    for i in range(x_cuts):
        row_start_time = time.time()
        x_min = x_slice * i
        x_max = x_slice * (i+1)
        x_max += int((i+1) == x_cuts) # expand edge at end

        col_test = apply_mask(test, 0, x_min, x_max)
        x_min -= x_border_aug
        x_max += x_border_aug
        col_train = apply_mask(train, 0, x_min, x_max)

        for j in range(y_cuts):
            y_min = y_slice * j
            y_max = y_slice * (j+1)
            y_max += int((j+1) == y_cuts) # expand edge at end

            row_test = apply_mask(col_test, 1, y_min, y_max)
            y_min -= y_border_aug
            y_max += y_border_aug
            row_train = apply_mask(col_train, 1, y_min, y_max)

            for t in range(t_cuts):
                #print(df_row_test.shape, df_row_train.shape)
                t_lim = time_dict[t]
                mask = (row_test[:, 2] >= t_lim[0])
                mask = mask & (row_test[:, 2] <= t_lim[1])
                mask = mask & (row_test[:, 3] >= t_lim[2])
                mask = mask & (row_test[:, 3] <= t_lim[3])
                cell_test = row_test[mask]
                mask = (row_train[:, 2] >= t_lim[4])
                mask = mask & (row_train[:, 2] <= t_lim[5])
                mask = mask & (row_train[:, 3] >= t_lim[6])
                mask = mask & (row_train[:, 3] <= t_lim[7])
                cell_train = row_train[mask]
                cell_pred = process_one_cell(cell_train, cell_test, 
                                             fw, th, n_neighbors)
                preds_list.append(cell_pred)
        elapsed = (time.time() - row_start_time)
        print('Row', i, 'completed in:', timedelta(seconds=elapsed))
    preds = np.vstack(preds_list)
    return preds

# Thank you Alex!
# From: https://www.kaggle.com/drarfc/facebook-v-predicting-check-ins/fastest-way-to-write-the-csv
def generate_submission(preds):    
    print('Writing submission file')
    with open('KNN_submission.csv', "w") as out:
        out.write("row_id,place_id\n")
        rows = ['']*8607230
        n=0
        for num in range(8607230):
            rows[n]='%d,%d %d %d\n' % (preds[num,0],preds[num,1],preds[num,2],preds[num,3])
            n=n+1
        out.writelines(rows)

def validation_split(df, val_start_day):
    day = df['time']//1440
    df_val = df.loc[(day>=val_start_day)].copy()
#    df = df.loc[(day<val_start_day)].copy()
    df = df.loc[(day>=val_start_day-365) & (day<val_start_day)].copy()
    return df, df_val
    
def remove_infrequent_places_df(df, th=5):
    place_counts = df.place_id.value_counts()
    mask = (place_counts[df.place_id.values] >= th).values
    df = df[mask]
    return df

def prepare_data(datapath, val_start_day, train_columns, test_columns, 
                 fw, th, off):
    val_label = None
    print('Loading train data')
    df_train = load_data(datapath + 'train.csv')
    if val_start_day > 0:
        # Create validation data
        df_train, df_test = validation_split(df_train, val_start_day)
        val_label = df_test['place_id'] 
        df_test.drop(['place_id'], axis=1, inplace=True)    
    print('Feature engineering on train')
    df_train.drop(['row_id'], axis=1, inplace=True)
    df_train = remove_infrequent_places_df(df_train, th)
    gc.collect()
    df_train = feature_engineering(df_train, off)
    df_train = apply_weights(df_train, fw)
    # reorder the columns so the place id is at the end
    train = df_train[train_columns].values
    del df_train
    gc.collect()
    if val_start_day == 0:
        print('Loading test data')
        df_test = load_data(datapath + 'test.csv') 
    print('Feature engineering on test')
    df_test = feature_engineering(df_test, off)
    df_test = apply_weights(df_test, fw)
    test = df_test[test_columns].values
    del df_test
    gc.collect()
    return train, test, val_label
        
def apply_weights(df, fw):
    df['accuracy'] *= fw[0]
    df['day_of_year_sin'] *= fw[1]
    df['day_of_year_cos'] *= fw[1]
    df['minute_sin'] *= fw[2]
    df['minute_cos'] *= fw[2]
    df['weekday_sin'] *= fw[3]
    df['weekday_cos'] *= fw[3]
    df.x *= fw[4]
    df.y *= fw[5]
    df['year'] *= fw[6]
    return df

def feature_engineering(df, off):
    minute =((df["time"]+off[0])//5)%288
    trig_arrays = time_trig(288)
    df['minute_sin'] = trig_arrays[0][minute]
    df['minute_cos'] = trig_arrays[1][minute]
    del minute
    day = ((df['time']+off[1])//1440)%365
    trig_arrays = time_trig(365)
    df['day_of_year_sin'] = trig_arrays[0][day]
    df['day_of_year_cos'] = trig_arrays[1][day]
    del day
    weekday = ((df['time']+off[2])//1440)%7
    trig_arrays = time_trig(7)
    df['weekday_sin'] = trig_arrays[0][weekday]
    df['weekday_cos'] = trig_arrays[1][weekday]
    del weekday
    df['year'] = (df['time']//525600).astype(float)
    df.drop(['time'], axis=1, inplace=True)
    df['accuracy'] = np.log10(df['accuracy']).astype(float)
    return df
    
print('Starting...')
start_time = time.time()
# Global variables
datapath = '../input/'
# Change val_start_day to zero to generate predictions
val_start_day = 455 # Day at which to cut validation
th = 5 # Threshold at which to cut places from train
fw = [134., 33.6, 64.4, 26., 2300, 5625, 55.6]
off = [444, 931, 421]

# Defining the size of the grid
x_cuts = 12 # number of cuts along x 
y_cuts = 25 # number of cuts along y
#TODO: More general solution for t_cuts. For now must be 4.
t_cuts = 4 # number of cuts along time. 
x_border_aug = 0.005 * fw[4] # expansion of x border on train 
y_border_aug = 0.004 * fw[5] # expansion of y border on train
time_aug = 2.9
n_neighbors = 0
columns = ['x', 'y', 'minute_sin', 'minute_cos', 'accuracy',
           'day_of_year_sin', 'day_of_year_cos', 
           'weekday_sin', 'weekday_cos', 'year']
train_columns = columns + ['place_id']
test_columns  = columns + ['row_id']

train, test, val_label = prepare_data(datapath, val_start_day,
                                      train_columns, test_columns, fw, th, off)

elapsed = (time.time() - start_time)
print('Data prepared in:', timedelta(seconds=elapsed))
    
preds = process_grid(train, test, x_cuts, y_cuts, t_cuts,
                     x_border_aug, y_border_aug, time_aug, 
                     fw, th, n_neighbors)
elapsed = (time.time() - start_time)
print('Predictions made in:', timedelta(seconds=elapsed))

if val_start_day > 0:
    preds = preds[preds[:, 0] > 0] # only use rows predicted
    labels = val_label.loc[preds[:, 0]].values
    score = mapkprecision(labels, preds[:, 1:])
    print('Final score:', score)
else:
    print('Pred shape:', preds.shape)
    generate_submission(preds)
elapsed = (time.time() - start_time)
print('Task completed in:', timedelta(seconds=elapsed))