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

__all__ = ['InteractionsDataset', 'InteractionsDataModule']

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

import sys
import os
from os import path as osp
import collections

import random
import numpy as np
import pandas as pd
from scipy.sparse import coo_matrix

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningDataModule

from ...utils.common_utils import *
from .. import base
from ...transforms.splitting import random_split, stratified_split

import warnings
warnings.filterwarnings('ignore')

# Cell
class InteractionsDataset(torch.utils.data.Dataset, base.Dataset):

    def __init__(self,
                 data_dir,
                 min_rating=None,
                 min_uc=5,
                 min_sc=5,
                 num_negative_samples=100,
                 max_samples=200,
                 data_type=None,
                 split_type='random',
                 val_p=0.2,
                 test_p=0.2,
                 seed=42,
                 *args,
                 **kwargs):
        """
        Args:
            data_dir: Where to save/load the data
            min_uc: minimum user count to keep in the data
            min_sc: minimum item count to keep in the data
            min_rating: minimum rating threshold to convert explicit feedback into implicit
            num_negative_samples: number of negative samples for each positive one
            max_samples: max samples limit
            data_type: train/valid/test
            split_type: data split method - stratified/random
            val_p: Percent (float) or number (int) of samples to use for the validation split
            test_p: Percent (float) or number (int) of samples to use for the test split
            seed: Random seed to be used for train/val/test splits
        """
        self.min_rating = min_rating
        self.min_uc = min_uc
        self.min_sc = min_sc
        self.num_negative_samples = num_negative_samples
        self.max_samples = 200
        self.data_type = data_type
        self.val_p = val_p if val_p is not None else 0.2
        self.test_p = test_p if test_p is not None else 0.2
        self.seed = seed
        self.split_type = split_type

        super().__init__(data_dir)

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

        self._process()

        if self.data_type is not None:
            self.load()

    @property
    def raw_file_names(self):
        raise NotImplementedError

    @property
    def processed_file_names(self):
        return ['data_train.pt',
                'data_valid_pos.pt',
                'data_valid_neg.pt',
                'data_test_pos.pt',
                'data_test_neg.pt']

    def download(self):
        raise NotImplementedError

    def make_implicit(self, df):
        "convert the explicit data to implicit by only keeping interactions with a rating >= min_rating"
        print('Turning into implicit ratings')
        df = df[df['rating'] >= self.min_rating].reset_index(drop=True)
        df['rating'] = 1
        return df

    def filter_triplets(self, df):
        print('Filtering triplets')
        if self.min_sc > 0 or self.min_uc > 0:
            item_sizes = df.groupby('sid').size()
            good_items = item_sizes.index[item_sizes >= self.min_sc]
            user_sizes = df.groupby('uid').size()
            good_users = user_sizes.index[user_sizes >= self.min_uc]
            while len(good_items) < len(item_sizes) or len(good_users) < len(user_sizes):
                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)]

                item_sizes = df.groupby('sid').size()
                good_items = item_sizes.index[item_sizes >= self.min_sc]
                user_sizes = df.groupby('uid').size()
                good_users = user_sizes.index[user_sizes >= self.min_uc]
        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 load_ratings_df(self):
        raise NotImplementedError

    @staticmethod
    def _convert_to_torch_sparse(mat):
        values = mat.data
        indices = np.vstack((mat.row, mat.col))

        i = torch.LongTensor(indices)
        v = torch.FloatTensor(values)
        shape = mat.shape

        return torch.sparse.FloatTensor(i, v, torch.Size(shape))

    def process(self):
        df = self.load_ratings_df()
        if self.min_rating:
            df = self.make_implicit(df)
        df = self.filter_triplets(df)
        df, umap, smap = self.densify_index(df)
        self.num_users = max(df.uid) + 1 # df.uid.nunique()
        self.num_items = max(df.sid) + 1 # df.sid.nunique()
        mat = coo_matrix((np.array(df.rating),
                          (np.array(df.uid), np.array(df.sid))),
                         shape=(self.num_users, self.num_items))

        if self.split_type == 'random':
            mat_train, mat_valid, mat_test = random_split(mat = mat,
                                                        val_p = self.val_p,
                                                        test_p = self.test_p,
                                                        seed = self.seed)
        elif self.split_type == 'stratified':
            mat_train, mat_valid, mat_test = stratified_split(mat = mat,
                                                        val_p = self.val_p,
                                                        test_p = self.test_p,
                                                        seed = self.seed)

        mat_train = self._convert_to_torch_sparse(mat_train)
        torch.save(mat_train, self.processed_paths[0])

        mat_valid_pos = self._convert_to_torch_sparse(mat_valid)._indices().T
        _, indices = np.unique(mat_valid_pos[:, 0], return_index=True)
        mat_valid_pos = mat_valid_pos[indices, :]
        torch.save(mat_valid_pos, self.processed_paths[1])

        pos_items = set(zip(mat_valid.row, mat_valid.col))
        mat_valid_neg = self._negative_sample(np.arange(mat_valid.shape[0]), pos_items)
        mat_valid_neg = torch.tensor(mat_valid_neg, dtype=torch.int)
        torch.save(mat_valid_neg, self.processed_paths[2])

        mat_test_pos = self._convert_to_torch_sparse(mat_test)._indices().T
        _, indices = np.unique(mat_test_pos[:, 0], return_index=True)
        mat_test_pos = mat_test_pos[indices, :]
        torch.save(mat_test_pos, self.processed_paths[3])

        pos_items = set(zip(mat_test.row, mat_test.col))
        mat_test_neg = self._negative_sample(np.arange(mat_test.shape[0]), pos_items)
        mat_test_neg = torch.tensor(mat_test_neg, dtype=torch.int)
        torch.save(mat_test_neg, self.processed_paths[4])

        return mat

    def todense(self) -> np.matrix:
        """Transforms sparse matrix to np.matrix, 2-d."""
        return self.mat.todense()

    def toarray(self) -> np.array:
        """Transforms sparse matrix to np.array, 2-d."""
        return self.mat.toarray()

    def head(self, n: int = 5) -> np.array:
        """Return the first ``n`` rows of the dense matrix as a np.array, 2-d."""
        n = self._prep_head_tail_n(n=n)
        return self.mat.tocsr()[range(n), :].toarray()

    def tail(self, n: int = 5) -> np.array:
        """Return the last ``n`` rows of the dense matrix as a np.array, 2-d."""
        n = self._prep_head_tail_n(n=n)
        return self.mat.tocsr()[range(-n, 0), :].toarray()

    def _prep_head_tail_n(self, n: int) -> int:
        """Ensure we don't run into an ``IndexError`` when using ``head`` or ``tail`` methods."""
        if n < 0:
            n = self.num_users + n
        if n > self.num_users:
            n = self.num_users
        return n

    def _negative_sample(self, user_id: Union[int, np.array], positive_items) -> np.array:
        """Generate negative samples for a ``user_id``."""
        if self.max_samples > 0:
            # if we are here, we are doing true negative sampling
            negative_item_ids_list = list()

            if not isinstance(user_id, collections.abc.Iterable):
                user_id = [user_id]

            for specific_user_id in user_id:
                # generate true negative samples for the ``user_id``
                samples_checked = 0
                temp_negative_item_ids_list = list()

                while len(temp_negative_item_ids_list) < self.num_negative_samples:
                    negative_item_id = random.choice(range(self.num_items))
                    # we have a negative sample, make sure the user has not interacted with it
                    # before, else we resample and try again
                    while (
                        (specific_user_id, negative_item_id) in positive_items
                        or negative_item_id in temp_negative_item_ids_list
                    ):
                        if samples_checked >= self.max_samples:
                            num_samples_left_to_generate = (
                                self.num_negative_samples - len(temp_negative_item_ids_list) - 1
                            )
                            temp_negative_item_ids_list += random.choices(
                                range(self.num_items), k=num_samples_left_to_generate
                            )
                            break

                        negative_item_id = random.choice(range(self.num_items))
                        samples_checked += 1

                    temp_negative_item_ids_list.append(negative_item_id)

                negative_item_ids_list += [np.array(temp_negative_item_ids_list)]

            if len(user_id) > 1:
                negative_item_ids_array = np.stack(negative_item_ids_list)
            else:
                negative_item_ids_array = negative_item_ids_list[0]
        else:
            # if we are here, we are doing approximate negative sampling
            if isinstance(user_id, collections.abc.Iterable):
                size = (len(user_id), self.num_negative_samples)
            else:
                size = (self.num_negative_samples,)

            negative_item_ids_array = np.random.randint(
                low=0,
                high=self.num_items,
                size=size,
            )

        return negative_item_ids_array

    def load(self):
        if self.data_type=='train':
            self.train = torch.load(self.processed_paths[0])
            self.train_pos = self.train._indices().T
            self.n_users, self.n_items = self.train.size()

            self.train_score = torch.sparse.sum(self.train, dim=0).to_dense().repeat((self.n_users, 1))
            self.train_score[self.train_pos[:, 0], self.train_pos[:, 1]] = 0
        elif self.data_type=='valid':
            self.valid_pos = torch.load(self.processed_paths[1])
            self.valid_neg = torch.load(self.processed_paths[2])
            self.n_users = self.valid_pos.shape[0]

            valid_items = []
            for u in range(self.n_users):
                items = torch.cat((self.valid_pos[u, 1].view(1), self.valid_neg[u]))
                valid_items.append(items)

            self.valid_items = torch.vstack(valid_items)
            self.valid_labels = torch.zeros(self.valid_items.shape)
            self.valid_labels[:, 0] += 1
        else:
            self.test_pos = torch.load(self.processed_paths[3])
            self.test_neg = torch.load(self.processed_paths[4])
            self.n_users = self.test_pos.shape[0]

            test_items = []
            for u in range(self.n_users):
                items = torch.cat((self.test_pos[u, 1].view(1), self.test_neg[u]))
                test_items.append(items)

            self.test_items = torch.vstack(test_items)
            self.test_labels = torch.zeros(self.test_items.shape)
            self.test_labels[:, 0] += 1

    def __len__(self):
            return self.n_users

    def __train__(self, index):
            return self.train_pos[index], self.train_score[self.train_pos[index][0]]

    def __valid__(self, index):
            return self.valid_pos[index], self.valid_items[index], self.valid_labels[index]

    def __test__(self, index):
            return self.test_pos[index], self.test_items[index], self.test_labels[index]

    def __getitem__(self, index):
        if self.data_type=='train':
            return self.__train__(index)
        elif self.data_type=='valid':
            return self.__valid__(index)
        else:
            return self.__test__(index)

