# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/datasets/datasets.base.ipynb (unless otherwise specified).

__all__ = ['SessionDataset', 'SessionDatasetv2', 'SessionDatasetv3', 'SessionDatasetv4', 'GraphDataset',
           'SessionGraphDataset', 'RatingDataset', 'TorchSessionDataset', 'TorchRatingDataset',
           'TorchPairwiseNegativeDataset']

# Cell
from typing import List, Optional, Callable, Union, Any, Tuple

import numpy as np
import pandas as pd
from tqdm import tqdm
tqdm.pandas()

from abc import *
from pathlib import Path
import os
import os.path as osp
from collections.abc import Sequence
import sys
import tempfile
import shutil
import pickle
import time
import csv
import math
import operator
import itertools
from collections import defaultdict
import datetime
from datetime import date, timezone, timedelta
from pandas import Timedelta

import torch

from ..utils.common_utils import download_url, extract_zip, extract_gz, makedirs

# Internal Cell
def to_list(value: Any) -> Sequence:
    if isinstance(value, Sequence) and not isinstance(value, str):
        return value
    else:
        return [value]

def files_exist(files: List[str]) -> bool:
    # NOTE: We return `False` in case `files` is empty, leading to a
    # re-processing of files on every instantiation.
    return len(files) != 0 and all([osp.exists(f) for f in files])

# Internal Cell
class Dataset:
    """Dataset base class
    """
    @property
    def raw_file_names(self) -> Union[str, List[str], Tuple]:
        r"""The name of the files in the :obj:`self.raw_dir` folder that must
        be present in order to skip downloading."""
        raise NotImplementedError

    @property
    def processed_file_names(self) -> Union[str, List[str], Tuple]:
        r"""The name of the files in the :obj:`self.processed_dir` folder that
        must be present in order to skip processing."""
        raise NotImplementedError

    def download(self):
        r"""Downloads the dataset to the :obj:`self.raw_dir` folder."""
        raise NotImplementedError

    def process(self):
        r"""Processes the dataset to the :obj:`self.processed_dir` folder."""
        raise NotImplementedError

    def __init__(self, root=None):
        self.root = root

        if 'download' in self.__class__.__dict__:
            self._download()

        # if 'process' in self.__class__.__dict__:
        #     self._process()

    @property
    def raw_dir(self) -> str:
        return osp.join(self.root, 'raw')

    @property
    def processed_dir(self) -> str:
        return osp.join(self.root, 'processed')

    @property
    def raw_paths(self) -> List[str]:
        r"""The absolute filepaths that must be present in order to skip
        downloading."""
        files = to_list(self.raw_file_names)
        return [osp.join(self.raw_dir, f) for f in files]

    @property
    def processed_paths(self) -> List[str]:
        r"""The absolute filepaths that must be present in order to skip
        processing."""
        files = to_list(self.processed_file_names)
        return [osp.join(self.processed_dir, f) for f in files]

    def _download(self):
        if files_exist(self.raw_paths):  # pragma: no cover
            return

        makedirs(self.raw_dir)
        self.download()

    def _process(self):
        if files_exist(self.processed_paths):  # pragma: no cover
            return

        print('Processing...', file=sys.stderr)

        makedirs(self.processed_dir)
        self.process()

        print('Done!', file=sys.stderr)

    def __repr__(self) -> str:
        arg_repr = str(len(self)) if len(self) > 1 else ''
        return f'{self.__class__.__name__}({arg_repr})'

# Cell
class SessionDataset(Dataset):
    r"""Session data base class.

    Args:
        min_session_length (int): Minimum number of items for a session to be valid
        min_item_support (int): Minimum number of interactions for an item to be valid
        eval_sec (int): these many seconds from the end will be taken as validation data

    References:
        1. https://github.com/Ethan-Yys/GRU4REC-pytorch-master/blob/master/preprocessing.py
    """
    def __init__(self,
                 root,
                 min_session_length: int = None,
                 min_item_support: int = None,
                 eval_sec: int = 86400,  # valid/test data is one day = 86400 seconds
                 ):
        super().__init__(root)
        self.min_session_length = min_session_length
        self.min_item_support = min_item_support
        self.eval_sec = eval_sec

        self._process()

    def load_ratings_df(self):
        r"""load raw dataset into pandas dataframe"""
        raise NotImplementedError

    def remove_short_sessions(self, df):
        r"""delete sessions where session length is less than `min_session_length`"""
        # groupby session id and get size of each session
        session_len = df.groupby('uid').size()
        df = df[np.in1d(df.uid, session_len[session_len >= self.min_session_length].index)]
        return df

    def remove_sparse_items(self, df):
        r"""delete records of items which appeared less than `min_item_support` times"""
        # groupby itemID and get size of each item
        item_len = df.groupby('sid').size()
        df = df[np.in1d(df.sid, item_len[item_len >= self.min_item_support].index)]
        # remove sessions of less than n interactions again
        df = self.remove_short_sessions(df)
        return df

    def split_df(self, df):
        timeMax = df.timestamp.max() #maximum time in all records
        sessionMaxTime = df.groupby('uid').timestamp.max() #group by uid and get the maximum time of each session
        sessionTrain = sessionMaxTime[sessionMaxTime < (timeMax - self.eval_sec)].index #training split is all sessions that ended before the `eval_sec` seconds
        sessionTest  = sessionMaxTime[sessionMaxTime >= (timeMax - self.eval_sec)].index #testing split is all sessions has records in the `eval_sec` seconds
        train = df[np.in1d(df.uid, sessionTrain)]
        test = df[np.in1d(df.uid, sessionTest)]
        # Delete records in testing split where items are not in training split
        test = test[np.in1d(test.sid, train.sid)]
        # Delete Sessions in testing split
        test = self.remove_short_sessions(test)
        print('Training Set has {} Events, {} Sessions, and {} Items\n\n'.format(
            len(train), train.uid.nunique(), train.sid.nunique()))
        print('Validation Set has {} Events, {} Sessions, and {} Items\n\n'.format(
            len(test), test.uid.nunique(), test.sid.nunique()))
        return train, test

