#
# 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 pickle
import re
import shutil
from dataclasses import dataclass
from glob import glob
from pathlib import Path
from typing import List, Optional, Tuple, Union

import numpy as np
from tqdm import tqdm

import merlin.io
from merlin.core.dispatch import get_lib
from merlin.models.utils.example_utils import workflow_fit_transform
from merlin.models.utils.nvt_utils import require_nvt
from merlin.schema import Tags

try:
    import nvtabular as nvt

    Workflow = nvt.Workflow
except ImportError:
    Workflow = None


def get_aliccp(
    path: Union[str, Path],
    overwrite: bool = False,
    transformed_name: str = "transformed",
    nvt_workflow: Optional[Workflow] = None,
    file_size: int = 100000,
    **kwargs,
) -> Tuple[merlin.io.Dataset, merlin.io.Dataset]:
    """
    Load the Ali-CCP: Alibaba Click and Conversion Prediction dataset.

    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. Download the raw data from
    [tianchi.aliyun.com](https://tianchi.aliyun.com/dataset/dataDetail?dataId=408#1).
    2. Unzip the raw data to a directory.
    3. Run `prepare_alliccp(data_dir)` to convert the raw data to parquet files.

    Downloading & preparing the data can take quite a while.
    In case you want to use this dataset to run our tutorials, you can also opt for synthetic data.
    Synthetic data can be generated by running::
        from merlin.models.data.synthetic import generate_data

        data = generate_data("aliccp")

    Note, running this function requires the
    [nvtabular](https://github.com/NVIDIA-Merlin/NVTabular) library.

    Parameters
    ----------
    path: Union[str, Path]
        Directory to load the raw data from.
    overwrite: bool
        Whether to overwrite the transformed data.
    transformed_name: str
        Name of the transformed data.
        This is useful when you want to load the transformed data
        multiple times using different workflows.
    nvt_workflow: Optional[Workflow]
        Workflow to transform the raw data.
        If None, the default workflow will be used from `default_aliccp_transformation`.
    file_size: int
        Size of the parquet files to be loaded when preparing the data.

    Returns
    -------
    train: merlin.io.Dataset
        Training dataset.
    valid: merlin.io.Dataset
        Test dataset.
    """

    require_nvt()

    p = Path(path)
    raw_path = p / "raw"
    if not raw_path.exists():
        raw_path.mkdir(parents=True)
        prepare_alliccp(path, output_dir=raw_path, file_size=file_size, **kwargs)

    nvt_path = p / transformed_name
    train_path, valid_path = nvt_path / "train", nvt_path / "valid"
    nvt_path_exists = train_path.exists() and valid_path.exists()
    if not nvt_path_exists or overwrite:
        transform_aliccp(raw_path, nvt_path, nvt_workflow=nvt_workflow)

    train = merlin.io.Dataset(str(train_path), engine="parquet")
    valid = merlin.io.Dataset(str(valid_path), engine="parquet")

    return train, valid


def prepare_alliccp(
    data_dir: Union[str, Path],
    convert_train: bool = True,
    convert_test: bool = True,
    file_size: int = 100000,
    max_num_rows: Optional[int] = None,
    pickle_common_features=True,
    output_dir: Optional[Union[str, Path]] = None,
):
    """
    Convert Ali-CPP data to parquet files.

    To download the raw the Ali-CCP training and test datasets visit
    [tianchi.aliyun.com](https://tianchi.aliyun.com/dataset/dataDetail?dataId=408#1).

    Parameters
    ----------
    data_dir: Union[str, Path]
        Directory to load the raw data from.
    convert_train: bool
        Whether to convert the training data.
    convert_test: bool
        Whether to convert the test data.
    file_size: int
        Number of rows to write to each parquet file.
    max_num_rows: int, optional
        Maximum number of rows to read from the raw data.
    pickle_common_features: bool
        Whether to pickle the common features.
        When enabled it will make the conversion faster if it would be run again.
    output_dir: Union[str, Path], optional
        Directory to write the converted data to.
        If not specified the data will be written to the same directory as the raw data.

    Returns
    -------
    data_dir
    """

    if convert_train:
        _convert_data(
            str(data_dir),
            file_size,
            is_train=True,
            max_num_rows=max_num_rows,
            pickle_common_features=pickle_common_features,
            output_dir=output_dir,
        )
    if convert_test:
        _convert_data(
            str(data_dir),
            file_size,
            is_train=False,
            max_num_rows=max_num_rows,
            pickle_common_features=pickle_common_features,
            output_dir=output_dir,
        )

    return data_dir


