{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp datasets.base"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Base dataset\n",
    "> Base class for dataset module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import List, Optional, Callable, Union, Any, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "tqdm.pandas()\n",
    "\n",
    "from abc import *\n",
    "from pathlib import Path\n",
    "import os\n",
    "import os.path as osp\n",
    "from collections.abc import Sequence\n",
    "import sys\n",
    "import tempfile\n",
    "import shutil\n",
    "import pickle\n",
    "import time\n",
    "import csv\n",
    "import math\n",
    "import operator\n",
    "import itertools\n",
    "from datetime import date, timezone, datetime, timedelta\n",
    "from pandas import Timedelta\n",
    "\n",
    "from recohut.utils.common_utils import download_url, extract_zip, makedirs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def download(url, savepath):\n",
    "    import wget\n",
    "    wget.download(url, str(savepath))\n",
    "\n",
    "\n",
    "def unzip(zippath, savepath):\n",
    "    import zipfile\n",
    "    zip = zipfile.ZipFile(zippath)\n",
    "    zip.extractall(savepath)\n",
    "    zip.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AbstractDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AbstractDataset(metaclass=ABCMeta):\n",
    "    def __init__(self, args):\n",
    "        self.min_rating = args.min_rating\n",
    "        self.min_uc = args.min_uc\n",
    "        self.min_sc = args.min_sc\n",
    "        self.split = args.split\n",
    "        self.RAW_DATASET_ROOT_FOLDER = args.RAW_DATASET_ROOT_FOLDER\n",
    "        self.PREP_DATASET_ROOT_FOLDER = args.PREP_DATASET_ROOT_FOLDER\n",
    "\n",
    "        assert self.min_uc >= 2, 'Need at least 2 ratings per user for validation and test'\n",
    "\n",
    "    @classmethod\n",
    "    @abstractmethod\n",
    "    def code(cls):\n",
    "        pass\n",
    "\n",
    "    @classmethod\n",
    "    def raw_code(cls):\n",
    "        return cls.code()\n",
    "\n",
    "    @abstractmethod\n",
    "    def preprocess(self):\n",
    "        pass\n",
    "\n",
    "    @abstractmethod\n",
    "    def load_ratings_df(self):\n",
    "        pass\n",
    "\n",
    "    def load_dataset(self):\n",
    "        self.preprocess()\n",
    "        dataset_path = self._get_preprocessed_dataset_path()\n",
    "        dataset = pickle.load(dataset_path.open('rb'))\n",
    "        return dataset\n",
    "\n",
    "    def filter_triplets(self, df):\n",
    "        print('Filtering triplets')\n",
    "        if self.min_sc > 0:\n",
    "            item_sizes = df.groupby('sid').size()\n",
    "            good_items = item_sizes.index[item_sizes >= self.min_sc]\n",
    "            df = df[df['sid'].isin(good_items)]\n",
    "\n",
    "        if self.min_uc > 0:\n",
    "            user_sizes = df.groupby('uid').size()\n",
    "            good_users = user_sizes.index[user_sizes >= self.min_uc]\n",
    "            df = df[df['uid'].isin(good_users)]\n",
    "        return df\n",
    "\n",
    "    def densify_index(self, df):\n",
    "        print('Densifying index')\n",
    "        umap = {u: i for i, u in enumerate(set(df['uid']), start=1)}\n",
    "        smap = {s: i for i, s in enumerate(set(df['sid']), start=1)}\n",
    "        df['uid'] = df['uid'].map(umap)\n",
    "        df['sid'] = df['sid'].map(smap)\n",
    "        return df, umap, smap\n",
    "\n",
    "    def split_df(self, df, user_count):\n",
    "        if self.split == 'leave_one_out':\n",
    "            print('Splitting')\n",
    "            user_group = df.groupby('uid')\n",
    "            user2items = user_group.progress_apply(\n",
    "                lambda d: list(d.sort_values(by=['timestamp', 'sid'])['sid']))\n",
    "            train, val, test = {}, {}, {}\n",
    "            for i in range(user_count):\n",
    "                user = i + 1\n",
    "                items = user2items[user]\n",
    "                train[user], val[user], test[user] = items[:-2], items[-2:-1], items[-1:]\n",
    "            return train, val, test\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "    def _get_rawdata_root_path(self):\n",
    "        return Path(self.RAW_DATASET_ROOT_FOLDER)\n",
    "\n",
    "    def _get_rawdata_folder_path(self):\n",
    "        root = self._get_rawdata_root_path()\n",
    "        return root.joinpath(self.raw_code())\n",
    "\n",
    "    def _get_preprocessed_root_path(self):\n",
    "        root = Path(self.PREP_DATASET_ROOT_FOLDER)\n",
    "        return root.joinpath(self.raw_code())\n",
    "\n",
    "    def _get_preprocessed_folder_path(self):\n",
    "        preprocessed_root = self._get_preprocessed_root_path()\n",
    "        # folder_name = '{}_min_rating{}-min_uc{}-min_sc{}-split{}' \\\n",
    "        #     .format(self.code(), self.min_rating, self.min_uc, self.min_sc, self.split)\n",
    "        # return preprocessed_root.joinpath(folder_name)\n",
    "        return preprocessed_root\n",
    "\n",
    "    def _get_preprocessed_dataset_path(self):\n",
    "        folder = self._get_preprocessed_folder_path()\n",
    "        return folder.joinpath('dataset.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AbstractDatasetv2(metaclass=ABCMeta):\n",
    "    def __init__(self, args):\n",
    "        self.args = args\n",
    "        self.min_rating = args.min_rating\n",
    "        self.min_uc = args.min_uc\n",
    "        self.min_sc = args.min_sc\n",
    "        self.split = args.split\n",
    "\n",
    "        assert self.min_uc >= 2, 'Need at least 2 ratings per user for validation and test'\n",
    "\n",
    "    @classmethod\n",
    "    @abstractmethod\n",
    "    def code(cls):\n",
    "        pass\n",
    "\n",
    "    @classmethod\n",
    "    def raw_code(cls):\n",
    "        return cls.code()\n",
    "\n",
    "    @classmethod\n",
    "    @abstractmethod\n",
    "    def url(cls):\n",
    "        pass\n",
    "\n",
    "    @classmethod\n",
    "    def is_zipfile(cls):\n",
    "        return True\n",
    "\n",
    "    @classmethod\n",
    "    def zip_file_content_is_folder(cls):\n",
    "        return True\n",
    "\n",
    "    @classmethod\n",
    "    def all_raw_file_names(cls):\n",
    "        return []\n",
    "\n",
    "    @abstractmethod\n",
    "    def load_ratings_df(self):\n",
    "        pass\n",
    "\n",
    "    def load_dataset(self):\n",
    "        self.preprocess()\n",
    "        dataset_path = self._get_preprocessed_dataset_path()\n",
    "        dataset = pickle.load(dataset_path.open('rb'))\n",
    "        return dataset\n",
    "\n",
    "    def preprocess(self):\n",
    "        dataset_path = self._get_preprocessed_dataset_path()\n",
    "        if dataset_path.is_file():\n",
    "            print('Already preprocessed. Skip preprocessing')\n",
    "            return\n",
    "        if not dataset_path.parent.is_dir():\n",
    "            dataset_path.parent.mkdir(parents=True)\n",
    "        self.maybe_download_raw_dataset()\n",
    "        df = self.load_ratings_df()\n",
    "        df = self.make_implicit(df)\n",
    "        df = self.filter_triplets(df)\n",
    "        df, umap, smap = self.densify_index(df)\n",
    "        train, val, test = self.split_df(df, len(umap))\n",
    "        dataset = {'train': train,\n",
    "                   'val': val,\n",
    "                   'test': test,\n",
    "                   'umap': umap,\n",
    "                   'smap': smap}\n",
    "        with dataset_path.open('wb') as f:\n",
    "            pickle.dump(dataset, f)\n",
    "\n",
    "    def maybe_download_raw_dataset(self):\n",
    "        folder_path = self._get_rawdata_folder_path()\n",
    "        if folder_path.is_dir() and\\\n",
    "           all(folder_path.joinpath(filename).is_file() for filename in self.all_raw_file_names()):\n",
    "            print('Raw data already exists. Skip downloading')\n",
    "            return\n",
    "        print(\"Raw file doesn't exist. Downloading...\")\n",
    "        if self.is_zipfile():\n",
    "            tmproot = Path(tempfile.mkdtemp())\n",
    "            tmpzip = tmproot.joinpath('file.zip')\n",
    "            tmpfolder = tmproot.joinpath('folder')\n",
    "            download(self.url(), tmpzip)\n",
    "            unzip(tmpzip, tmpfolder)\n",
    "            if self.zip_file_content_is_folder():\n",
    "                tmpfolder = tmpfolder.joinpath(os.listdir(tmpfolder)[0])\n",
    "            shutil.move(tmpfolder, folder_path)\n",
    "            shutil.rmtree(tmproot)\n",
    "            print()\n",
    "        else:\n",
    "            tmproot = Path(tempfile.mkdtemp())\n",
    "            tmpfile = tmproot.joinpath('file')\n",
    "            download(self.url(), tmpfile)\n",
    "            folder_path.mkdir(parents=True)\n",
    "            shutil.move(tmpfile, folder_path.joinpath('ratings.csv'))\n",
    "            shutil.rmtree(tmproot)\n",
    "            print()\n",
    "\n",
    "    def make_implicit(self, df):\n",
    "        print('Turning into implicit ratings')\n",
    "        df = df[df['rating'] >= self.min_rating]\n",
    "        # return df[['uid', 'sid', 'timestamp']]\n",
    "        return df\n",
    "\n",
    "    def filter_triplets(self, df):\n",
    "        print('Filtering triplets')\n",
    "        if self.min_sc > 0:\n",
    "            item_sizes = df.groupby('sid').size()\n",
    "            good_items = item_sizes.index[item_sizes >= self.min_sc]\n",
    "            df = df[df['sid'].isin(good_items)]\n",
    "\n",
    "        if self.min_uc > 0:\n",
    "            user_sizes = df.groupby('uid').size()\n",
    "            good_users = user_sizes.index[user_sizes >= self.min_uc]\n",
    "            df = df[df['uid'].isin(good_users)]\n",
    "\n",
    "        return df\n",
    "\n",
    "    def densify_index(self, df):\n",
    "        print('Densifying index')\n",
    "        umap = {u: i for i, u in enumerate(set(df['uid']))}\n",
    "        smap = {s: i for i, s in enumerate(set(df['sid']))}\n",
    "        df['uid'] = df['uid'].map(umap)\n",
    "        df['sid'] = df['sid'].map(smap)\n",
    "        return df, umap, smap\n",
    "\n",
    "    def split_df(self, df, user_count):\n",
    "        if self.args.split == 'leave_one_out':\n",
    "            print('Splitting')\n",
    "            user_group = df.groupby('uid')\n",
    "            user2items = user_group.progress_apply(lambda d: list(d.sort_values(by='timestamp')['sid']))\n",
    "            train, val, test = {}, {}, {}\n",
    "            for user in range(user_count):\n",
    "                items = user2items[user]\n",
    "                train[user], val[user], test[user] = items[:-2], items[-2:-1], items[-1:]\n",
    "            return train, val, test\n",
    "        elif self.args.split == 'holdout':\n",
    "            print('Splitting')\n",
    "            np.random.seed(self.args.dataset_split_seed)\n",
    "            eval_set_size = self.args.eval_set_size\n",
    "\n",
    "            # Generate user indices\n",
    "            permuted_index = np.random.permutation(user_count)\n",
    "            train_user_index = permuted_index[                :-2*eval_set_size]\n",
    "            val_user_index   = permuted_index[-2*eval_set_size:  -eval_set_size]\n",
    "            test_user_index  = permuted_index[  -eval_set_size:                ]\n",
    "\n",
    "            # Split DataFrames\n",
    "            train_df = df.loc[df['uid'].isin(train_user_index)]\n",
    "            val_df   = df.loc[df['uid'].isin(val_user_index)]\n",
    "            test_df  = df.loc[df['uid'].isin(test_user_index)]\n",
    "\n",
    "            # DataFrame to dict => {uid : list of sid's}\n",
    "            train = dict(train_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))\n",
    "            val   = dict(val_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))\n",
    "            test  = dict(test_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))\n",
    "            return train, val, test\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "    def _get_rawdata_root_path(self):\n",
    "        return Path(self.args.RAW_DATASET_ROOT_FOLDER)\n",
    "\n",
    "    def _get_rawdata_folder_path(self):\n",
    "        root = self._get_rawdata_root_path()\n",
    "        return root.joinpath(self.raw_code())\n",
    "\n",
    "    def _get_preprocessed_root_path(self):\n",
    "        root = self._get_rawdata_root_path()\n",
    "        return root.joinpath('preprocessed')\n",
    "\n",
    "    def _get_preprocessed_folder_path(self):\n",
    "        preprocessed_root = self._get_preprocessed_root_path()\n",
    "        folder_name = '{}_min_rating{}-min_uc{}-min_sc{}-split{}' \\\n",
    "            .format(self.code(), self.min_rating, self.min_uc, self.min_sc, self.split)\n",
    "        return preprocessed_root.joinpath(folder_name)\n",
    "\n",
    "    def _get_preprocessed_dataset_path(self):\n",
    "        folder = self._get_preprocessed_folder_path()\n",
    "        return folder.joinpath('dataset.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def to_list(value: Any) -> Sequence:\n",
    "    if isinstance(value, Sequence) and not isinstance(value, str):\n",
    "        return value\n",
    "    else:\n",
    "        return [value]\n",
    "\n",
    "def files_exist(files: List[str]) -> bool:\n",
    "    # NOTE: We return `False` in case `files` is empty, leading to a\n",
    "    # re-processing of files on every instantiation.\n",
    "    return len(files) != 0 and all([osp.exists(f) for f in files])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class Dataset:\n",
    "    \"\"\"Dataset base class\n",
    "    \"\"\"\n",
    "    @property\n",
    "    def raw_file_names(self) -> Union[str, List[str], Tuple]:\n",
    "        r\"\"\"The name of the files in the :obj:`self.raw_dir` folder that must\n",
    "        be present in order to skip downloading.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> Union[str, List[str], Tuple]:\n",
    "        r\"\"\"The name of the files in the :obj:`self.processed_dir` folder that\n",
    "        must be present in order to skip processing.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def download(self):\n",
    "        r\"\"\"Downloads the dataset to the :obj:`self.raw_dir` folder.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def process(self):\n",
    "        r\"\"\"Processes the dataset to the :obj:`self.processed_dir` folder.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def __init__(self, root=None):\n",
    "        self.root = root\n",
    "\n",
    "        if 'download' in self.__class__.__dict__:\n",
    "            self._download()\n",
    "\n",
    "        # if 'process' in self.__class__.__dict__:\n",
    "        #     self._process()\n",
    "\n",
    "    @property\n",
    "    def raw_dir(self) -> str:\n",
    "        return osp.join(self.root, 'raw')\n",
    "\n",
    "    @property\n",
    "    def processed_dir(self) -> str:\n",
    "        return osp.join(self.root, 'processed')\n",
    "\n",
    "    @property\n",
    "    def raw_paths(self) -> List[str]:\n",
    "        r\"\"\"The absolute filepaths that must be present in order to skip\n",
    "        downloading.\"\"\"\n",
    "        files = to_list(self.raw_file_names)\n",
    "        return [osp.join(self.raw_dir, f) for f in files]\n",
    "\n",
    "    @property\n",
    "    def processed_paths(self) -> List[str]:\n",
    "        r\"\"\"The absolute filepaths that must be present in order to skip\n",
    "        processing.\"\"\"\n",
    "        files = to_list(self.processed_file_names)\n",
    "        return [osp.join(self.processed_dir, f) for f in files]\n",
    "\n",
    "    def _download(self):\n",
    "        if files_exist(self.raw_paths):  # pragma: no cover\n",
    "            return\n",
    "\n",
    "        makedirs(self.raw_dir)\n",
    "        self.download()\n",
    "\n",
    "    def _process(self):\n",
    "        if files_exist(self.processed_paths):  # pragma: no cover\n",
    "            return\n",
    "\n",
    "        print('Processing...', file=sys.stderr)\n",
    "\n",
    "        makedirs(self.processed_dir)\n",
    "        self.process()\n",
    "\n",
    "        print('Done!', file=sys.stderr)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        arg_repr = str(len(self)) if len(self) > 1 else ''\n",
    "        return f'{self.__class__.__name__}({arg_repr})'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SessionDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDataset(Dataset):\n",
    "    r\"\"\"Session dataset base class.\n",
    "\n",
    "    Args:\n",
    "        root (string): Root directory where the dataset should be saved.\n",
    "        process_method (string):\n",
    "            last: last day => test set\n",
    "            last_min_date: last day => test set, but from a minimal date onwards\n",
    "            days_test: last N days => test set\n",
    "            slice: create multiple train-test-combinations with a sliding window approach\n",
    "        min_date (string): Minimum date\n",
    "        session_length (int): Session time length :default = 30 * 60 #30 minutes\n",
    "        min_session_length (int): Minimum number of items for a session to be valid\n",
    "        min_item_support (int): Minimum number of interactions for an item to be valid\n",
    "        num_slices (int): Offset in days from the first date in the data set\n",
    "        days_offset (int): Number of days the training start date is shifted after creating one slice\n",
    "        days_shift (int): Days shift\n",
    "        days_train (int): Days in train set in each slice\n",
    "        days_test (int): Days in test set in each slice\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, root, process_method, min_date=None,\n",
    "                 session_length=None, min_session_length=None, min_item_support=None,\n",
    "                 num_slices=None, days_offset=None, days_shift=None, days_train=None,\n",
    "                 days_test=None, data=None):\n",
    "        super().__init__(root)\n",
    "        self.process_method = process_method\n",
    "        self.min_date = min_date\n",
    "        self.session_length = session_length\n",
    "        self.min_session_length = min_session_length\n",
    "        self.min_item_support = min_item_support\n",
    "        self.num_slices = num_slices\n",
    "        self.days_offset = days_offset\n",
    "        self.days_shift = days_shift\n",
    "        self.days_train = days_train\n",
    "        self.days_test = days_test\n",
    "        self.data = None\n",
    "\n",
    "        self._process()\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def load(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def filter_data(self): \n",
    "        data = self.data\n",
    "\n",
    "        #filter session length\n",
    "        session_lengths = data.groupby('SessionId').size()\n",
    "        data = data[np.in1d(data.SessionId, session_lengths[session_lengths>1].index)]\n",
    "        \n",
    "        #filter item support\n",
    "        item_supports = data.groupby('ItemId').size()\n",
    "        data = data[np.in1d(data.ItemId, item_supports[item_supports>= self.min_item_support].index)]\n",
    "        \n",
    "        #filter session length\n",
    "        session_lengths = data.groupby('SessionId').size()\n",
    "        data = data[np.in1d(data.SessionId, session_lengths[session_lengths>= self.min_session_length].index)]\n",
    "        \n",
    "        #output\n",
    "        data_start = datetime.fromtimestamp(data.Time.min(), timezone.utc)\n",
    "        data_end = datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        \n",
    "        print('Filtered data set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}\\n\\n'.\n",
    "              format(len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.date().isoformat(), data_end.date().isoformat()))\n",
    "    \n",
    "        self.data = data\n",
    "        \n",
    "    def filter_min_date(self):\n",
    "        data = self.data\n",
    "\n",
    "        min_datetime = datetime.strptime(self.min_date + ' 00:00:00', '%Y-%m-%d %H:%M:%S')\n",
    "        \n",
    "        #filter\n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        session_keep = session_max_times[session_max_times > min_datetime.timestamp()].index\n",
    "        \n",
    "        data = data[np.in1d(data.SessionId, session_keep)]\n",
    "        \n",
    "        #output\n",
    "        data_start = datetime.fromtimestamp(data.Time.min(), timezone.utc)\n",
    "        data_end = datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        \n",
    "        print('Filtered data set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}\\n\\n'.\n",
    "              format(len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.date().isoformat(), data_end.date().isoformat()))\n",
    "        \n",
    "        self.data = data\n",
    "\n",
    "    def split_data_org(self):\n",
    "        data = self.data\n",
    "        tmax = data.Time.max()\n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        session_train = session_max_times[session_max_times < tmax-86400].index\n",
    "        session_test = session_max_times[session_max_times >= tmax-86400].index\n",
    "        train = data[np.in1d(data.SessionId, session_train)]\n",
    "        test = data[np.in1d(data.SessionId, session_test)]\n",
    "        test = test[np.in1d(test.ItemId, train.ItemId)]\n",
    "        tslength = test.groupby('SessionId').size()\n",
    "        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]\n",
    "        print('Full train set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(train), train.SessionId.nunique(), train.ItemId.nunique()))\n",
    "        train.to_csv(osp.join(self.processed_dir,'events_train_full.txt'), sep='\\t', index=False)\n",
    "        print('Test set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(test), test.SessionId.nunique(), test.ItemId.nunique()))\n",
    "        test.to_csv(osp.join(self.processed_dir,'events_test.txt'), sep='\\t', index=False)\n",
    "        \n",
    "        tmax = train.Time.max()\n",
    "        session_max_times = train.groupby('SessionId').Time.max()\n",
    "        session_train = session_max_times[session_max_times < tmax-86400].index\n",
    "        session_valid = session_max_times[session_max_times >= tmax-86400].index\n",
    "        train_tr = train[np.in1d(train.SessionId, session_train)]\n",
    "        valid = train[np.in1d(train.SessionId, session_valid)]\n",
    "        valid = valid[np.in1d(valid.ItemId, train_tr.ItemId)]\n",
    "        tslength = valid.groupby('SessionId').size()\n",
    "        valid = valid[np.in1d(valid.SessionId, tslength[tslength>=2].index)]\n",
    "        print('Train set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(train_tr), train_tr.SessionId.nunique(), train_tr.ItemId.nunique()))\n",
    "        train_tr.to_csv(osp.join(self.processed_dir,'events_train_tr.txt'), sep='\\t', index=False)\n",
    "        print('Validation set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(valid), valid.SessionId.nunique(), valid.ItemId.nunique()))\n",
    "        valid.to_csv(osp.join(self.processed_dir,'events_train_valid.txt'), sep='\\t', index=False)\n",
    "\n",
    "    def split_data(self):\n",
    "        data = self.data\n",
    "        data_end = datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        test_from = data_end - timedelta(self.days_test)\n",
    "        \n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        session_train = session_max_times[session_max_times < test_from.timestamp()].index\n",
    "        session_test = session_max_times[session_max_times >= test_from.timestamp()].index\n",
    "        train = data[np.in1d(data.SessionId, session_train)]\n",
    "        test = data[np.in1d(data.SessionId, session_test)]\n",
    "        test = test[np.in1d(test.ItemId, train.ItemId)]\n",
    "        tslength = test.groupby('SessionId').size()\n",
    "        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]\n",
    "        print('Full train set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(train), train.SessionId.nunique(), train.ItemId.nunique()))\n",
    "        train.to_csv(osp.join(self.processed_dir,'events_train_full.txt'), sep='\\t', index=False)\n",
    "        print('Test set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(test), test.SessionId.nunique(), test.ItemId.nunique()))\n",
    "        test.to_csv(osp.join(self.processed_dir,'events_test.txt'), sep='\\t', index=False)\n",
    "\n",
    "    def slice_data(self):\n",
    "        for slice_id in range(0, self.num_slices):\n",
    "            self.split_data_slice(slice_id, self.days_offset+(slice_id*self.days_shift))\n",
    "\n",
    "    def split_data_slice(self, slice_id, days_offset):\n",
    "        data = self.data\n",
    "        data_start = datetime.fromtimestamp(data.Time.min(), timezone.utc)\n",
    "        data_end = datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        \n",
    "        print('Full data set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}'.\n",
    "            format(slice_id, len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.isoformat(), data_end.isoformat()))    \n",
    "        \n",
    "        start = datetime.fromtimestamp(data.Time.min(), timezone.utc ) + timedelta(days_offset) \n",
    "        middle =  start + timedelta(self.days_train)\n",
    "        end =  middle + timedelta(self.days_test)\n",
    "        \n",
    "        #prefilter the timespan\n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        greater_start = session_max_times[session_max_times >= start.timestamp()].index\n",
    "        lower_end = session_max_times[session_max_times <= end.timestamp()].index\n",
    "        data_filtered = data[np.in1d(data.SessionId, greater_start.intersection(lower_end))]\n",
    "        \n",
    "        print('Slice data set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {} / {}'.\n",
    "            format( slice_id, len(data_filtered), data_filtered.SessionId.nunique(), data_filtered.ItemId.nunique(), start.date().isoformat(), middle.date().isoformat(), end.date().isoformat() ) )\n",
    "        \n",
    "        #split to train and test\n",
    "        session_max_times = data_filtered.groupby('SessionId').Time.max()\n",
    "        sessions_train = session_max_times[session_max_times < middle.timestamp()].index\n",
    "        sessions_test = session_max_times[session_max_times >= middle.timestamp()].index\n",
    "        \n",
    "        train = data[np.in1d(data.SessionId, sessions_train)]\n",
    "        \n",
    "        print('Train set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}'.\n",
    "            format( slice_id, len(train), train.SessionId.nunique(), train.ItemId.nunique(), start.date().isoformat(), middle.date().isoformat() ) )\n",
    "        \n",
    "        train.to_csv(osp.join(self.processed_dir,'events_train_full.'+str(slice_id)+'.txt'), sep='\\t', index=False)\n",
    "        \n",
    "        test = data[np.in1d(data.SessionId, sessions_test)]\n",
    "        test = test[np.in1d(test.ItemId, train.ItemId)]\n",
    "        \n",
    "        tslength = test.groupby('SessionId').size()\n",
    "        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]\n",
    "        \n",
    "        print('Test set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {} \\n\\n'.\n",
    "            format( slice_id, len(test), test.SessionId.nunique(), test.ItemId.nunique(), middle.date().isoformat(), end.date().isoformat() ) )\n",
    "        \n",
    "        test.to_csv(osp.join(self.processed_dir,'events_test.'+str(slice_id)+'.txt'), sep='\\t', index=False)\n",
    "        \n",
    "    def process(self):\n",
    "        self.load()\n",
    "        self.filter_data()\n",
    "        if self.process_method == 'last':\n",
    "            self.split_data_org()\n",
    "        elif self.process_method == 'last_min_date':\n",
    "            self.filter_min_date()\n",
    "            self.split_data_org()\n",
    "        elif self.process_method == 'days_test':\n",
    "            self.split_data()\n",
    "        elif self.process_method == 'slice':\n",
    "            self.slice_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDatasetv2(Dataset):\n",
    "    def __init__(self, root, column_names):\n",
    "        super().__init__(root)\n",
    "        self.SESSION_ID = column_names['SESSION_ID']\n",
    "        self.ITEM_ID = column_names['ITEM_ID']\n",
    "        self.TIMEFRAME = column_names['TIMEFRAME']\n",
    "        self.EVENT_DATE = column_names['EVENT_DATE']\n",
    "\n",
    "        self._process()\n",
    "    \n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return ['train.txt','test.txt','all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def process(self):\n",
    "        with open(self.raw_paths[0], \"r\") as f:\n",
    "            reader = csv.DictReader(f, delimiter=';')\n",
    "            sess_clicks = {}\n",
    "            sess_date = {}\n",
    "            ctr = 0\n",
    "            curid = -1\n",
    "            curdate = None\n",
    "            for data in reader:\n",
    "                sessid = data[self.SESSION_ID]\n",
    "                if curdate and not curid == sessid:\n",
    "                    date = ''\n",
    "                    date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))\n",
    "                    sess_date[curid] = date\n",
    "                curid = sessid\n",
    "                item = data[self.ITEM_ID], int(data[self.TIMEFRAME])\n",
    "                curdate = ''\n",
    "                curdate = data[self.EVENT_DATE]\n",
    "\n",
    "                if sessid in sess_clicks:\n",
    "                    sess_clicks[sessid] += [item]\n",
    "                else:\n",
    "                    sess_clicks[sessid] = [item]\n",
    "                ctr += 1\n",
    "            date = ''\n",
    "            date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))\n",
    "            for i in list(sess_clicks):\n",
    "                sorted_clicks = sorted(sess_clicks[i], key=operator.itemgetter(1))\n",
    "                sess_clicks[i] = [c[0] for c in sorted_clicks]\n",
    "            sess_date[curid] = date\n",
    "\n",
    "        print(\"-- Reading data\")\n",
    "\n",
    "        # Filter out length 1 sessions\n",
    "        for s in list(sess_clicks):\n",
    "            if len(sess_clicks[s]) == 1:\n",
    "                del sess_clicks[s]\n",
    "                del sess_date[s]\n",
    "\n",
    "        # Count number of times each item appears\n",
    "        iid_counts = {}\n",
    "        for s in sess_clicks:\n",
    "            seq = sess_clicks[s]\n",
    "            for iid in seq:\n",
    "                if iid in iid_counts:\n",
    "                    iid_counts[iid] += 1\n",
    "                else:\n",
    "                    iid_counts[iid] = 1\n",
    "\n",
    "        sorted_counts = sorted(iid_counts.items(), key=operator.itemgetter(1))\n",
    "\n",
    "        length = len(sess_clicks)\n",
    "        for s in list(sess_clicks):\n",
    "            curseq = sess_clicks[s]\n",
    "            filseq = list(filter(lambda i: iid_counts[i] >= 5, curseq))\n",
    "            if len(filseq) < 2:\n",
    "                del sess_clicks[s]\n",
    "                del sess_date[s]\n",
    "            else:\n",
    "                sess_clicks[s] = filseq\n",
    "\n",
    "        # Split out test set based on dates\n",
    "        dates = list(sess_date.items())\n",
    "        maxdate = dates[0][1]\n",
    "\n",
    "        for _, date in dates:\n",
    "            if maxdate < date:\n",
    "                maxdate = date\n",
    "\n",
    "        # 7 days for test\n",
    "        splitdate = 0\n",
    "        splitdate = maxdate - 86400 * 7\n",
    "\n",
    "        print('Splitting date', splitdate)\n",
    "        tra_sess = filter(lambda x: x[1] < splitdate, dates)\n",
    "        tes_sess = filter(lambda x: x[1] > splitdate, dates)\n",
    "\n",
    "        # Sort sessions by date\n",
    "        tra_sess = sorted(tra_sess, key=operator.itemgetter(1))\n",
    "        tes_sess = sorted(tes_sess, key=operator.itemgetter(1))\n",
    "\n",
    "        print(len(tra_sess))\n",
    "        print(len(tes_sess))\n",
    "        print(tra_sess[:3])\n",
    "        print(tes_sess[:3])\n",
    "\n",
    "        print(\"-- Splitting train set and test set\")\n",
    "\n",
    "        item_dict = {}\n",
    "        # Convert training sessions to sequences and renumber items to start from 1\n",
    "        def obtian_tra():\n",
    "            train_ids = []\n",
    "            train_seqs = []\n",
    "            train_dates = []\n",
    "            item_ctr = 1\n",
    "            for s, date in tra_sess:\n",
    "                seq = sess_clicks[s]\n",
    "                outseq = []\n",
    "                for i in seq:\n",
    "                    if i in item_dict:\n",
    "                        outseq += [item_dict[i]]\n",
    "                    else:\n",
    "                        outseq += [item_ctr]\n",
    "                        item_dict[i] = item_ctr\n",
    "                        item_ctr += 1\n",
    "                if len(outseq) < 2:  # Doesn't occur\n",
    "                    continue\n",
    "                train_ids += [s]\n",
    "                train_dates += [date]\n",
    "                train_seqs += [outseq]\n",
    "            print(item_ctr)     # 43098, 37484\n",
    "            return train_ids, train_dates, train_seqs\n",
    "\n",
    "        # Convert test sessions to sequences, ignoring items that do not appear in training set\n",
    "        def obtian_tes():\n",
    "            test_ids = []\n",
    "            test_seqs = []\n",
    "            test_dates = []\n",
    "            for s, date in tes_sess:\n",
    "                seq = sess_clicks[s]\n",
    "                outseq = []\n",
    "                for i in seq:\n",
    "                    if i in item_dict:\n",
    "                        outseq += [item_dict[i]]\n",
    "                if len(outseq) < 2:\n",
    "                    continue\n",
    "                test_ids += [s]\n",
    "                test_dates += [date]\n",
    "                test_seqs += [outseq]\n",
    "            return test_ids, test_dates, test_seqs\n",
    "\n",
    "        tra_ids, tra_dates, tra_seqs = obtian_tra()\n",
    "        tes_ids, tes_dates, tes_seqs = obtian_tes()\n",
    "\n",
    "        def process_seqs(iseqs, idates):\n",
    "            out_seqs = []\n",
    "            out_dates = []\n",
    "            labs = []\n",
    "            ids = []\n",
    "            for id, seq, date in zip(range(len(iseqs)), iseqs, idates):\n",
    "                for i in range(1, len(seq)):\n",
    "                    tar = seq[-i]\n",
    "                    labs += [tar]\n",
    "                    out_seqs += [seq[:-i]]\n",
    "                    out_dates += [date]\n",
    "                    ids += [id]\n",
    "            return out_seqs, out_dates, labs, ids\n",
    "\n",
    "        tr_seqs, tr_dates, tr_labs, tr_ids = process_seqs(tra_seqs, tra_dates)\n",
    "        te_seqs, te_dates, te_labs, te_ids = process_seqs(tes_seqs, tes_dates)\n",
    "        tra = (tr_seqs, tr_labs)\n",
    "        tes = (te_seqs, te_labs)\n",
    "\n",
    "        print(len(tr_seqs))\n",
    "        print(len(te_seqs))\n",
    "        print(tr_seqs[:3], tr_dates[:3], tr_labs[:3])\n",
    "        print(te_seqs[:3], te_dates[:3], te_labs[:3])\n",
    "\n",
    "        all = 0\n",
    "\n",
    "        for seq in tra_seqs:\n",
    "            all += len(seq)\n",
    "        for seq in tes_seqs:\n",
    "            all += len(seq)\n",
    "        print('avg length: ', all/(len(tra_seqs) + len(tes_seqs) * 1.0))\n",
    "\n",
    "        pickle.dump(tra, open(self.processed_paths[0], 'wb'))\n",
    "        pickle.dump(tes, open(self.processed_paths[1], 'wb'))\n",
    "        pickle.dump(tra_seqs, open(self.processed_paths[2], 'wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDatasetv3(Dataset):\n",
    "    def __init__(self, root):\n",
    "        super().__init__(root)\n",
    "\n",
    "        self._process()\n",
    "    \n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return ['train.txt','test.txt','all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def process(self):\n",
    "        def get_session_id(df, interval):\n",
    "            df_prev = df.shift()\n",
    "            is_new_session = (df.userId != df_prev.userId) | (\n",
    "                df.timestamp - df_prev.timestamp > interval\n",
    "            )\n",
    "            session_id = is_new_session.cumsum() - 1\n",
    "            return session_id\n",
    "\n",
    "        def group_sessions(df, interval):\n",
    "            sessionId = get_session_id(df, interval)\n",
    "            df = df.assign(sessionId=sessionId)\n",
    "            return df\n",
    "\n",
    "        def filter_short_sessions(df, min_len=2):\n",
    "            session_len = df.groupby('sessionId', sort=False).size()\n",
    "            long_sessions = session_len[session_len >= min_len].index\n",
    "            df_long = df[df.sessionId.isin(long_sessions)]\n",
    "            return df_long\n",
    "\n",
    "        def filter_infreq_items(df, min_support=5):\n",
    "            item_support = df.groupby('itemId', sort=False).size()\n",
    "            freq_items = item_support[item_support >= min_support].index\n",
    "            df_freq = df[df.itemId.isin(freq_items)]\n",
    "            return df_freq\n",
    "\n",
    "        def filter_until_all_long_and_freq(df, min_len=2, min_support=5):\n",
    "            while True:\n",
    "                df_long = filter_short_sessions(df, min_len)\n",
    "                df_freq = filter_infreq_items(df_long, min_support)\n",
    "                if len(df_freq) == len(df):\n",
    "                    break\n",
    "                df = df_freq\n",
    "            return df\n",
    "\n",
    "        def truncate_long_sessions(df, max_len=20, is_sorted=False):\n",
    "            if not is_sorted:\n",
    "                df = df.sort_values(['sessionId', 'timestamp'])\n",
    "            itemIdx = df.groupby('sessionId').cumcount()\n",
    "            df_t = df[itemIdx < max_len]\n",
    "            return df_t\n",
    "\n",
    "        def update_id(df, field):\n",
    "            labels = pd.factorize(df[field])[0]\n",
    "            kwargs = {field: labels}\n",
    "            df = df.assign(**kwargs)\n",
    "            return df\n",
    "\n",
    "        def remove_immediate_repeats(df):\n",
    "            df_prev = df.shift()\n",
    "            is_not_repeat = (df.sessionId != df_prev.sessionId) | (df.itemId != df_prev.itemId)\n",
    "            df_no_repeat = df[is_not_repeat]\n",
    "            return df_no_repeat\n",
    "\n",
    "        def reorder_sessions_by_endtime(df):\n",
    "            endtime = df.groupby('sessionId', sort=False).timestamp.max()\n",
    "            df_endtime = endtime.sort_values().reset_index()\n",
    "            oid2nid = dict(zip(df_endtime.sessionId, df_endtime.index))\n",
    "            sessionId_new = df.sessionId.map(oid2nid)\n",
    "            df = df.assign(sessionId=sessionId_new)\n",
    "            df = df.sort_values(['sessionId', 'timestamp'])\n",
    "            return df\n",
    "\n",
    "        def keep_top_n_items(df, n):\n",
    "            item_support = df.groupby('itemId', sort=False).size()\n",
    "            top_items = item_support.nlargest(n).index\n",
    "            df_top = df[df.itemId.isin(top_items)]\n",
    "            return df_top\n",
    "\n",
    "        def split_by_time(df, timedelta):\n",
    "            max_time = df.timestamp.max()\n",
    "            end_time = df.groupby('sessionId').timestamp.max()\n",
    "            split_time = max_time - timedelta\n",
    "            train_sids = end_time[end_time < split_time].index\n",
    "            df_train = df[df.sessionId.isin(train_sids)]\n",
    "            df_test = df[~df.sessionId.isin(train_sids)]\n",
    "            return df_train, df_test\n",
    "\n",
    "        def train_test_split(df, test_split=0.2):\n",
    "            endtime = df.groupby('sessionId', sort=False).timestamp.max()\n",
    "            endtime = endtime.sort_values()\n",
    "            num_tests = int(len(endtime) * test_split)\n",
    "            test_session_ids = endtime.index[-num_tests:]\n",
    "            df_train = df[~df.sessionId.isin(test_session_ids)]\n",
    "            df_test = df[df.sessionId.isin(test_session_ids)]\n",
    "            return df_train, df_test\n",
    "\n",
    "        def save_sessions(df, filepath):\n",
    "            df = reorder_sessions_by_endtime(df)\n",
    "            sessions = df.groupby('sessionId').itemId.apply(lambda x: ','.join(map(str, x)))\n",
    "            sessions.to_csv(filepath, sep='\\t', header=False, index=False)\n",
    "\n",
    "        def save_dataset(df_train, df_test):\n",
    "            # filter items in test but not in train\n",
    "            df_test = df_test[df_test.itemId.isin(df_train.itemId.unique())]\n",
    "            df_test = filter_short_sessions(df_test)\n",
    "\n",
    "            print(f'No. of Clicks: {len(df_train) + len(df_test)}')\n",
    "            print(f'No. of Items: {df_train.itemId.nunique()}')\n",
    "\n",
    "            # update itemId\n",
    "            train_itemId_new, uniques = pd.factorize(df_train.itemId)\n",
    "            df_train = df_train.assign(itemId=train_itemId_new)\n",
    "            oid2nid = {oid: i for i, oid in enumerate(uniques)}\n",
    "            test_itemId_new = df_test.itemId.map(oid2nid)\n",
    "            df_test = df_test.assign(itemId=test_itemId_new)\n",
    "\n",
    "            print('saving dataset')\n",
    "            save_sessions(df_train, self.processed_paths[0])\n",
    "            save_sessions(df_test, self.processed_paths[1])\n",
    "            num_items = len(uniques)\n",
    "            with open(self.processed_paths[2], 'w') as f:\n",
    "                f.write(str(num_items))\n",
    "\n",
    "        df = pd.read_csv(\n",
    "            osp.join(self.raw_dir,self.raw_file_names),\n",
    "            sep='\\t',\n",
    "            header=None,\n",
    "            names=['userId', 'timestamp', 'itemId'],\n",
    "            usecols=[0, 1, 4],\n",
    "            parse_dates=['timestamp'],\n",
    "            infer_datetime_format=True,\n",
    "        )\n",
    "\n",
    "        print('start preprocessing')\n",
    "        df = df.dropna()\n",
    "        df = update_id(df, 'userId')\n",
    "        df = update_id(df, 'itemId')\n",
    "        df = df.sort_values(['userId', 'timestamp'])\n",
    "\n",
    "        df = group_sessions(df, Timedelta(days=1))\n",
    "        df = remove_immediate_repeats(df)\n",
    "        df = truncate_long_sessions(df, is_sorted=True)\n",
    "        df = keep_top_n_items(df, n=30000)\n",
    "        df = filter_until_all_long_and_freq(df)\n",
    "        df_train, df_test = train_test_split(df, test_split=0.2)\n",
    "        save_dataset(df_train, df_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GraphData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def data_masks(all_usr_pois, item_tail):\n",
    "    us_lens = [len(upois) for upois in all_usr_pois]\n",
    "    len_max = max(us_lens)\n",
    "    us_pois = [upois + item_tail * (len_max - le) for upois, le in zip(all_usr_pois, us_lens)]\n",
    "    us_msks = [[1] * le + [0] * (len_max - le) for le in us_lens]\n",
    "    return us_pois, us_msks, len_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class GraphData():\n",
    "    def __init__(self, data, shuffle=False, graph=None):\n",
    "        inputs = data[0]\n",
    "        inputs, mask, len_max = data_masks(inputs, [0])\n",
    "        self.inputs = np.asarray(inputs)\n",
    "        self.mask = np.asarray(mask)\n",
    "        self.len_max = len_max\n",
    "        self.targets = np.asarray(data[1])\n",
    "        self.length = len(inputs)\n",
    "        self.shuffle = shuffle\n",
    "        self.graph = graph\n",
    "\n",
    "    def generate_batch(self, batch_size):\n",
    "        if self.shuffle:\n",
    "            shuffled_arg = np.arange(self.length)\n",
    "            np.random.shuffle(shuffled_arg)\n",
    "            self.inputs = self.inputs[shuffled_arg]\n",
    "            self.mask = self.mask[shuffled_arg]\n",
    "            self.targets = self.targets[shuffled_arg]\n",
    "        n_batch = int(self.length / batch_size)\n",
    "        if self.length % batch_size != 0:\n",
    "            n_batch += 1\n",
    "        slices = np.split(np.arange(n_batch * batch_size), n_batch)\n",
    "        slices[-1] = slices[-1][:(self.length - batch_size * (n_batch - 1))]\n",
    "        return slices\n",
    "\n",
    "    def get_slice(self, i):\n",
    "        inputs, mask, targets = self.inputs[i], self.mask[i], self.targets[i]\n",
    "        items, n_node, A, alias_inputs = [], [], [], []\n",
    "        for u_input in inputs:\n",
    "            n_node.append(len(np.unique(u_input)))\n",
    "        max_n_node = np.max(n_node)\n",
    "        for u_input in inputs:\n",
    "            node = np.unique(u_input)\n",
    "            items.append(node.tolist() + (max_n_node - len(node)) * [0])\n",
    "            u_A = np.zeros((max_n_node, max_n_node))\n",
    "            for i in np.arange(len(u_input) - 1):\n",
    "                if u_input[i + 1] == 0:\n",
    "                    break\n",
    "                u = np.where(node == u_input[i])[0][0]\n",
    "                v = np.where(node == u_input[i + 1])[0][0]\n",
    "                u_A[u][v] = 1\n",
    "            u_sum_in = np.sum(u_A, 0)\n",
    "            u_sum_in[np.where(u_sum_in == 0)] = 1\n",
    "            u_A_in = np.divide(u_A, u_sum_in)\n",
    "            u_sum_out = np.sum(u_A, 1)\n",
    "            u_sum_out[np.where(u_sum_out == 0)] = 1\n",
    "            u_A_out = np.divide(u_A.transpose(), u_sum_out)\n",
    "            u_A = np.concatenate([u_A_in, u_A_out]).transpose()\n",
    "            A.append(u_A)\n",
    "            alias_inputs.append([np.where(node == i)[0][0] for i in u_input])\n",
    "        return alias_inputs, A, items, mask, targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([0]), array([1]), array([2]), array([3]), array([4])]\n",
      "[array([0, 1]), array([2, 3]), array([4])]\n",
      "[[1 2 3]\n",
      " [2 3 4]\n",
      " [1 2 4]\n",
      " [2 3 0]\n",
      " [1 0 0]]\n",
      "[array([0]), array([1]), array([2]), array([3]), array([4])]\n",
      "[array([0, 1]), array([2, 3]), array([4])]\n",
      "[[1 2 4]\n",
      " [2 3 4]\n",
      " [1 0 0]\n",
      " [2 3 0]\n",
      " [1 2 3]]\n"
     ]
    }
   ],
   "source": [
    "train_data = ([[1, 2, 3], [2, 3, 4], [1, 2, 4], [2, 3], [1]], \n",
    "              [4, 5, 5, 4, 2])\n",
    "\n",
    "tds = GraphData(train_data, shuffle=False)\n",
    "print(tds.generate_batch(1))\n",
    "print(tds.generate_batch(2))\n",
    "print(tds.inputs)\n",
    "\n",
    "tds = GraphData(train_data, shuffle=True)\n",
    "print(tds.generate_batch(1))\n",
    "print(tds.generate_batch(2))\n",
    "print(tds.inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-18 09:38:50\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.104+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "IPython: 5.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "!pip install -q watermark\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