# Cell
class SessionDatasetv2(Dataset):
    def __init__(self, root, column_names):
        super().__init__(root)
        self.SESSION_ID = column_names['SESSION_ID']
        self.ITEM_ID = column_names['ITEM_ID']
        self.TIMEFRAME = column_names['TIMEFRAME']
        self.EVENT_DATE = column_names['EVENT_DATE']

        self._process()

    @property
    def raw_file_names(self) -> str:
        raise NotImplementedError

    @property
    def processed_file_names(self) -> str:
        return ['train.txt','test.txt','all_train_seq.txt']

    def download(self):
        raise NotImplementedError

    def process(self):
        with open(self.raw_paths[0], "r") as f:
            reader = csv.DictReader(f, delimiter=';')
            sess_clicks = {}
            sess_date = {}
            ctr = 0
            curid = -1
            curdate = None
            for data in reader:
                sessid = data[self.SESSION_ID]
                if curdate and not curid == sessid:
                    date = ''
                    date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))
                    sess_date[curid] = date
                curid = sessid
                item = data[self.ITEM_ID], int(data[self.TIMEFRAME])
                curdate = ''
                curdate = data[self.EVENT_DATE]

                if sessid in sess_clicks:
                    sess_clicks[sessid] += [item]
                else:
                    sess_clicks[sessid] = [item]
                ctr += 1
            date = ''
            date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))
            for i in list(sess_clicks):
                sorted_clicks = sorted(sess_clicks[i], key=operator.itemgetter(1))
                sess_clicks[i] = [c[0] for c in sorted_clicks]
            sess_date[curid] = date

        print("-- Reading data")

        # Filter out length 1 sessions
        for s in list(sess_clicks):
            if len(sess_clicks[s]) == 1:
                del sess_clicks[s]
                del sess_date[s]

        # Count number of times each item appears
        iid_counts = {}
        for s in sess_clicks:
            seq = sess_clicks[s]
            for iid in seq:
                if iid in iid_counts:
                    iid_counts[iid] += 1
                else:
                    iid_counts[iid] = 1

        sorted_counts = sorted(iid_counts.items(), key=operator.itemgetter(1))

        length = len(sess_clicks)
        for s in list(sess_clicks):
            curseq = sess_clicks[s]
            filseq = list(filter(lambda i: iid_counts[i] >= 5, curseq))
            if len(filseq) < 2:
                del sess_clicks[s]
                del sess_date[s]
            else:
                sess_clicks[s] = filseq

        # Split out test set based on dates
        dates = list(sess_date.items())
        maxdate = dates[0][1]

        for _, date in dates:
            if maxdate < date:
                maxdate = date

        # 7 days for test
        splitdate = 0
        splitdate = maxdate - 86400 * 7

        print('Splitting date', splitdate)
        tra_sess = filter(lambda x: x[1] < splitdate, dates)
        tes_sess = filter(lambda x: x[1] > splitdate, dates)

        # Sort sessions by date
        tra_sess = sorted(tra_sess, key=operator.itemgetter(1))
        tes_sess = sorted(tes_sess, key=operator.itemgetter(1))

        print(len(tra_sess))
        print(len(tes_sess))
        print(tra_sess[:3])
        print(tes_sess[:3])

        print("-- Splitting train set and test set")

        item_dict = {}
        # Convert training sessions to sequences and renumber items to start from 1
        def obtian_tra():
            train_ids = []
            train_seqs = []
            train_dates = []
            item_ctr = 1
            for s, date in tra_sess:
                seq = sess_clicks[s]
                outseq = []
                for i in seq:
                    if i in item_dict:
                        outseq += [item_dict[i]]
                    else:
                        outseq += [item_ctr]
                        item_dict[i] = item_ctr
                        item_ctr += 1
                if len(outseq) < 2:  # Doesn't occur
                    continue
                train_ids += [s]
                train_dates += [date]
                train_seqs += [outseq]
            print(item_ctr)     # 43098, 37484
            return train_ids, train_dates, train_seqs

        # Convert test sessions to sequences, ignoring items that do not appear in training set
        def obtian_tes():
            test_ids = []
            test_seqs = []
            test_dates = []
            for s, date in tes_sess:
                seq = sess_clicks[s]
                outseq = []
                for i in seq:
                    if i in item_dict:
                        outseq += [item_dict[i]]
                if len(outseq) < 2:
                    continue
                test_ids += [s]
                test_dates += [date]
                test_seqs += [outseq]
            return test_ids, test_dates, test_seqs

        tra_ids, tra_dates, tra_seqs = obtian_tra()
        tes_ids, tes_dates, tes_seqs = obtian_tes()

        def process_seqs(iseqs, idates):
            out_seqs = []
            out_dates = []
            labs = []
            ids = []
            for id, seq, date in zip(range(len(iseqs)), iseqs, idates):
                for i in range(1, len(seq)):
                    tar = seq[-i]
                    labs += [tar]
                    out_seqs += [seq[:-i]]
                    out_dates += [date]
                    ids += [id]
            return out_seqs, out_dates, labs, ids

        tr_seqs, tr_dates, tr_labs, tr_ids = process_seqs(tra_seqs, tra_dates)
        te_seqs, te_dates, te_labs, te_ids = process_seqs(tes_seqs, tes_dates)
        tra = (tr_seqs, tr_labs)
        tes = (te_seqs, te_labs)

        print(len(tr_seqs))
        print(len(te_seqs))
        print(tr_seqs[:3], tr_dates[:3], tr_labs[:3])
        print(te_seqs[:3], te_dates[:3], te_labs[:3])

        all = 0

        for seq in tra_seqs:
            all += len(seq)
        for seq in tes_seqs:
            all += len(seq)
        print('avg length: ', all/(len(tra_seqs) + len(tes_seqs) * 1.0))

        pickle.dump(tra, open(self.processed_paths[0], 'wb'))
        pickle.dump(tes, open(self.processed_paths[1], 'wb'))
        pickle.dump(tra_seqs, open(self.processed_paths[2], 'wb'))