def default_aliccp_transformation(add_target_encoding=True, **kwargs):
    import nvtabular as nvt
    from nvtabular import ops as nvt_ops

    cat = lambda: nvt_ops.Categorify(dtype="int32")  # noqa: E731

    user_id = ["user_id"] >> cat() >> nvt_ops.TagAsUserID()
    item_id = ["item_id"] >> cat() >> nvt_ops.TagAsItemID()

    item_features = (
        ["item_category", "item_shop", "item_brand", "item_intention"]
        >> cat()
        >> nvt_ops.TagAsItemFeatures()
    )

    user_features = (
        [
            "user_shops",
            "user_profile",
            "user_group",
            "user_gender",
            "user_age",
            "user_consumption_1",
            "user_consumption_2",
            "user_is_occupied",
            "user_geography",
            "user_intentions",
            "user_brands",
            "user_categories",
        ]
        >> cat()
        >> nvt_ops.TagAsUserFeatures()
    )

    user_item_features = (
        [
            "user_item_categories",
            "user_item_shops",
            "user_item_brands",
            "user_item_intentions",
        ]
        >> cat()
        >> nvt_ops.AddMetadata(tags=["user_item"])
    )

    context_features = ["position"] >> cat() >> nvt_ops.AddMetadata(tags=[Tags.CONTEXT])

    targets = ["click", "conversion"] >> nvt_ops.AddMetadata(
        tags=[Tags.BINARY_CLASSIFICATION, "target"]
    )
    outputs = (
        user_id
        + item_id
        + item_features
        + user_features
        + user_item_features
        + context_features
        + targets
    )

    if add_target_encoding:
        continuous = (
            ["user_id", "item_id"]
            >> nvt_ops.TargetEncoding(["click"], kfold=1, p_smooth=20)
            >> nvt_ops.Normalize()
        )
        outputs += continuous

    return nvt.Workflow(outputs)


def transform_aliccp(
    data: Union[str, Path, Tuple[merlin.io.Dataset, merlin.io.Dataset]],
    output_path: Union[str, Path],
    nvt_workflow=None,
    **kwargs,
):
    nvt_workflow = nvt_workflow or default_aliccp_transformation(**locals())

    if isinstance(data, (str, Path)):
        _train = glob(str(data / "train/*.parquet"))
        _valid = glob(str(data / "test/*.parquet"))
    elif (
        isinstance(data, tuple)
        and len(data) == 2
        and all(isinstance(x, merlin.io.Dataset) for x in data)
    ):
        _train, _valid = data
    else:
        raise ValueError("data must be a path or a tuple of train and valid datasets")

    workflow_fit_transform(nvt_workflow, _train, _valid, str(output_path), **kwargs)


@dataclass
class _Feature:
    name: str
    id: str
    tags: List[Union[str, Tags]]
    description: str


class _Features:
    def __init__(self):
        self.features: List[_Feature] = [
            # User
            _Feature("user_id", "101", [Tags.USER_ID, Tags.USER], "User ID"),
            _Feature(
                "user_categories",
                "109_14",
                [Tags.USER],
                "User historical behaviors of category ID and count",
            ),
            _Feature(
                "user_shops",
                "110_14",
                [Tags.USER],
                "User historical behaviors of shop ID and count",
            ),
            _Feature(
                "user_brands",
                "127_14",
                [Tags.USER],
                "User historical behaviors of brand ID and count",
            ),
            _Feature(
                "user_intentions",
                "150_14",
                [Tags.USER],
                "User historical behaviors of intention node ID and count",
            ),
            _Feature("user_profile", "121", [Tags.USER], "Categorical ID of User Profile"),
            _Feature("user_group", "122", [Tags.USER], "Categorical group ID of User Profile"),
            _Feature("user_gender", "124", [Tags.USER], "Users Gender ID"),
            _Feature("user_age", "125", [Tags.USER], "Users Age ID"),
            _Feature("user_consumption_1", "126", [Tags.USER], "Users Consumption Level Type I"),
            _Feature("user_consumption_2", "127", [Tags.USER], "Users Consumption Level Type II"),
            _Feature(
                "user_is_occupied", "128", [Tags.USER], "Users Occupation: whether or not to work"
            ),
            _Feature("user_geography", "129", [Tags.USER], "Users Geography Information"),
            # Item
            _Feature("item_id", "205", [Tags.ITEM, Tags.ITEM_ID], "Item ID"),
            _Feature(
                "item_category", "206", [Tags.ITEM], "Category ID to which the item belongs to"
            ),
            _Feature("item_shop", "207", [Tags.ITEM], "Shop ID to which item belongs to"),
            _Feature(
                "item_intention", "210", [Tags.ITEM], "Intention node ID which the item belongs to"
            ),
            _Feature("item_brand", "216", [Tags.ITEM], "Brand ID of the item"),
            # User-Item
            _Feature(
                "user_item_categories",
                "508",
                ["user_item"],
                "The combination of features with 109_14 and 206",
            ),
            _Feature(
                "user_item_shops",
                "509",
                ["user_item"],
                "The combination of features with 110_14 and 207",
            ),
            _Feature(
                "user_item_brands",
                "702",
                ["user_item"],
                "The combination of features with 127_14 and 216",
            ),
            _Feature(
                "user_item_intentions",
                "853",
                ["user_item"],
                "The combination of features with 150_14 and 210",
            ),
            # Context
            _Feature("position", "301", [Tags.CONTEXT], "A categorical expression of position"),
        ]

    @property
    def by_id(self):
        return {feature.id: feature.name for feature in self.features}


