#
# 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 glob
import os
from pathlib import Path
from typing import Optional, Tuple, Union

import numpy as np

import merlin.io
from merlin.core.utils import download_file
from merlin.datasets import BASE_PATH
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_criteo(
    path: Union[str, Path] = None,
    num_days: int = 2,
    overwrite: bool = False,
    transformed_name: str = "transformed",
    nvt_workflow: Optional[Workflow] = None,
    valid_num_days: int = 1,
    **kwargs,
) -> Tuple[merlin.io.Dataset, merlin.io.Dataset]:
    """
    Get the Criteo 1TB Click Logs dataset, more info can be found
    [here](https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/).

    This dataset contains feature values and click feedback for millions of display
    ads. Its purpose is to benchmark algorithms for clickthrough rate (CTR) prediction.
    It is similar, but larger, to the dataset released for the Display Advertising
    Challenge hosted by Kaggle: https://www.kaggle.com/c/criteo-display-ad-challenge

    This function will download the data if not already present.
    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("criteo")

    Parameters
    ----------
    path: str or Path, optional
        Path to the dataset. If not set will default to
        the 'merlin-models-data` directory in your home folder
    num_days: int, optional
        Number of days to download, should be between 2 & 23. Default is 2.
        The last day will be used for the valid set, the rest for train.
    overwrite: bool, optional
        Whether to overwrite the transformed data.
    transformed_name: str, optional
        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_criteo_transform`.
    valid_num_days: int.
        Number of days to use for the validation set. This will be the last day(s).
        Default is 1. This is similar to all papers published on the Criteo dataset.

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

    require_nvt()

    if path is None:
        p = Path(BASE_PATH) / "criteo"
    else:
        p = Path(path)

    raw_data_path = p / "orig"
    maybe_download(raw_data_path, num_days)

    variant_path = p / f"{num_days}-days"
    raw_path = variant_path / "raw"
    if not variant_path.exists():
        raw_path.mkdir(parents=True)
        prepare_criteo(raw_data_path, raw_path, num_days=num_days)

    nvt_path = variant_path / 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_criteo(
            raw_path, variant_path, workflow=nvt_workflow, valid_num_days=valid_num_days, **kwargs
        )

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

    return train, valid


def maybe_download(destination: Union[str, Path], num_days: int):
    if num_days < 2 or num_days > 23:
        raise ValueError(
            str(num_days)
            + " is not supported. A minimum of 2 days are "
            + "required and a maximum of 24 (0-23 days) are available"
        )

    # Create input dir if not exists
    if not os.path.exists(destination):
        os.makedirs(destination)

    # Iterate over days
    for i in range(0, num_days):
        file = os.path.join(destination, "day_" + str(i) + ".gz")
        # Download file, if there is no .gz, .csv or .parquet file
        if not (
            os.path.exists(file)
            or os.path.exists(
                file.replace(".gz", ".parquet").replace("crit_orig", "converted/criteo/")
            )
            or os.path.exists(file.replace(".gz", ""))
        ):
            download_file(
                "https://storage.googleapis.com/criteo-cail-datasets/day_" + str(i) + ".gz",
                file,
            )


def prepare_criteo(
    raw_data_path: Union[str, Path],
    output_path: Union[str, Path],
    num_days: int,
    part_mem_fraction=0.1,
    **kwargs,
):
    """Convert the raw data to parquet format.

    Parameters
    ----------
    raw_data_path: str or Path
        Path to the raw data.
    output_path: str or Path
        Path to the output directory.
    num_days: int
        Number of days to use.
    part_mem_fraction: float, optional
        Fraction of memory to use for the conversion.
    client:

    Returns
    -------
    output_path: str or Path
    """

    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # Specify column names
    cont_names = ["I" + str(x) for x in range(1, 14)]
    cat_names = ["C" + str(x) for x in range(1, 27)]
    cols = ["label"] + cont_names + cat_names

    # Specify column dtypes. Note that "hex" means that
    # the values will be hexadecimal strings that should
    # be converted to int32
    dtypes = {}
    dtypes["label"] = np.int32
    for x in cont_names:
        dtypes[x] = np.int32
    for x in cat_names:
        dtypes[x] = "hex"

    # Create an NVTabular Dataset from a CSV-file glob
    file_list = glob.glob(os.path.join(str(raw_data_path), "day_*[!.gz]"))
    file_list = sorted(file_list)[:num_days]
    dataset = nvt.Dataset(
        file_list,
        engine="csv",
        names=cols,
        part_mem_fraction=part_mem_fraction,
        sep="\t",
        dtypes=dtypes,
        **kwargs,
    )

    dataset.to_parquet(str(output_path), preserve_files=True)

    return output_path


def default_criteo_transform(num_buckets=10000000, **kwargs) -> Workflow:
    from nvtabular import ops as nvt_ops

    continuous = ["I" + str(x) for x in range(1, 14)]
    categorical = ["C" + str(x) for x in list(range(1, 21)) + list(range(22, 27))]
    # It's a reasonable guess that C21 is the item-id col.
    # This is only used to generate synthetic data
    item_id = ["C21"] >> nvt_ops.AddMetadata(tags=[Tags.ITEM_ID])
    targets = ["label"] >> nvt_ops.AddMetadata(tags=["target", Tags.BINARY_CLASSIFICATION])

    categorify_op = nvt_ops.Categorify(max_size=num_buckets)
    cat_features = categorical + item_id >> categorify_op
    cont_features = (
        continuous >> nvt_ops.FillMissing() >> nvt_ops.Clip(min_value=0) >> nvt_ops.Normalize()
    )
    features = cat_features + cont_features + targets

    workflow = nvt.Workflow(features)

    return workflow


def transform_criteo(
    raw_data_path: Union[str, Path],
    output_path: Union[str, Path],
    workflow: Optional[Workflow] = None,
    num_buckets=10000000,
    valid_num_days=1,
    **kwargs,
):
    workflow = workflow or default_criteo_transform(**locals())
    file_list = glob.glob(os.path.join(str(raw_data_path), "day_*[!.gz]"))
    workflow_fit_transform(
        workflow,
        file_list[:-valid_num_days],
        file_list[-valid_num_days],
        str(output_path),
    )

    return output_path