# Cell
class SessionDatasetv3(Dataset):
    def __init__(self, root):
        super().__init__(root)

        self._process()

    @property
    def raw_file_names(self) -> str:
        raise NotImplementedError

    @property
    def processed_file_names(self) -> str:
        return ['train.txt','test.txt','all_train_seq.txt']

    def download(self):
        raise NotImplementedError

    def process(self):
        def get_session_id(df, interval):
            df_prev = df.shift()
            is_new_session = (df.userId != df_prev.userId) | (
                df.timestamp - df_prev.timestamp > interval
            )
            session_id = is_new_session.cumsum() - 1
            return session_id

        def group_sessions(df, interval):
            sessionId = get_session_id(df, interval)
            df = df.assign(sessionId=sessionId)
            return df

        def filter_short_sessions(df, min_len=2):
            session_len = df.groupby('sessionId', sort=False).size()
            long_sessions = session_len[session_len >= min_len].index
            df_long = df[df.sessionId.isin(long_sessions)]
            return df_long

        def filter_infreq_items(df, min_support=5):
            item_support = df.groupby('itemId', sort=False).size()
            freq_items = item_support[item_support >= min_support].index
            df_freq = df[df.itemId.isin(freq_items)]
            return df_freq

        def filter_until_all_long_and_freq(df, min_len=2, min_support=5):
            while True:
                df_long = filter_short_sessions(df, min_len)
                df_freq = filter_infreq_items(df_long, min_support)
                if len(df_freq) == len(df):
                    break
                df = df_freq
            return df

        def truncate_long_sessions(df, max_len=20, is_sorted=False):
            if not is_sorted:
                df = df.sort_values(['sessionId', 'timestamp'])
            itemIdx = df.groupby('sessionId').cumcount()
            df_t = df[itemIdx < max_len]
            return df_t

        def update_id(df, field):
            labels = pd.factorize(df[field])[0]
            kwargs = {field: labels}
            df = df.assign(**kwargs)
            return df

        def remove_immediate_repeats(df):
            df_prev = df.shift()
            is_not_repeat = (df.sessionId != df_prev.sessionId) | (df.itemId != df_prev.itemId)
            df_no_repeat = df[is_not_repeat]
            return df_no_repeat

        def reorder_sessions_by_endtime(df):
            endtime = df.groupby('sessionId', sort=False).timestamp.max()
            df_endtime = endtime.sort_values().reset_index()
            oid2nid = dict(zip(df_endtime.sessionId, df_endtime.index))
            sessionId_new = df.sessionId.map(oid2nid)
            df = df.assign(sessionId=sessionId_new)
            df = df.sort_values(['sessionId', 'timestamp'])
            return df

        def keep_top_n_items(df, n):
            item_support = df.groupby('itemId', sort=False).size()
            top_items = item_support.nlargest(n).index
            df_top = df[df.itemId.isin(top_items)]
            return df_top

        def split_by_time(df, timedelta):
            max_time = df.timestamp.max()
            end_time = df.groupby('sessionId').timestamp.max()
            split_time = max_time - timedelta
            train_sids = end_time[end_time < split_time].index
            df_train = df[df.sessionId.isin(train_sids)]
            df_test = df[~df.sessionId.isin(train_sids)]
            return df_train, df_test

        def train_test_split(df, test_split=0.2):
            endtime = df.groupby('sessionId', sort=False).timestamp.max()
            endtime = endtime.sort_values()
            num_tests = int(len(endtime) * test_split)
            test_session_ids = endtime.index[-num_tests:]
            df_train = df[~df.sessionId.isin(test_session_ids)]
            df_test = df[df.sessionId.isin(test_session_ids)]
            return df_train, df_test

        def save_sessions(df, filepath):
            df = reorder_sessions_by_endtime(df)
            sessions = df.groupby('sessionId').itemId.apply(lambda x: ','.join(map(str, x)))
            sessions.to_csv(filepath, sep='\t', header=False, index=False)

        def save_dataset(df_train, df_test):
            # filter items in test but not in train
            df_test = df_test[df_test.itemId.isin(df_train.itemId.unique())]
            df_test = filter_short_sessions(df_test)

            print(f'No. of Clicks: {len(df_train) + len(df_test)}')
            print(f'No. of Items: {df_train.itemId.nunique()}')

            # update itemId
            train_itemId_new, uniques = pd.factorize(df_train.itemId)
            df_train = df_train.assign(itemId=train_itemId_new)
            oid2nid = {oid: i for i, oid in enumerate(uniques)}
            test_itemId_new = df_test.itemId.map(oid2nid)
            df_test = df_test.assign(itemId=test_itemId_new)

            print('saving dataset')
            save_sessions(df_train, self.processed_paths[0])
            save_sessions(df_test, self.processed_paths[1])
            num_items = len(uniques)
            with open(self.processed_paths[2], 'w') as f:
                f.write(str(num_items))

        df = pd.read_csv(
            osp.join(self.raw_dir,self.raw_file_names),
            sep='\t',
            header=None,
            names=['userId', 'timestamp', 'itemId'],
            usecols=[0, 1, 4],
            parse_dates=['timestamp'],
            infer_datetime_format=True,
        )

        print('start preprocessing')
        df = df.dropna()
        df = update_id(df, 'userId')
        df = update_id(df, 'itemId')
        df = df.sort_values(['userId', 'timestamp'])

        df = group_sessions(df, Timedelta(days=1))
        df = remove_immediate_repeats(df)
        df = truncate_long_sessions(df, is_sorted=True)
        df = keep_top_n_items(df, n=30000)
        df = filter_until_all_long_and_freq(df)
        df_train, df_test = train_test_split(df, test_split=0.2)
        save_dataset(df_train, df_test)

