# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
import pandas as pd
import numpy as np

from datetime import datetime, timezone, timedelta
# -

# ### Load datasets

events_df = pd.read_pickle("data/01_events_df.pkl")

category_tree_df = pd.read_pickle("data/01_category_tree_df.pkl")

item_properties_df = pd.read_pickle("data/01_item_properties_df.pkl")

# ### Data Preparation

# Define variables
TIME_KEY = "timestamp"
USER_KEY = "visitorid"
ITEM_KEY = "itemid"
SESSION_KEY = "sessionid"

# Remove rows with NA userId
events_prepared_df = events_df[~np.isnan(events_df[USER_KEY])].copy()

# Truncate milli seconds
events_prepared_df[TIME_KEY] = (events_prepared_df[TIME_KEY] / 1000).astype(int)

# #### Introduce Sessions

# Sort data by user and time
events_prepared_df.sort_values(by=[USER_KEY, TIME_KEY], ascending=True, inplace=True)

# Compute the time difference between queries
tdiff = np.diff(events_prepared_df[TIME_KEY].values)

# Check which of them are bigger then session threshold
SESSION_THRESHOLD = 30 * 60
split_session = tdiff > SESSION_THRESHOLD
split_session = np.r_[True, split_session]

# Check when the user changes its data
new_user = events_prepared_df[USER_KEY].values[1:] != events_prepared_df[USER_KEY].values[:-1]
new_user = np.r_[True, new_user]

# A new sessions stars when at least one of the two conditions is verified
new_session = np.logical_or(new_user, split_session)

# Compute the session ids
session_ids = np.cumsum(new_session)
events_prepared_df[SESSION_KEY] = session_ids

events_prepared_df.head()


def print_stats(data):
    data_start = datetime.fromtimestamp(data[TIME_KEY].min(), timezone.utc)
    data_end = datetime.fromtimestamp(data[TIME_KEY].max(), timezone.utc)
    print('\tEvents: {}\n\tUsers: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {}'.
          format(len(data), data[USER_KEY].nunique(),
                 data[SESSION_KEY].nunique(), data[ITEM_KEY].nunique(),
                 data_start.date().isoformat(), data_end.date().isoformat()))


print("Raw dataset:")
print_stats(events_prepared_df)

# #### Filter

# Keep items with >=5 interactions
MIN_ITEM_SUPPORT = 5
item_pop = events_prepared_df[ITEM_KEY].value_counts()
good_items = item_pop[item_pop >= MIN_ITEM_SUPPORT].index
events_prepared_df = events_prepared_df[events_prepared_df[ITEM_KEY].isin(good_items)]

# Remove sessions with length < 2
MIN_SESSION_LENGTH = 2
session_length = events_prepared_df[SESSION_KEY].value_counts()
good_sessions = session_length[session_length >= MIN_SESSION_LENGTH].index
events_prepared_df = events_prepared_df[events_prepared_df[SESSION_KEY].isin(good_sessions)]

# let's keep only returning users (with >= 2 sessions)
# need to be 3, because we need at least 1 for each training, validation and test set
MIN_USER_SESSIONS = 3
MAX_USER_SESSIONS = None
sess_per_user = events_prepared_df.groupby(USER_KEY)[SESSION_KEY].nunique()
if MAX_USER_SESSIONS is None:  # no filter for max number of sessions for each user
    good_users = sess_per_user[(sess_per_user >= MIN_USER_SESSIONS)].index
else:
    good_users = sess_per_user[(sess_per_user >= MIN_USER_SESSIONS) & (sess_per_user < MAX_USER_SESSIONS)].index
events_prepared_df = events_prepared_df[events_prepared_df[USER_KEY].isin(good_users)]

print("Filtered dataset:")
print_stats(events_prepared_df)

# ### Create single train / test split

CLEAN_TEST = True
def last_session_out_split(data, min_session_length):
    """
    last-session-out split
    assign the last session of every user to the test set and the remaining ones to the training set
    """
    sessions = data.sort_values(by=[USER_KEY, TIME_KEY]).groupby(USER_KEY)[SESSION_KEY]
    last_session = sessions.last()
    train = data[~data[SESSION_KEY].isin(last_session.values)].copy()
    test = data[data[SESSION_KEY].isin(last_session.values)].copy()
    if CLEAN_TEST:
        train_items = train[ITEM_KEY].unique()
        test = test[test[ITEM_KEY].isin(train_items)]
        
        #  Remove sessions in test shorter than min_session_length
        slen = test[SESSION_KEY].value_counts()
        good_sessions = slen[slen >= min_session_length].index
        test = test[test[SESSION_KEY].isin(good_sessions)].copy()
        train = train.reset_index(drop=True)
        test = test.reset_index(drop=True)
    return train, test


# assign the last session of every user to the test set and the remaining ones to the training set
train_sessions, test_sessions = last_session_out_split(events_prepared_df, MIN_SESSION_LENGTH)
validation_train_sessions, validation_test_sessions = last_session_out_split(train_sessions, MIN_SESSION_LENGTH)

print("Training set:")
print_stats(train_sessions)

print("Test set:")
print_stats(test_sessions)

print("Validation training set:")
print_stats(validation_train_sessions)

print("Validation test set:")
print_stats(validation_test_sessions)

# ### Store Datasets

events_prepared_df.to_pickle("data/02_events_prepared_df.pkl")

train_sessions.to_pickle("data/02_train_sessions.pkl")

test_sessions.to_pickle("data/02_test_sessions.pkl")

validation_train_sessions.to_pickle("data/02_validation_train_sessions.pkl")

validation_test_sessions.to_pickle("data/02_validation_test_sessions.pkl")

StarcoderdataJupyterScriptsDedupFiltered