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

__all__ = ['ML1mDataset', 'ML1mDataModule', 'ML100kDataset']

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

import pandas as pd

from ..utils.common_utils import *
from .bases.common import Dataset
from .bases.interactions import InteractionsDataset, InteractionsDataModule

# Cell
class ML1mDataset(InteractionsDataset):
    url = "http://files.grouplens.org/datasets/movielens/ml-1m.zip"

    @property
    def raw_file_names(self):
        return 'ratings.dat'

    def download(self):
        path = download_url(self.url, self.raw_dir)
        extract_zip(path, self.raw_dir)
        from shutil import move, rmtree
        move(os.path.join(self.raw_dir, 'ml-1m', self.raw_file_names), self.raw_dir)
        rmtree(os.path.join(self.raw_dir, 'ml-1m'))
        os.unlink(path)

    def load_ratings_df(self):
        df = pd.read_csv(self.raw_paths[0], sep='::', header=None, engine='python')
        df.columns = ['uid', 'sid', 'rating', 'timestamp']
        # drop duplicate user-item pair records, keeping recent ratings only
        df.drop_duplicates(subset=['uid', 'sid'], keep='last', inplace=True)
        return df

# Cell
class ML1mDataModule(InteractionsDataModule):
    dataset_cls = ML1mDataset

# Cell
class ML100kDataset(Dataset):
    url = 'https://files.grouplens.org/datasets/movielens/ml-100k.zip'

    def __init__(self, root):
        super().__init__(root)

    @property
    def raw_file_names(self) -> str:
        return ['u1.base', 'u1.test', 'u4.test', 'allbut.pl', 'u.item',
                'ua.test', 'u.occupation', 'u3.test', 'u5.base', 'ub.test',
                'u2.test', 'u3.base', 'u.genre', 'u.data', 'u4.base',
                'u5.test', 'u.info', 'README', 'ub.base', 'mku.sh', 'u2.base',
                'u.user', 'ua.base']

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

    def download(self):
        path = download_url(self.url, self.raw_dir)
        extract_zip(path, self.raw_dir)
        from shutil import move, rmtree
        file_names = os.listdir(osp.join(self.raw_dir, 'ml-100k'))
        for file_name in file_names:
            move(osp.join(self.raw_dir, 'ml-100k', file_name), self.raw_dir)
        rmtree(osp.join(self.raw_dir, 'ml-100k'))
        os.unlink(path)

    def process(self):
        raise NotImplementedError