# Cell
class SessionDatasetv4(Dataset):
    r"""Session dataset base class.

    Args:
        root (string): Root directory where the dataset should be saved.
        process_method (string):
            last: last day => test set
            last_min_date: last day => test set, but from a minimal date onwards
            days_test: last N days => test set
            slice: create multiple train-test-combinations with a sliding window approach
        min_date (string): Minimum date
        session_length (int): Session time length :default = 30 * 60 #30 minutes
        min_session_length (int): Minimum number of items for a session to be valid
        min_item_support (int): Minimum number of interactions for an item to be valid
        num_slices (int): Offset in days from the first date in the data set
        days_offset (int): Number of days the training start date is shifted after creating one slice
        days_shift (int): Days shift
        days_train (int): Days in train set in each slice
        days_test (int): Days in test set in each slice
    """

    def __init__(self, root, process_method, min_date=None,
                 session_length=None, min_session_length=None, min_item_support=None,
                 num_slices=None, days_offset=None, days_shift=None, days_train=None,
                 days_test=None, data=None):
        super().__init__(root)
        self.process_method = process_method
        self.min_date = min_date
        self.session_length = session_length
        self.min_session_length = min_session_length
        self.min_item_support = min_item_support
        self.num_slices = num_slices
        self.days_offset = days_offset
        self.days_shift = days_shift
        self.days_train = days_train
        self.days_test = days_test
        self.data = None

        self._process()

    @property
    def raw_file_names(self) -> str:
        raise NotImplementedError

    @property
    def processed_file_names(self) -> str:
        raise NotImplementedError

    def download(self):
        raise NotImplementedError

    def load(self):
        raise NotImplementedError

    def filter_data(self):
        data = self.data

        #filter session length
        session_lengths = data.groupby('SessionId').size()
        data = data[np.in1d(data.SessionId, session_lengths[session_lengths>1].index)]

        #filter item support
        item_supports = data.groupby('ItemId').size()
        data = data[np.in1d(data.ItemId, item_supports[item_supports>= self.min_item_support].index)]

        #filter session length
        session_lengths = data.groupby('SessionId').size()
        data = data[np.in1d(data.SessionId, session_lengths[session_lengths>= self.min_session_length].index)]

        #output
        data_start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc)
        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)

        print('Filtered data set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {}\n\n'.
              format(len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.date().isoformat(), data_end.date().isoformat()))

        self.data = data

    def filter_min_date(self):
        data = self.data

        min_datetime = datetime.datetime.strptime(self.min_date + ' 00:00:00', '%Y-%m-%d %H:%M:%S')

        #filter
        session_max_times = data.groupby('SessionId').Time.max()
        session_keep = session_max_times[session_max_times > min_datetime.datetime.timestamp()].index

        data = data[np.in1d(data.SessionId, session_keep)]

        #output
        data_start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc)
        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)

        print('Filtered data set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {}\n\n'.
              format(len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.date().isoformat(), data_end.date().isoformat()))

        self.data = data

    def split_data_org(self):
        data = self.data
        tmax = data.Time.max()
        session_max_times = data.groupby('SessionId').Time.max()
        session_train = session_max_times[session_max_times < tmax-86400].index
        session_test = session_max_times[session_max_times >= tmax-86400].index
        train = data[np.in1d(data.SessionId, session_train)]
        test = data[np.in1d(data.SessionId, session_test)]
        test = test[np.in1d(test.ItemId, train.ItemId)]
        tslength = test.groupby('SessionId').size()
        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]
        print('Full train set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}'.format(len(train), train.SessionId.nunique(), train.ItemId.nunique()))
        train.to_csv(osp.join(self.processed_dir,'events_train_full.txt'), sep='\t', index=False)
        print('Test set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}'.format(len(test), test.SessionId.nunique(), test.ItemId.nunique()))
        test.to_csv(osp.join(self.processed_dir,'events_test.txt'), sep='\t', index=False)

        tmax = train.Time.max()
        session_max_times = train.groupby('SessionId').Time.max()
        session_train = session_max_times[session_max_times < tmax-86400].index
        session_valid = session_max_times[session_max_times >= tmax-86400].index
        train_tr = train[np.in1d(train.SessionId, session_train)]
        valid = train[np.in1d(train.SessionId, session_valid)]
        valid = valid[np.in1d(valid.ItemId, train_tr.ItemId)]
        tslength = valid.groupby('SessionId').size()
        valid = valid[np.in1d(valid.SessionId, tslength[tslength>=2].index)]
        print('Train set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}'.format(len(train_tr), train_tr.SessionId.nunique(), train_tr.ItemId.nunique()))
        train_tr.to_csv(osp.join(self.processed_dir,'events_train_tr.txt'), sep='\t', index=False)
        print('Validation set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}'.format(len(valid), valid.SessionId.nunique(), valid.ItemId.nunique()))
        valid.to_csv(osp.join(self.processed_dir,'events_train_valid.txt'), sep='\t', index=False)

    def split_data(self):
        data = self.data
        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)
        test_from = data_end - timedelta(self.days_test)

        session_max_times = data.groupby('SessionId').Time.max()
        session_train = session_max_times[session_max_times < test_from.timestamp()].index
        session_test = session_max_times[session_max_times >= test_from.timestamp()].index
        train = data[np.in1d(data.SessionId, session_train)]
        test = data[np.in1d(data.SessionId, session_test)]
        test = test[np.in1d(test.ItemId, train.ItemId)]
        tslength = test.groupby('SessionId').size()
        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]
        print('Full train set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}'.format(len(train), train.SessionId.nunique(), train.ItemId.nunique()))
        train.to_csv(osp.join(self.processed_dir,'events_train_full.txt'), sep='\t', index=False)
        print('Test set\n\tEvents: {}\n\tSessions: {}\n\tItems: {}'.format(len(test), test.SessionId.nunique(), test.ItemId.nunique()))
        test.to_csv(osp.join(self.processed_dir,'events_test.txt'), sep='\t', index=False)

    def slice_data(self):
        for slice_id in range(0, self.num_slices):
            self.split_data_slice(slice_id, self.days_offset+(slice_id*self.days_shift))

    def split_data_slice(self, slice_id, days_offset):
        data = self.data
        data_start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc)
        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)

        print('Full data set {}\n\tEvents: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {}'.
            format(slice_id, len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.isoformat(), data_end.isoformat()))

        start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc ) + timedelta(days_offset)
        middle =  start + timedelta(self.days_train)
        end =  middle + timedelta(self.days_test)

        #prefilter the timespan
        session_max_times = data.groupby('SessionId').Time.max()
        greater_start = session_max_times[session_max_times >= start.timestamp()].index
        lower_end = session_max_times[session_max_times <= end.timestamp()].index
        data_filtered = data[np.in1d(data.SessionId, greater_start.intersection(lower_end))]

        print('Slice data set {}\n\tEvents: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {} / {}'.
            format( slice_id, len(data_filtered), data_filtered.SessionId.nunique(), data_filtered.ItemId.nunique(), start.date().isoformat(), middle.date().isoformat(), end.date().isoformat() ) )

        #split to train and test
        session_max_times = data_filtered.groupby('SessionId').Time.max()
        sessions_train = session_max_times[session_max_times < middle.timestamp()].index
        sessions_test = session_max_times[session_max_times >= middle.timestamp()].index

        train = data[np.in1d(data.SessionId, sessions_train)]

        print('Train set {}\n\tEvents: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {}'.
            format( slice_id, len(train), train.SessionId.nunique(), train.ItemId.nunique(), start.date().isoformat(), middle.date().isoformat() ) )

        train.to_csv(osp.join(self.processed_dir,'events_train_full.'+str(slice_id)+'.txt'), sep='\t', index=False)

        test = data[np.in1d(data.SessionId, sessions_test)]
        test = test[np.in1d(test.ItemId, train.ItemId)]

        tslength = test.groupby('SessionId').size()
        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]

        print('Test set {}\n\tEvents: {}\n\tSessions: {}\n\tItems: {}\n\tSpan: {} / {} \n\n'.
            format( slice_id, len(test), test.SessionId.nunique(), test.ItemId.nunique(), middle.date().isoformat(), end.date().isoformat() ) )

        test.to_csv(osp.join(self.processed_dir,'events_test.'+str(slice_id)+'.txt'), sep='\t', index=False)

    def process(self):
        self.load()
        self.filter_data()
        if self.process_method == 'last':
            self.split_data_org()
        elif self.process_method == 'last_min_date':
            self.filter_min_date()
            self.split_data_org()
        elif self.process_method == 'days_test':
            self.split_data()
        elif self.process_method == 'slice':
            self.slice_data()

