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

__all__ = ['SampleDataset', 'SampleDatasetv2']

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

import os
import os.path as osp
from collections.abc import Sequence
import sys
import csv
import pickle
import math
import operator

import numpy as np
from datetime import timezone, datetime, timedelta
import time

from .bases import common as base
from ..utils.common_utils import download_url

# Cell
class SampleDataset(base.SessionDatasetv2):
    url = 'https://github.com/RecoHut-Datasets/sample_session/raw/v1/sample_train-item-views.csv'

    def __init__(self, root, column_names={'SESSION_ID':'session_id',
                                        'ITEM_ID': 'item_id',
                                        'TIMEFRAME': 'timeframe',
                                        'EVENT_DATE': 'eventdate'}):
        super().__init__(root, column_names)

    @property
    def raw_file_names(self) -> str:
        return 'sample_train-item-views.csv'

    def download(self):
        path = download_url(self.url, self.raw_dir)

# 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

# Internal Cell
def split_validation(train_set, valid_portion):
    train_set_x, train_set_y = train_set
    n_samples = len(train_set_x)
    sidx = np.arange(n_samples, dtype='int32')
    np.random.shuffle(sidx)
    n_train = int(np.round(n_samples * (1. - valid_portion)))
    valid_set_x = [train_set_x[s] for s in sidx[n_train:]]
    valid_set_y = [train_set_y[s] for s in sidx[n_train:]]
    train_set_x = [train_set_x[s] for s in sidx[:n_train]]
    train_set_y = [train_set_y[s] for s in sidx[:n_train]]

    return (train_set_x, train_set_y), (valid_set_x, valid_set_y)

# Cell
class SampleDatasetv2(base.SessionGraphDataset):
    train_url = "https://github.com/RecoHut-Datasets/sample_session/raw/v2/train.txt"
    test_url = "https://github.com/RecoHut-Datasets/sample_session/raw/v2/test.txt"
    all_train_seq_url = "https://github.com/RecoHut-Datasets/sample_session/raw/v2/all_train_seq.txt"

    def __init__(self, root, shuffle=False, n_node=309, is_train=True):
        self.n_node = n_node
        self.shuffle = shuffle
        self.is_train = is_train
        super().__init__(root, shuffle, n_node)

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

    def download(self):
        download_url(self.all_train_seq_url, self.raw_dir)
        if self.is_train:
            download_url(self.train_url, self.raw_dir)
        else:
            download_url(self.test_url, self.raw_dir)