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

__all__ = ['TmallDataset']

# Cell
import numpy as np

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

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

    def __init__(self, root, shuffle=False, n_node=40727, 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)