# Cell
class GraphDataset():
    def __init__(self, data, shuffle=False, graph=None):
        inputs = data[0]
        inputs, mask, len_max = self.data_masks(inputs, [0])
        self.inputs = np.asarray(inputs)
        self.mask = np.asarray(mask)
        self.len_max = len_max
        self.targets = np.asarray(data[1])
        self.length = len(inputs)
        self.shuffle = shuffle
        self.graph = graph

    def generate_batch(self, batch_size):
        if self.shuffle:
            shuffled_arg = np.arange(self.length)
            np.random.shuffle(shuffled_arg)
            self.inputs = self.inputs[shuffled_arg]
            self.mask = self.mask[shuffled_arg]
            self.targets = self.targets[shuffled_arg]
        n_batch = int(self.length / batch_size)
        if self.length % batch_size != 0:
            n_batch += 1
        slices = np.split(np.arange(n_batch * batch_size), n_batch)
        slices[-1] = slices[-1][:(self.length - batch_size * (n_batch - 1))]
        return slices

    def get_slice(self, i):
        inputs, mask, targets = self.inputs[i], self.mask[i], self.targets[i]
        items, n_node, A, alias_inputs = [], [], [], []
        for u_input in inputs:
            n_node.append(len(np.unique(u_input)))
        max_n_node = np.max(n_node)
        for u_input in inputs:
            node = np.unique(u_input)
            items.append(node.tolist() + (max_n_node - len(node)) * [0])
            u_A = np.zeros((max_n_node, max_n_node))
            for i in np.arange(len(u_input) - 1):
                if u_input[i + 1] == 0:
                    break
                u = np.where(node == u_input[i])[0][0]
                v = np.where(node == u_input[i + 1])[0][0]
                u_A[u][v] = 1
            u_sum_in = np.sum(u_A, 0)
            u_sum_in[np.where(u_sum_in == 0)] = 1
            u_A_in = np.divide(u_A, u_sum_in)
            u_sum_out = np.sum(u_A, 1)
            u_sum_out[np.where(u_sum_out == 0)] = 1
            u_A_out = np.divide(u_A.transpose(), u_sum_out)
            u_A = np.concatenate([u_A_in, u_A_out]).transpose()
            A.append(u_A)
            alias_inputs.append([np.where(node == i)[0][0] for i in u_input])
        return alias_inputs, A, items, mask, targets

    @staticmethod
    def data_masks(all_usr_pois, item_tail):
        us_lens = [len(upois) for upois in all_usr_pois]
        len_max = max(us_lens)
        us_pois = [upois + item_tail * (len_max - le) for upois, le in zip(all_usr_pois, us_lens)]
        us_msks = [[1] * le + [0] * (len_max - le) for le in us_lens]
        return us_pois, us_msks, len_max

