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

__all__ = ['AbstractDataset', 'AbstractDatasetv2', 'SessionDataset', 'SessionDatasetv2', 'SessionDatasetv3',
           'GraphData']

# 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 datetime import date, timezone, datetime, timedelta
from pandas import Timedelta

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

# Internal Cell
def download(url, savepath):
    import wget
    wget.download(url, str(savepath))


def unzip(zippath, savepath):
    import zipfile
    zip = zipfile.ZipFile(zippath)
    zip.extractall(savepath)
    zip.close()

# Cell
class AbstractDataset(metaclass=ABCMeta):
    def __init__(self, args):
        self.min_rating = args.min_rating
        self.min_uc = args.min_uc
        self.min_sc = args.min_sc
        self.split = args.split
        self.RAW_DATASET_ROOT_FOLDER = args.RAW_DATASET_ROOT_FOLDER
        self.PREP_DATASET_ROOT_FOLDER = args.PREP_DATASET_ROOT_FOLDER

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

    @classmethod
    @abstractmethod
    def code(cls):
        pass

    @classmethod
    def raw_code(cls):
        return cls.code()

    @abstractmethod
    def preprocess(self):
        pass

    @abstractmethod
    def load_ratings_df(self):
        pass

    def load_dataset(self):
        self.preprocess()
        dataset_path = self._get_preprocessed_dataset_path()
        dataset = pickle.load(dataset_path.open('rb'))
        return dataset

    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 densify_index(self, df):
        print('Densifying index')
        umap = {u: i for i, u in enumerate(set(df['uid']), start=1)}
        smap = {s: i for i, s in enumerate(set(df['sid']), start=1)}
        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'])['sid']))
            train, val, test = {}, {}, {}
            for i in range(user_count):
                user = i + 1
                items = user2items[user]
                train[user], val[user], test[user] = items[:-2], items[-2:-1], items[-1:]
            return train, val, test
        else:
            raise NotImplementedError

    def _get_rawdata_root_path(self):
        return Path(self.RAW_DATASET_ROOT_FOLDER)

    def _get_rawdata_folder_path(self):
        root = self._get_rawdata_root_path()
        return root.joinpath(self.raw_code())

    def _get_preprocessed_root_path(self):
        root = Path(self.PREP_DATASET_ROOT_FOLDER)
        return root.joinpath(self.raw_code())

    def _get_preprocessed_folder_path(self):
        preprocessed_root = self._get_preprocessed_root_path()
        # folder_name = '{}_min_rating{}-min_uc{}-min_sc{}-split{}' \
        #     .format(self.code(), self.min_rating, self.min_uc, self.min_sc, self.split)
        # return preprocessed_root.joinpath(folder_name)
        return preprocessed_root

    def _get_preprocessed_dataset_path(self):
        folder = self._get_preprocessed_folder_path()
        return folder.joinpath('dataset.pkl')

# Cell
class AbstractDatasetv2(metaclass=ABCMeta):
    def __init__(self, args):
        self.args = args
        self.min_rating = args.min_rating
        self.min_uc = args.min_uc
        self.min_sc = args.min_sc
        self.split = args.split

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

    @classmethod
    @abstractmethod
    def code(cls):
        pass

    @classmethod
    def raw_code(cls):
        return cls.code()

    @classmethod
    @abstractmethod
    def url(cls):
        pass

    @classmethod
    def is_zipfile(cls):
        return True

    @classmethod
    def zip_file_content_is_folder(cls):
        return True

    @classmethod
    def all_raw_file_names(cls):
        return []

    @abstractmethod
    def load_ratings_df(self):
        pass

    def load_dataset(self):
        self.preprocess()
        dataset_path = self._get_preprocessed_dataset_path()
        dataset = pickle.load(dataset_path.open('rb'))
        return dataset

    def preprocess(self):
        dataset_path = self._get_preprocessed_dataset_path()
        if dataset_path.is_file():
            print('Already preprocessed. Skip preprocessing')
            return
        if not dataset_path.parent.is_dir():
            dataset_path.parent.mkdir(parents=True)
        self.maybe_download_raw_dataset()
        df = self.load_ratings_df()
        df = self.make_implicit(df)
        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 dataset_path.open('wb') as f:
            pickle.dump(dataset, f)

    def maybe_download_raw_dataset(self):
        folder_path = self._get_rawdata_folder_path()
        if folder_path.is_dir() and\
           all(folder_path.joinpath(filename).is_file() for filename in self.all_raw_file_names()):
            print('Raw data already exists. Skip downloading')
            return
        print("Raw file doesn't exist. Downloading...")
        if self.is_zipfile():
            tmproot = Path(tempfile.mkdtemp())
            tmpzip = tmproot.joinpath('file.zip')
            tmpfolder = tmproot.joinpath('folder')
            download(self.url(), tmpzip)
            unzip(tmpzip, tmpfolder)
            if self.zip_file_content_is_folder():
                tmpfolder = tmpfolder.joinpath(os.listdir(tmpfolder)[0])
            shutil.move(tmpfolder, folder_path)
            shutil.rmtree(tmproot)
            print()
        else:
            tmproot = Path(tempfile.mkdtemp())
            tmpfile = tmproot.joinpath('file')
            download(self.url(), tmpfile)
            folder_path.mkdir(parents=True)
            shutil.move(tmpfile, folder_path.joinpath('ratings.csv'))
            shutil.rmtree(tmproot)
            print()

    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 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.args.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.args.split == 'holdout':
            print('Splitting')
            np.random.seed(self.args.dataset_split_seed)
            eval_set_size = self.args.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 _get_rawdata_root_path(self):
        return Path(self.args.RAW_DATASET_ROOT_FOLDER)

    def _get_rawdata_folder_path(self):
        root = self._get_rawdata_root_path()
        return root.joinpath(self.raw_code())

    def _get_preprocessed_root_path(self):
        root = self._get_rawdata_root_path()
        return root.joinpath('preprocessed')

    def _get_preprocessed_folder_path(self):
        preprocessed_root = self._get_preprocessed_root_path()
        folder_name = '{}_min_rating{}-min_uc{}-min_sc{}-split{}' \
            .format(self.code(), self.min_rating, self.min_uc, self.min_sc, self.split)
        return preprocessed_root.joinpath(folder_name)

    def _get_preprocessed_dataset_path(self):
        folder = self._get_preprocessed_folder_path()
        return folder.joinpath('dataset.pkl')

# 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 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.fromtimestamp(data.Time.min(), timezone.utc)
        data_end = 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.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.timestamp()].index

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

        #output
        data_start = datetime.fromtimestamp(data.Time.min(), timezone.utc)
        data_end = 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.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.fromtimestamp(data.Time.min(), timezone.utc)
        data_end = 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.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 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)

# Internal Cell
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 GraphData():
    def __init__(self, data, shuffle=False, graph=None):
        inputs = data[0]
        inputs, mask, len_max = 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