#
# Copyright (c) 2021, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os

import merlin.io
from merlin.core.dispatch import get_lib

DATETIME_CONVERTION = "ms"

df_lib = get_lib()


def process_item_features(input_path, category_coverage_min=0.8):
    df = df_lib.read_csv(os.path.join(input_path, "item_features.csv"))
    tmp = df["feature_category_id"].value_counts() / df.item_id.nunique()
    categories_to_keep = [3, 4, 5, 17, 24, 30, 45, 46, 53, 55, 58, 63, 65, 73]
    df = df[
        (df.feature_category_id.isin(categories_to_keep))
        | (df.feature_category_id.isin(tmp[tmp >= category_coverage_min].index))
    ]
    df = df[~df.feature_category_id.isin([30, 4, 46, 28, 53, 1])]
    df = df.pivot(index="item_id", columns="feature_category_id").reset_index()
    df.columns = ["item_id"] + df.columns.get_level_values(1)[1:].tolist()
    feat_columns = [col for col in df.columns if isinstance(col, int)]
    df[feat_columns] = df[feat_columns].fillna(-1).astype("int")
    df.columns = ["f_" + str(col) if isinstance(col, int) else str(col) for col in df.columns]
    return df


# add timestamp and day features
def process_date_column(ddf):
    ddf["date"] = ddf["date"].astype(f"datetime64[{DATETIME_CONVERTION}]")
    ddf["timestamp"] = ddf["date"].astype("int64")
    ddf = ddf.sort_values(["session_id", "date"]).reset_index(drop=True)
    ddf["day"] = (ddf["date"] - ddf["date"].min()).dt.days
    return ddf


def get_dressipi2022(input_path):
    """
    Load the dressipi dataset from the RecSys 2022 Challenge.

    We can't download this dataset automatically so this needs to be downloaded manually
    To use this function, prepare the data by following these 3 steps:
    1. Sign up and download the data from
    [dressipi-recsys2022.com](https://www.dressipi-recsys2022.com/).
    2. Unzip the raw data to a directory.
    3. Run `get_dressipi2022(data_dir)` to convert the data to Merlin datasets

    In case you want to use this dataset to run our examples, you can also opt for synthetic data.
    Synthetic data can be generated by running::

        from merlin.datasets.synthetic import generate_data
        train, valid = generate_data("dressipi2022-preprocessed")

    Parameters
    ----------
    path: Union[str, Path]
        Directory to load the raw data from.
    Returns
    -------
    train: merlin.io.Dataset
        Training dataset.
    valid: merlin.io.Dataset
        Test dataset.
    """
    if not os.path.exists(input_path):
        print(
            "Dressipi dataset not found. Please follow the above instructions to download \
            the data."
        )
        return

    # get the item features
    item_features = process_item_features(input_path)

    # load data
    sessions = df_lib.read_csv(os.path.join(input_path, "train_sessions.csv"))
    purchases = df_lib.read_csv(os.path.join(input_path, "train_purchases.csv"))

    # merge session data with item features
    sessions = df_lib.merge(sessions, item_features, on="item_id", how="left")
    purchases = df_lib.merge(purchases, item_features, on="item_id", how="left")

    # add timestamp and day features, and convert the format of date to ms
    sessions = process_date_column(sessions)
    purchases = process_date_column(purchases)
    purchases = purchases.rename(columns={"item_id": "purchase_id"})

    # Split into train and validation set
    train_session = (
        sessions.loc[sessions.day <= (sessions.day.max() - 30)].copy().reset_index(drop=True)
    )
    valid_session = (
        sessions.loc[sessions.day > (sessions.day.max() - 30)].copy().reset_index(drop=True)
    )
    train_session.shape, valid_session.shape

    # Merge with train + valid purchases
    train_purchases = purchases[
        purchases.session_id.isin(train_session.session_id.unique().values.tolist())
    ]
    valid_purchases = purchases[
        purchases.session_id.isin(valid_session.session_id.unique().values.tolist())
    ]

    train_purchases = train_purchases[["session_id", "purchase_id"]]
    train = df_lib.merge(train_session, train_purchases, on="session_id", how="left")

    valid_purchases = valid_purchases[["session_id", "purchase_id"]]
    valid = df_lib.merge(valid_session, valid_purchases, on="session_id", how="left")

    return merlin.io.Dataset(train), merlin.io.Dataset(valid)