# Cell
class SessionGraphDataset(Dataset):
    """
    References
        1. COTREC session-based recommender model training. https://t.ly/cXTH.
    """
    def __init__(self, root, shuffle=False, n_node=None):
        super().__init__(root)
        self.n_node = n_node
        self.shuffle = shuffle
        self.process()

    @property
    def raw_file_names(self) -> str:
        return ['data.txt', 'all_seq.txt']

    @property
    def processed_file_names(self) -> str:
        pass

    def download(self):
        raise NotImplementedError

    def process(self):
        import pickle
        data = pickle.load(open(self.raw_paths[0], 'rb'))
        all_seq = pickle.load(open(self.raw_paths[1], 'rb'))
        self.raw = np.asarray(data[0])
        self.targets = np.asarray(data[1])
        self.length = len(self.raw)
        adj = self.data_masks(all_seq, self.n_node)
        self.adjacency = adj.multiply(1.0/adj.sum(axis=0).reshape(1, -1))


    def get_overlap(self, sessions):
        matrix = np.zeros((len(sessions), len(sessions)))
        for i in range(len(sessions)):
            seq_a = set(sessions[i])
            seq_a.discard(0)
            for j in range(i+1, len(sessions)):
                seq_b = set(sessions[j])
                seq_b.discard(0)
                overlap = seq_a.intersection(seq_b)
                ab_set = seq_a | seq_b
                matrix[i][j] = float(len(overlap))/float(len(ab_set))
                matrix[j][i] = matrix[i][j]
        # matrix = self.dropout(matrix, 0.2)
        matrix = matrix + np.diag([1.0]*len(sessions))
        degree = np.sum(np.array(matrix), 1)
        degree = np.diag(1.0/degree)
        return matrix, degree

    def generate_batch(self, batch_size):
        if self.shuffle:
            shuffled_arg = np.arange(self.length)
            np.random.shuffle(shuffled_arg)
            self.raw = self.raw[shuffled_arg]
            self.targets = self.targets[shuffled_arg]
        n_batch = int(self.length / batch_size)
        if self.length % batch_size != 0:
            n_batch += 1
        slices = np.split(np.arange(n_batch * batch_size), n_batch)
        slices[-1] = np.arange(self.length-batch_size, self.length)
        return slices

    def get_slice(self, index):
        items, num_node = [], []
        inp = self.raw[index]
        for session in inp:
            num_node.append(len(np.nonzero(session)[0]))
        max_n_node = np.max(num_node)
        session_len = []
        reversed_sess_item = []
        mask = []
        # item_set = set()
        for session in inp:
            nonzero_elems = np.nonzero(session)[0]
            # item_set.update(set([t-1 for t in session]))
            session_len.append([len(nonzero_elems)])
            items.append(session + (max_n_node - len(nonzero_elems)) * [0])
            mask.append([1]*len(nonzero_elems) + (max_n_node - len(nonzero_elems)) * [0])
            reversed_sess_item.append(list(reversed(session)) + (max_n_node - len(nonzero_elems)) * [0])
        # item_set = list(item_set)
        # index_list = [item_set.index(a) for a in self.targets[index]-1]
        diff_mask = np.ones(shape=[100, self.n_node]) * (1/(self.n_node - 1))
        for count, value in enumerate(self.targets[index]-1):
            diff_mask[count][value] = 1
        return self.targets[index]-1, session_len,items, reversed_sess_item, mask, diff_mask

    @staticmethod
    def data_masks(all_sessions, n_node):
        adj = dict()
        for sess in all_sessions:
            for i, item in enumerate(sess):
                if i == len(sess)-1:
                    break
                else:
                    if sess[i] - 1 not in adj.keys():
                        adj[sess[i]-1] = dict()
                        adj[sess[i]-1][sess[i]-1] = 1
                        adj[sess[i]-1][sess[i+1]-1] = 1
                    else:
                        if sess[i+1]-1 not in adj[sess[i]-1].keys():
                            adj[sess[i] - 1][sess[i + 1] - 1] = 1
                        else:
                            adj[sess[i]-1][sess[i+1]-1] += 1
        row, col, data = [], [], []
        for i in adj.keys():
            item = adj[i]
            for j in item.keys():
                row.append(i)
                col.append(j)
                data.append(adj[i][j])
        from scipy.sparse import coo_matrix
        coo = coo_matrix((data, (row, col)), shape=(n_node, n_node))
        return coo