def _convert_common_features(common_path, pickle_path=None):
    common = {}

    with open(common_path, "r") as common_features:
        for csv_line in tqdm(common_features, desc="Reading common features..."):
            line = csv_line.strip().split(",")
            kv = np.array(re.split("[]", line[2]))
            keys = kv[range(0, len(kv), 3)]
            values = kv[range(1, len(kv), 3)]
            common[line[0]] = dict(zip(keys, values))

        if pickle_path:
            with open(pickle_path, "wb") as handle:
                pickle.dump(common, handle, protocol=pickle.HIGHEST_PROTOCOL)

    return common


# TODO: Optimize this function, right now it's too slow.
def _convert_data(
    data_dir,
    file_size,
    is_train=True,
    pickle_common_features=True,
    max_num_rows=None,
    output_dir=None,
):
    data_type = "train" if is_train else "test"
    output_dir = output_dir or data_dir

    common_path = os.path.join(data_dir, data_type, f"common_features_{data_type}.csv")
    path = os.path.join(data_dir, data_type, f"sample_skeleton_{data_type}.csv")
    common_features_path = os.path.join(data_dir, data_type, "common_features.pickle")

    common = {}
    if pickle_common_features:
        if os.path.exists(common_features_path):
            with open(common_features_path, "rb") as f:
                common = pickle.load(f)

    if not common:
        pickle_path = common_features_path if pickle_common_features else None
        common = _convert_common_features(common_path, pickle_path)

    current = []
    by_id = _Features().by_id

    out_dir = os.path.join(str(output_dir), data_type)
    tmp_dir = os.path.join(out_dir, "tmp")
    if not os.path.exists(tmp_dir):
        os.makedirs(tmp_dir)

    with open(path, "r") as skeleton:
        for i, csv_line in tqdm(enumerate(skeleton), desc="Processing data..."):
            if max_num_rows and i >= max_num_rows:
                break

            line = csv_line.strip().split(",")
            if line[1] == "0" and line[2] == "1":
                continue
            kv = np.array(re.split("[]", line[5]))
            key = kv[range(0, len(kv), 3)]
            value = kv[range(1, len(kv), 3)]
            feat_dict = dict(zip(key, value))
            feat_dict.update(common[line[3]])
            feat_dict["click"] = int(line[1])
            feat_dict["conversion"] = int(line[2])

            current.append(feat_dict)

            if i > 0 and i % file_size == 0:
                df = get_lib().DataFrame(current)
                cols = []
                for col in list(df.columns):
                    if col == "click" or col == "conversion":
                        cols.append(col)
                    else:
                        cols.append(by_id[col])

                df.columns = cols
                index = int((i / file_size) - 1)
                df.to_parquet(
                    os.path.join(tmp_dir, f"{data_type}_{index}.parquet"),
                    overwrite=True,
                )
                current = []

    tmp_files = glob(os.path.join(tmp_dir, f"{data_type}_*.parquet"))
    dtypes = {f.name: "int32" for f in _Features().features}
    merlin.io.Dataset(tmp_files, dtypes=dtypes).to_parquet(out_dir)
    shutil.rmtree(tmp_dir)
