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

__all__ = ['TianchiDataset']

# Cell
import os
import os.path as osp
import sys
import numpy as np
import pandas as pd

from .base import Dataset
from ..utils.common_utils import download_url, extract_zip, makedirs

# Cell
class TianchiDataset(Dataset):
    data_id = '16XaYNZHn4Sb33sAI7J3N_QkScVNmvYGa'

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

        self._process()

    @property
    def raw_file_names(self) -> str:
        return ['user.csv', 'item.csv', 'user_behavior.csv']

    @property
    def processed_file_names(self) -> str:
        return 'tianchi.csv'

    def download(self):
        from google_drive_downloader import GoogleDriveDownloader as gdd
        from shutil import move, rmtree

        path = osp.join(self.raw_dir, 'tianchi_raw.zip')
        gdd.download_file_from_google_drive(self.data_id, path)
        extract_zip(path, self.raw_dir)
        for f in self.raw_file_names:
            move(osp.join(self.raw_dir, 'ECommAI_EUIR_round2_train_20190816', f),
                 osp.join(self.raw_dir, f))
        rmtree(osp.join(self.raw_dir, 'ECommAI_EUIR_round2_train_20190816'))
        os.unlink(path)

    @staticmethod
    def bucket_age(age):
        if age < 30:
            return 1
        elif age < 40:
            return 2
        elif age < 50:
            return 3
        else:
            return 4

    def process(self):
        # 1. loading the data into memory

        user_feat = pd.read_csv(self.raw_paths[0], header=None,
                                names=["user", "sex", "age", "pur_power"])
        item_feat = pd.read_csv(self.raw_paths[1], header=None,
                                names=["item", "category", "shop", "brand"])
        behavior = pd.read_csv(self.raw_paths[2], header=None,
                            names=["user", "item", "behavior", "time"])

        # 2. sorting values chronologically and dropping duplicate records

        behavior = behavior.sort_values(by="time").reset_index(drop=True)
        behavior = behavior.drop_duplicates(subset=["user", "item", "behavior"])

        # 3. Choosing 60K random users with short journey and 20K with long journey
        user_counts = behavior.groupby("user")[["user"]].count().rename(
            columns={"user": "count_user"}
        ).sort_values("count_user", ascending=False)

        short_users = np.array(
            user_counts[
                (user_counts.count_user > 5) & (user_counts.count_user <= 50)
            ].index
        )
        long_users = np.array(
            user_counts[
                (user_counts.count_user > 50) & (user_counts.count_user <= 200)
            ].index
        )
        short_chosen_users = np.random.choice(short_users, 60000, replace=False)
        long_chosen_users = np.random.choice(long_users, 20000, replace=False)
        chosen_users = np.concatenate([short_chosen_users, long_chosen_users])

        behavior = behavior[behavior.user.isin(chosen_users)]
        print(f"n_users: {behavior.user.nunique()}, "
            f"n_items: {behavior.item.nunique()}, "
            f"behavior length: {len(behavior)}")

        # 4. merge with all features, bucketizing the age and saving the processed data
        behavior = behavior.merge(user_feat, on="user")
        behavior = behavior.merge(item_feat, on="item")
        behavior["age"] = behavior["age"].apply(self.bucket_age)
        behavior = behavior.sort_values(by="time").reset_index(drop=True)
        behavior.to_csv(self.processed_paths[0], header=None, index=False)