# Cell
class RatingDataset(Dataset):
    r"""Interaction data with rating feedback

    Args:
        root: data folder path
        min_uc: minimum user count to keep in the data
        min_sc: minimum item count to keep in the data
        split: data split method - leave_one_out/holdout
        min_rating: minimum rating threshold to convert explicit feedback into implicit

    References:
        1. https://github.com/Yueeeeeeee/RecSys-Extraction-Attack/tree/main/datasets
    """
    def __init__(self, root, min_uc, min_sc, split='leave_one_out', dataset_split_seed=42,
                 eval_set_size=None, min_rating=None, iterative_triplet=False):
        super().__init__(root)
        self.min_rating = min_rating
        self.min_uc = min_uc
        self.min_sc = min_sc
        self.split = split
        self.dataset_split_seed = dataset_split_seed
        self.eval_set_size = eval_set_size
        self.iterative_triplet = iterative_triplet

        assert self.min_uc >= 2, 'Need at least 2 ratings per user for validation and test'

        self._process()

    def load_ratings_df(self):
        r"""load raw dataset into pandas dataframe"""
        raise NotImplementedError

    def load_processed(self):
        return pickle.load(open(self.processed_paths[0], 'rb'))

    @property
    def processed_file_names(self):
        return 'dataset.pkl'

    def make_implicit(self, df):
        print('Turning into implicit ratings')
        df = df[df['rating'] >= self.min_rating]
        # return df[['uid', 'sid', 'timestamp']]
        return df

    def filter_triplets(self, df):
        print('Filtering triplets')
        if self.min_sc > 0:
            item_sizes = df.groupby('sid').size()
            good_items = item_sizes.index[item_sizes >= self.min_sc]
            df = df[df['sid'].isin(good_items)]

        if self.min_uc > 0:
            user_sizes = df.groupby('uid').size()
            good_users = user_sizes.index[user_sizes >= self.min_uc]
            df = df[df['uid'].isin(good_users)]

        return df

    def filter_triplets_iteratively(self, df):
        print('Filtering triplets')
        if self.min_sc > 0 or self.min_uc > 0:
            item_sizes = df.groupby('sid').size()
            good_items = item_sizes.index[item_sizes >= self.min_sc]
            user_sizes = df.groupby('uid').size()
            good_users = user_sizes.index[user_sizes >= self.min_uc]
            while len(good_items) < len(item_sizes) or len(good_users) < len(user_sizes):
                if self.min_sc > 0:
                    item_sizes = df.groupby('sid').size()
                    good_items = item_sizes.index[item_sizes >= self.min_sc]
                    df = df[df['sid'].isin(good_items)]

                if self.min_uc > 0:
                    user_sizes = df.groupby('uid').size()
                    good_users = user_sizes.index[user_sizes >= self.min_uc]
                    df = df[df['uid'].isin(good_users)]

                item_sizes = df.groupby('sid').size()
                good_items = item_sizes.index[item_sizes >= self.min_sc]
                user_sizes = df.groupby('uid').size()
                good_users = user_sizes.index[user_sizes >= self.min_uc]

        return df

    def densify_index(self, df):
        print('Densifying index')
        umap = {u: i for i, u in enumerate(set(df['uid']))}
        smap = {s: i for i, s in enumerate(set(df['sid']))}
        df['uid'] = df['uid'].map(umap)
        df['sid'] = df['sid'].map(smap)
        return df, umap, smap

    def split_df(self, df, user_count):
        if self.split == 'leave_one_out':
            print('Splitting')
            user_group = df.groupby('uid')
            user2items = user_group.progress_apply(lambda d: list(d.sort_values(by='timestamp')['sid']))
            train, val, test = {}, {}, {}
            for user in range(user_count):
                items = user2items[user]
                train[user], val[user], test[user] = items[:-2], items[-2:-1], items[-1:]
            return train, val, test
        elif self.split == 'holdout':
            print('Splitting')
            np.random.seed(self.dataset_split_seed)
            eval_set_size = self.eval_set_size

            # Generate user indices
            permuted_index = np.random.permutation(user_count)
            train_user_index = permuted_index[                :-2*eval_set_size]
            val_user_index   = permuted_index[-2*eval_set_size:  -eval_set_size]
            test_user_index  = permuted_index[  -eval_set_size:                ]

            # Split DataFrames
            train_df = df.loc[df['uid'].isin(train_user_index)]
            val_df   = df.loc[df['uid'].isin(val_user_index)]
            test_df  = df.loc[df['uid'].isin(test_user_index)]

            # DataFrame to dict => {uid : list of sid's}
            train = dict(train_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))
            val   = dict(val_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))
            test  = dict(test_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))
            return train, val, test
        else:
            raise NotImplementedError

    def process(self):
        df = self.load_ratings_df()
        if self.iterative_triplet:
            df = self.filter_triplets_iteratively(df)
        else:
            df = self.filter_triplets(df)
        df, umap, smap = self.densify_index(df)
        train, val, test = self.split_df(df, len(umap))
        dataset = {'train': train,
                   'val': val,
                   'test': test,
                   'umap': umap,
                   'smap': smap}
        with open(self.processed_paths[0], 'wb') as f:
            pickle.dump(dataset, f)