# Cell
class InteractionsDataModule(LightningDataModule):

    dataset_cls: str = ""

    def __init__(self,
                 data_dir: Optional[str] = None,
                 num_workers: int = 0,
                 normalize: bool = False,
                 batch_size: int = 32,
                 shuffle: bool = True,
                 pin_memory: bool = True,
                 drop_last: bool = False,
                 *args,
                 **kwargs) -> None:
        """
        Args:
            data_dir: Where to save/load the data
            num_workers: How many workers to use for loading data
            normalize: If true applies rating normalize
            batch_size: How many samples per batch to load
            shuffle: If true shuffles the train data every epoch
            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before
                        returning them
            drop_last: If true drops the last incomplete batch
        """
        super().__init__(data_dir)

        self.data_dir = data_dir if data_dir is not None else os.getcwd()
        self.num_workers = num_workers
        self.normalize = normalize
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.pin_memory = pin_memory
        self.drop_last = drop_last
        self.kwargs = kwargs

    def prepare_data(self, *args: Any, **kwargs: Any) -> None:
        """Saves files to data_dir."""
        self.data = self.dataset_cls(self.data_dir, **self.kwargs)

    def setup(self, stage: Optional[str] = None) -> None:
        """Creates train, val, and test dataset."""
        if stage == "fit" or stage is None:
            self.dataset_train = self.dataset_cls(self.data_dir, **self.kwargs, data_type='train')
            self.dataset_val = self.dataset_cls(self.data_dir, **self.kwargs, data_type='valid')
        if stage == "test" or stage is None:
            self.dataset_test = self.dataset_cls(self.data_dir, **self.kwargs, data_type='test')

    def train_dataloader(self, *args: Any, **kwargs: Any) -> DataLoader:
        """The train dataloader."""
        return self._data_loader(self.dataset_train, shuffle=self.shuffle)

    def val_dataloader(self, *args: Any, **kwargs: Any) -> Union[DataLoader, List[DataLoader]]:
        """The val dataloader."""
        return self._data_loader(self.dataset_val)

    def test_dataloader(self, *args: Any, **kwargs: Any) -> Union[DataLoader, List[DataLoader]]:
        """The test dataloader."""
        return self._data_loader(self.dataset_test)

    def _data_loader(self, dataset: Dataset, shuffle: bool = False) -> DataLoader:
        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            drop_last=self.drop_last,
            pin_memory=self.pin_memory,
        )