# Cell
class TorchSessionDataset(torch.utils.data.Dataset):
    def __init__(self, fpath, maxlen, is_train=True):
            [train, valid, test, itemnum] = self.data_partition(fpath)
            print("Number of sessions:",len(train)+len(valid)+len(test))
            print("Number of items:", itemnum)

            action = 0
            for i in train:
                action += np.count_nonzero(i)
            for i in valid:
                action += np.count_nonzero(i)
            for i in test:
                action += np.count_nonzero(i)

            print("Number of actions:", action)
            print("Average length of sessions:", action/(len(train)+len(valid)+len(test)))

            self.data = train if is_train else test
            self.maxlen = maxlen
            self.itemnum = itemnum
            self.is_train = is_train

    def __len__(self):
            return len(self.data)

    def __train__(self, index):
            session = np.asarray(self.data[index], dtype=np.int64)
            if len(session) > self.maxlen:
                session = session[-self.maxlen:]
            else:
                session = np.pad(session, (self.maxlen-len(session), 0), 'constant', constant_values=0)
            curr_seq = session[:-1]
            curr_pos = session[1:]
            return curr_seq, curr_pos

    def __test__(self, index):
            session = self.data[index]
            seq = np.zeros([self.maxlen], dtype=np.int64)
            idx = self.maxlen - 1
            for i in reversed(session[:-1]): #everything except the last one
                seq[idx] = i
                idx -= 1
                if idx == -1: break
            return seq, session[-1]-1 #index of the item in the list of all items

    def __getitem__(self, index):
            if self.is_train:
                return self.__train__(index)
            else:
                return self.__test__(index)

    @staticmethod
    def data_partition(fname, percentage=[0.1, 0.2]):
        itemnum = 0

        sessions = defaultdict(list)
        session_train = []
        session_valid = []
        session_test = []
        # assume user/item index starting from 1
        session_id = 0
        f = open(fname, 'r')
        total_length = 0
        max_length = 0
        for line in f:

            items = [int(l) for l in line.rstrip().split(',')]

            if len(items) < 5: continue
            total_length += len(items)

            if max_length< len(items):
                max_length = len(items)

            itemnum = max(max(items), itemnum)
            sessions[session_id].append(items)
            session_id += 1

        print("Avg length:", total_length/session_id)
        print("Maximum length:", max_length)

        valid_perc = percentage[0]
        test_perc = percentage[1]

        total_sessions = session_id

        shuffle_indices = np.random.permutation(range(total_sessions)) #

        train_index = int(total_sessions*(1 - valid_perc - test_perc))
        valid_index = int(total_sessions*(1 - test_perc))

        if (train_index == valid_index): valid_index += 1 #break the tie

        train_indices = shuffle_indices[:train_index]
        valid_indices = shuffle_indices[train_index:valid_index]
        test_indices = shuffle_indices[valid_index:]

        for i in train_indices:
            session_train.extend(sessions[i])
        for i in valid_indices:
            session_valid.extend(sessions[i])
        for i in test_indices:
            session_test.extend(sessions[i])

        return [np.asarray(session_train), np.asarray(session_valid), np.asarray(session_test), itemnum]

    @staticmethod
    def nextitnet_format(fname, maxlen):

        sessions = []

        # assume user/item index starting from 1
        f = open(fname, 'r')

        for line in f:

            items = [int(l) for l in line.rstrip().split(',')]

            if len(items) < 5: continue

            seq = np.zeros([maxlen], dtype=np.int32)

            idx = maxlen - 1

            for i in reversed(items):
                seq[idx] = i
                idx -= 1
                if idx == -1: break

            sessions.append(seq)

        print("number of session:", len(sessions))

        return sessions

    @staticmethod
    def gru_format(fname, user_train, user_valid, user_test):

        session_id = 0
        train = []
        for session in user_train:
            for item in session:
                train.append([session_id, item, 0])
            session_id += 1

        valid = []
        for session in user_valid:
            for item in session:
                valid.append([session_id, item, 0])
            session_id += 1

        test = []
        for session in user_test:
            for item in session:
                test.append([session_id, item, 0])
            session_id += 1

        train_data = pd.DataFrame(train, columns= ['SessionId', 'ItemId', 'Time'])
        valid_data = pd.DataFrame(valid, columns= ['SessionId', 'ItemId', 'Time'])
        test_data = pd.DataFrame(test, columns= ['SessionId', 'ItemId', 'Time'])

        return train_data, valid_data, test_data

# Cell
class TorchRatingDataset(torch.utils.data.Dataset):
    """Wrapper, convert <user, item, rating> Tensor into Pytorch Dataset."""

    def __init__(self, user_tensor, item_tensor, target_tensor):
        """Init UserItemRatingDataset Class.
        Args:
            target_tensor: torch.Tensor, the corresponding rating for <user, item> pair.
        """
        self.user_tensor = user_tensor
        self.item_tensor = item_tensor
        self.target_tensor = target_tensor

    def __getitem__(self, index):
        """Get an item from dataset."""
        return (
            self.user_tensor[index],
            self.item_tensor[index],
            self.target_tensor[index],
        )

    def __len__(self):
        """Get the size of the dataset."""
        return self.user_tensor.size(0)

# Cell
class TorchPairwiseNegativeDataset(torch.utils.data.Dataset):
    """Wrapper, convert <user, pos_item, neg_item> Tensor into Pytorch Dataset."""

    def __init__(self, user_tensor, pos_item_tensor, neg_item_tensor):
        """Init PairwiseNegativeDataset Class.
        Args:
            target_tensor: torch.Tensor, the corresponding rating for <user, item> pair.
        """
        self.user_tensor = user_tensor
        self.pos_item_tensor = pos_item_tensor
        self.neg_item_tensor = neg_item_tensor

    def __getitem__(self, index):
        """Get an item from the dataset."""
        return (
            self.user_tensor[index],
            self.pos_item_tensor[index],
            self.neg_item_tensor[index],
        )

    def __len__(self):
        """Get the size of the dataset."""
        return self.user_tensor.size(0)