# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the BSD 3-Clause License  (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# 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.
#
# Description: generate inputs and targets for the DLRM benchmark
#
# Utility function(s) to download and pre-process public data sets
#   - Criteo Kaggle Display Advertising Challenge Dataset
#     https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset
#   - Criteo Terabyte Dataset
#     https://labs.criteo.com/2013/12/download-terabyte-click-logs
#
# After downloading dataset, run:
#   getCriteoAdData(
#       datafile="<path-to-train.txt>",
#       o_filename=kaggleAdDisplayChallenge_processed.npz,
#       max_ind_range=-1,
#       sub_sample_rate=0.0,
#       days=7,
#       data_split='train',
#       randomize='total',
#       criteo_kaggle=True,
#       memory_map=False
#   )
#   getCriteoAdData(
#       datafile="<path-to-day_{0,...,23}>",
#       o_filename=terabyte_processed.npz,
#       max_ind_range=-1,
#       sub_sample_rate=0.0,
#       days=24,
#       data_split='train',
#       randomize='total',
#       criteo_kaggle=False,
#       memory_map=False
#   )

from __future__ import absolute_import, division, print_function, unicode_literals

import sys
# import os
from os import path
from multiprocessing import Process, Manager
# import io
# from io import StringIO
# import collections as coll

import numpy as np


def convertUStringToDistinctIntsDict(mat, convertDicts, counts):
    # Converts matrix of unicode strings into distinct integers.
    #
    # Inputs:
    #     mat (np.array): array of unicode strings to convert
    #     convertDicts (list): dictionary for each column
    #     counts (list): number of different categories in each column
    #
    # Outputs:
    #     out (np.array): array of output integers
    #     convertDicts (list): dictionary for each column
    #     counts (list): number of different categories in each column

    # check if convertDicts and counts match correct length of mat
    if len(convertDicts) != mat.shape[1] or len(counts) != mat.shape[1]:
        print("Length of convertDicts or counts does not match input shape")
        print("Generating convertDicts and counts...")

        convertDicts = [{} for _ in range(mat.shape[1])]
        counts = [0 for _ in range(mat.shape[1])]

    # initialize output
    out = np.zeros(mat.shape)

    for j in range(mat.shape[1]):
        for i in range(mat.shape[0]):
            # add to convertDict and increment count
            if mat[i, j] not in convertDicts[j]:
                convertDicts[j][mat[i, j]] = counts[j]
                counts[j] += 1
            out[i, j] = convertDicts[j][mat[i, j]]

    return out, convertDicts, counts


def convertUStringToDistinctIntsUnique(mat, mat_uni, counts):
    # mat is an array of 0,...,# samples, with each being 26 categorical features

    # check if mat_unique and counts match correct length of mat
    if len(mat_uni) != mat.shape[1] or len(counts) != mat.shape[1]:
        print("Length of mat_unique or counts does not match input shape")
        print("Generating mat_unique and counts...")

        mat_uni = [np.array([]) for _ in range(mat.shape[1])]
        counts = [0 for _ in range(mat.shape[1])]

    # initialize output
    out = np.zeros(mat.shape)
    ind_map = [np.array([]) for _ in range(mat.shape[1])]

    # find out and assign unique ids to features
    for j in range(mat.shape[1]):
        m = mat_uni[j].size
        mat_concat = np.concatenate((mat_uni[j], mat[:, j]))
        mat_uni[j], ind_map[j] = np.unique(mat_concat, return_inverse=True)
        out[:, j] = ind_map[j][m:]
        counts[j] = mat_uni[j].size

    return out, mat_uni, counts


def processCriteoAdData(d_path, d_file, npzfile, i, convertDicts, pre_comp_counts):
    # Process Kaggle Display Advertising Challenge or Terabyte Dataset
    # by converting unicode strings in X_cat to integers and
    # converting negative integer values in X_int.
    #
    # Loads data in the form "{kaggle|terabyte}_day_i.npz" where i is the day.
    #
    # Inputs:
    #   d_path (str): path for {kaggle|terabyte}_day_i.npz files
    #   i (int): splits in the dataset (typically 0 to 7 or 0 to 24)

    # process data if not all files exist
    filename_i = npzfile + "_{0}_processed.npz".format(i)

    if path.exists(filename_i):
        print("Using existing " + filename_i, end="\n")
    else:
        print("Not existing " + filename_i)
        with np.load(npzfile + "_{0}.npz".format(i)) as data:
            # categorical features
            '''
            # Approach 1a: using empty dictionaries
            X_cat, convertDicts, counts = convertUStringToDistinctIntsDict(
                data["X_cat"], convertDicts, counts
            )
            '''
            '''
            # Approach 1b: using empty np.unique
            X_cat, convertDicts, counts = convertUStringToDistinctIntsUnique(
                data["X_cat"], convertDicts, counts
            )
            '''
            # Approach 2a: using pre-computed dictionaries
            X_cat_t = np.zeros(data["X_cat_t"].shape)
            for j in range(26):
                for k, x in enumerate(data["X_cat_t"][j, :]):
                    X_cat_t[j, k] = convertDicts[j][x]
            # continuous features
            X_int = data["X_int"]
            X_int[X_int < 0] = 0
            # targets
            y = data["y"]

        np.savez_compressed(
            filename_i,
            # X_cat = X_cat,
            X_cat=np.transpose(X_cat_t),  # transpose of the data
            X_int=X_int,
            y=y,
        )
        print("Processed " + filename_i, end="\n")
    # sanity check (applicable only if counts have been pre-computed & are re-computed)
    # for j in range(26):
    #    if pre_comp_counts[j] != counts[j]:
    #        sys.exit("ERROR: Sanity check on counts has failed")
    # print("\nSanity check on counts passed")

    return


def concatCriteoAdData(
        d_path,
        d_file,
        npzfile,
        trafile,
        days,
        data_split,
        randomize,
        total_per_file,
        total_count,
        memory_map,
        o_filename
):
    # Concatenates different days and saves the result.
    #
    # Inputs:
    #   days (int): total number of days in the dataset (typically 7 or 24)
    #   d_path (str): path for {kaggle|terabyte}_day_i.npz files
    #   o_filename (str): output file name
    #
    # Output:
    #   o_file (str): output file path

    if memory_map:
        # dataset break up per fea
        # tar_fea = 1   # single target
        den_fea = 13  # 13 dense  features
        spa_fea = 26  # 26 sparse features
        # tad_fea = tar_fea + den_fea
        # tot_fea = tad_fea + spa_fea
        # create offset per file
        offset_per_file = np.array([0] + [x for x in total_per_file])
        for i in range(days):
            offset_per_file[i + 1] += offset_per_file[i]

        '''
        # Approach 1, 2 and 3 use indices, while Approach 4 does not use them
        # create indices
        indices = np.arange(total_count)
        if data_split == "none":
            if randomize == "total":
                indices = np.random.permutation(indices)
        else:
            indices = np.array_split(indices, offset_per_file[1:-1])

            # randomize train data (per day)
            if randomize == "day":  # or randomize == "total":
                for i in range(len(indices) - 1):
                    indices[i] = np.random.permutation(indices[i])
                print("Randomized indices per day ...")

            train_indices = np.concatenate(indices[:-1])
            test_indices = indices[-1]

            # randomize train data (across days)
            if randomize == "total":
                train_indices = np.random.permutation(train_indices)
                print("Randomized indices across days ...")

            indices = np.concatenate((train_indices, test_indices))
        # no reordering
        # indices = np.arange(total_count)
        '''
        '''
        # Approach 1: simple and slow (no grouping is used)
        # check if data already exists
        recreate_flag = False
        for j in range(tot_fea):
            filename_j = trafile + "_{0}_reordered.npy".format(j)
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # load, reorder and concatenate data (memmap all reordered files per feature)
        if recreate_flag:
            # init reordered files (.npy appended automatically)
            z = np.zeros((total_count))
            for j in range(tot_fea):
                filename_j = trafile + "_{0}_reordered".format(j)
                np.save(filename_j, z)
                print("Creating " + filename_j)

            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
                with np.load(filename_i) as data:
                    X_cat_t = np.transpose(data["X_cat"])
                    X_int_t = np.transpose(data["X_int"])
                    y = data["y"]
                size = len(y)
                # sanity check
                if total_per_file[i] != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                # print(filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #     + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r+')
                    if j < tar_fea:
                        fj[indices[start:end]] = y
                    elif tar_fea <= j and j < tad_fea:
                        fj[indices[start:end]] = X_int_t[j - tar_fea, :]
                    else:
                        fj[indices[start:end]] = X_cat_t[j - tad_fea, :]
                    del fj
        else:
            print("Reordered fea files already exist, skipping ...")

        # check if data already exists
        recreate_flag = False
        for i in range(days):
            filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if path.exists(filename_i):
                print("Using existing " + filename_i)
            else:
                recreate_flag = True
        # split reordered data by files (memmap all reordered files per feature)
        # on the day boundary del the file object and memmap again
        if recreate_flag:
            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
                size = total_per_file[i]
                X_int_t = np.zeros((den_fea, size))
                X_cat_t = np.zeros((spa_fea, size))
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                print("Creating " + filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #     + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r')
                    if j < tar_fea:
                        y = fj[start:end]
                    elif tar_fea <= j and j < tad_fea:
                        X_int_t[j - tar_fea, :] = fj[start:end]
                    else:
                        X_cat_t[j - tad_fea, :] = fj[start:end]
                    del fj

                np.savez_compressed(
                    filename_i,
                    X_cat=np.transpose(X_cat_t),  # transpose of the data
                    X_int=np.transpose(X_int_t),  # transpose of the data
                    y=y,
                )
        else:
            print("Reordered day files already exist, skipping ...")
        '''
        '''
        # Approach 2: group days
        # check if data already exists
        recreate_flag = False
        for j in range(tot_fea):
            filename_j = trafile + "_{0}_reordered.npy".format(j)
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # load, reorder and concatenate data (memmap all reordered files per feature)
        if recreate_flag:
            # init reordered files (.npy appended automatically)
            z = np.zeros((total_count))
            for j in range(tot_fea):
                filename_j = trafile + "_{0}_reordered".format(j)
                np.save(filename_j, z)
                print("Creating " + filename_j)

            group_day = 3  # e.g. 8, 4 or 3
            group_num = days // group_day
            file_group = [i*group_day for i in range(group_num)] + [days]
            for ii in range(group_num):
                # for last may be group_size != group_num, therefore reset it below
                group_size = file_group[ii + 1] - file_group[ii]
                X_cat_t = [0]*group_size
                X_int_t = [0]*group_size
                y = [0]*group_size
                start = [0]*group_size
                end  = [0]*group_size
                for ig in range(group_size):
                    i = file_group[ii] + ig
                    filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
                    # setup start and end ranges
                    start[ig] = offset_per_file[i]
                    end[ig] = offset_per_file[i + 1]
                    # print(filename_i)
                    # load a group of files
                    with np.load(filename_i) as data:
                        X_cat_t[ig] = np.transpose(data["X_cat"])
                        X_int_t[ig] = np.transpose(data["X_int"])
                        y[ig] = data["y"]
                    # sanity check
                    if total_per_file[i] != len(y[ig]):
                        sys.exit("ERROR: sanity check on number of samples failed")
                # print("start=" + str(start) + " end=" + str(end)
                #  + " diff=" + str(end[ig]-start[ig]) + "=" + str(total_per_file[i]))

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r+')
                    for ig in range(group_size):
                        if j < tar_fea:
                            fj[indices[start[ig]:end[ig]]] = y[ig]
                        elif tar_fea <= j and j < tad_fea:
                            fj[indices[start[ig]:end[ig]]] = X_int_t[ig][j - tar_fea, :]
                        else:
                            fj[indices[start[ig]:end[ig]]] = X_cat_t[ig][j - tad_fea, :]
                    del fj
        else:
            print("Reordered fea files already exist, skipping ...")

        # check if data already exists
        recreate_flag = False
        for i in range(days):
            filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if path.exists(filename_i):
                print("Using existing " + filename_i)
            else:
                recreate_flag = True
        # split reordered data by files (memmap all reordered files per feature)
        # on the day boundary del the file object and memmap again
        if recreate_flag:
            for ii in range(group_num):
                # for last may be group_size != group_num, therefore reset it below
                group_size = file_group[ii + 1] - file_group[ii]
                X_cat_t= []; X_int_t = []
                for ig in range(group_size):
                    i = file_group[ii] + ig
                    X_int_t.append(np.zeros((den_fea, total_per_file[i])))
                    X_cat_t.append(np.zeros((spa_fea, total_per_file[i])))
                y = [0]*group_size
                start = [0]*group_size
                end  = [0]*group_size

                for j in range(tot_fea):
                    filename_j = trafile + "_{0}_reordered.npy".format(j)
                    fj = np.load(filename_j, mmap_mode='r')
                    # load a group of files
                    for ig in range(group_size):
                        i = file_group[ii] + ig
                        # setup start and end ranges
                        start[ig] = offset_per_file[i]
                        end[ig] = offset_per_file[i + 1]
                        # load data for the group of files
                        if j < tar_fea:
                            y[ig] = fj[start[ig]:end[ig]]
                        elif tar_fea <= j and j < tad_fea:
                            X_int_t[ig][j - tar_fea, :] = fj[start[ig]:end[ig]]
                        else:
                            X_cat_t[ig][j - tad_fea, :] = fj[start[ig]:end[ig]]
                    del fj

                for ig in range(group_size):
                    i = file_group[ii] + ig
                    filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
                    print("Creating " + filename_i)
                    np.savez_compressed(
                        filename_i,
                        X_cat=np.transpose(X_cat_t[ig]),  # transpose of the data
                        X_int=np.transpose(X_int_t[ig]),  # transpose of the data
                        y=y[ig],
                    )
        else:
            print("Reordered day files already exist, skipping ...")
        '''
        '''
        # Approach 3: group features
        # check if data already exists
        group_fea = 5  # e.g. 8, 5 or 4
        group_num = tot_fea // group_fea
        if tot_fea % group_fea != 0:  # sanity check
            sys.exit("ERROR: the group_fea must divided tot_fea evenly.")
        recreate_flag = False
        for jn in range(group_num):
            filename_j = trafile + "_{0}_reordered{1}.npy".format(
                jn, group_fea
            )
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # load, reorder and concatenate data (memmap all reordered files per feature)
        if recreate_flag:
            # init reordered files (.npy appended automatically)
            z = np.zeros((group_fea, total_count))
            for jn in range(group_num):
                filename_j = trafile + "_{0}_reordered{1}".format(
                    jn, group_fea
                )
                np.save(filename_j, z)
                print("Creating " + filename_j)

            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
                with np.load(filename_i) as data:
                    X_cat_t = np.transpose(data["X_cat"])
                    X_int_t = np.transpose(data["X_int"])
                    y = data["y"]
                size = len(y)
                # sanity check
                if total_per_file[i] != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                # print(filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #      + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for jn in range(group_num):
                    filename_j = trafile + "_{0}_reordered{1}.npy".format(
                        jn, group_fea
                    )
                    fj = np.load(filename_j, mmap_mode='r+')
                    for jg in range(group_fea):
                        j = jn * group_fea + jg
                        # print("j=" + str(j) + " jn=" + str(jn) + " jg=" + str(jg))
                        if j < tar_fea:
                            fj[jg, indices[start:end]] = y
                        elif tar_fea <= j and j < tad_fea:
                            fj[jg, indices[start:end]] = X_int_t[j - tar_fea, :]
                        else:
                            fj[jg, indices[start:end]] = X_cat_t[j - tad_fea, :]
                    del fj
        else:
            print("Reordered fea files already exist, skipping ...")

        # check if data already exists
        recreate_flag = False
        for i in range(days):
            filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if path.exists(filename_i):
                print("Using existing" + filename_i)
            else:
                recreate_flag = True
        # split reordered data by files (memmap all reordered files per feature)
        # on the day boundary del the file object and memmap again
        if recreate_flag:
            for i in range(days):
                filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
                size = total_per_file[i]
                X_int_t = np.zeros((den_fea, size))
                X_cat_t = np.zeros((spa_fea, size))
                # setup start and end ranges
                start = offset_per_file[i]
                end = offset_per_file[i + 1]
                print("Creating " + filename_i)
                # print("start=" + str(start) + " end=" + str(end)
                #      + " diff=" + str(end - start) + "=" + str(total_per_file[i]))

                for jn in range(group_num):
                    filename_j = trafile + "_{0}_reordered{1}.npy".format(
                        jn, group_fea
                    )
                    fj = np.load(filename_j, mmap_mode='r')
                    for jg in range(group_fea):
                        j = jn * group_fea + jg
                        # print("j=" + str(j) + " jn=" + str(jn) + " jg=" + str(jg))
                        if j < tar_fea:
                            y = fj[jg, start:end]
                        elif tar_fea <= j and j < tad_fea:
                            X_int_t[j - tar_fea, :] = fj[jg, start:end]
                        else:
                            X_cat_t[j - tad_fea, :] = fj[jg, start:end]
                    del fj

                np.savez_compressed(
                    filename_i,
                    X_cat=np.transpose(X_cat_t),  # transpose of the data
                    X_int=np.transpose(X_int_t),  # transpose of the data
                    y=y,
                )

        else:
            print("Reordered day files already exist, skipping ...")
        '''

        # Approach 4: Fisher-Yates-Rao (FYR) shuffle algorithm
        # 1st pass of FYR shuffle
        # check if data already exists
        recreate_flag = False
        for j in range(days):
            filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
            filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
            filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
            if (
                path.exists(filename_j_y)
                and path.exists(filename_j_d)
                and path.exists(filename_j_s)
            ):
                print(
                    "Using existing\n"
                    + filename_j_y + "\n"
                    + filename_j_d + "\n"
                    + filename_j_s
                )
            else:
                recreate_flag = True
        # reorder across buckets using sampling
        if recreate_flag:
            # init intermediate files (.npy appended automatically)
            for j in range(days):
                filename_j_y = npzfile + "_{0}_intermediate_y".format(j)
                filename_j_d = npzfile + "_{0}_intermediate_d".format(j)
                filename_j_s = npzfile + "_{0}_intermediate_s".format(j)
                np.save(filename_j_y, np.zeros((total_per_file[j])))
                np.save(filename_j_d, np.zeros((total_per_file[j], den_fea)))
                np.save(filename_j_s, np.zeros((total_per_file[j], spa_fea)))
            # start processing files
            total_counter = [0] * days
            for i in range(days):
                filename_i = npzfile + "_{0}_processed.npz".format(i)
                with np.load(filename_i) as data:
                    X_cat = data["X_cat"]
                    X_int = data["X_int"]
                    y = data["y"]
                size = len(y)
                # sanity check
                if total_per_file[i] != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # debug prints
                print("Reordering (1st pass) " + filename_i)

                # create buckets using sampling of random ints
                # from (discrete) uniform distribution
                buckets = []
                for _j in range(days):
                    buckets.append([])
                counter = [0] * days
                days_to_sample = days if data_split == "none" else days - 1
                if randomize == "total":
                    rand_u = np.random.randint(low=0, high=days_to_sample, size=size)
                    for k in range(size):
                        # sample and make sure elements per buckets do not overflow
                        if data_split == "none" or i < days - 1:
                            # choose bucket
                            p = rand_u[k]
                            # retry of the bucket is full
                            while total_counter[p] + counter[p] >= total_per_file[p]:
                                p = np.random.randint(low=0, high=days_to_sample)
                        else:  # preserve the last day/bucket if needed
                            p = i
                        buckets[p].append(k)
                        counter[p] += 1
                else:  # randomize is day or none
                    for k in range(size):
                        # do not sample, preserve the data in this bucket
                        p = i
                        buckets[p].append(k)
                        counter[p] += 1

                # sanity check
                if np.sum(counter) != size:
                    sys.exit("ERROR: sanity check on number of samples failed")
                # debug prints
                # print(counter)
                # print(str(np.sum(counter)) + " = " + str(size))
                # print([len(x) for x in buckets])
                # print(total_counter)

                # partially feel the buckets
                for j in range(days):
                    filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
                    filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
                    filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
                    start = total_counter[j]
                    end = total_counter[j] + counter[j]
                    # target buckets
                    fj_y = np.load(filename_j_y, mmap_mode='r+')
                    # print("start=" + str(start) + " end=" + str(end)
                    #       + " end - start=" + str(end - start) + " "
                    #       + str(fj_y[start:end].shape) + " "
                    #       + str(len(buckets[j])))
                    fj_y[start:end] = y[buckets[j]]
                    del fj_y
                    # dense buckets
                    fj_d = np.load(filename_j_d, mmap_mode='r+')
                    # print("start=" + str(start) + " end=" + str(end)
                    #       + " end - start=" + str(end - start) + " "
                    #       + str(fj_d[start:end, :].shape) + " "
                    #       + str(len(buckets[j])))
                    fj_d[start:end, :] = X_int[buckets[j], :]
                    del fj_d
                    # sparse buckets
                    fj_s = np.load(filename_j_s, mmap_mode='r+')
                    # print("start=" + str(start) + " end=" + str(end)
                    #       + " end - start=" + str(end - start) + " "
                    #       + str(fj_s[start:end, :].shape) + " "
                    #       + str(len(buckets[j])))
                    fj_s[start:end, :] = X_cat[buckets[j], :]
                    del fj_s
                    # update counters for next step
                    total_counter[j] += counter[j]

        # 2nd pass of FYR shuffle
        # check if data already exists
        for j in range(days):
            filename_j = npzfile + "_{0}_reordered.npz".format(j)
            if path.exists(filename_j):
                print("Using existing " + filename_j)
            else:
                recreate_flag = True
        # reorder within buckets
        if recreate_flag:
            for j in range(days):
                filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
                filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
                filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
                fj_y = np.load(filename_j_y)
                fj_d = np.load(filename_j_d)
                fj_s = np.load(filename_j_s)

                indices = range(total_per_file[j])
                if randomize == "day" or randomize == "total":
                    if data_split == "none" or j < days - 1:
                        indices = np.random.permutation(range(total_per_file[j]))

                filename_r = npzfile + "_{0}_reordered.npz".format(j)
                print("Reordering (2nd pass) " + filename_r)
                np.savez_compressed(
                    filename_r,
                    X_cat=fj_s[indices, :],
                    X_int=fj_d[indices, :],
                    y=fj_y[indices],
                )

        '''
        # sanity check (under no reordering norms should be zero)
        for i in range(days):
            filename_i_o = npzfile + "_{0}_processed.npz".format(i)
            print(filename_i_o)
            with np.load(filename_i_o) as data_original:
                X_cat_o = data_original["X_cat"]
                X_int_o = data_original["X_int"]
                y_o = data_original["y"]
            filename_i_r = npzfile + "_{0}_reordered.npz".format(i)
            print(filename_i_r)
            with np.load(filename_i_r) as data_reordered:
                X_cat_r = data_reordered["X_cat"]
                X_int_r = data_reordered["X_int"]
                y_r = data_reordered["y"]
            print(np.linalg.norm(y_o - y_r))
            print(np.linalg.norm(X_int_o - X_int_r))
            print(np.linalg.norm(X_cat_o - X_cat_r))
        '''

    else:
        print("Concatenating multiple days into %s.npz file" % str(d_path + o_filename))

        # load and concatenate data
        for i in range(days):
            filename_i = npzfile + "_{0}_processed.npz".format(i)
            with np.load(filename_i) as data:
                if i == 0:
                    X_cat = data["X_cat"]
                    X_int = data["X_int"]
                    y = data["y"]
                else:
                    X_cat = np.concatenate((X_cat, data["X_cat"]))
                    X_int = np.concatenate((X_int, data["X_int"]))
                    y = np.concatenate((y, data["y"]))
            print("Loaded day:", i, "y = 1:", len(y[y == 1]), "y = 0:", len(y[y == 0]))

        with np.load(d_path + d_file + "_fea_count.npz") as data:
            counts = data["counts"]
        print("Loaded counts!")

        np.savez_compressed(
            d_path + o_filename + ".npz",
            X_cat=X_cat,
            X_int=X_int,
            y=y,
            counts=counts,
        )

    return d_path + o_filename + ".npz"


def transformCriteoAdData(X_cat, X_int, y, days, data_split, randomize, total_per_file):
    # Transforms Criteo Kaggle or terabyte data by applying log transformation
    # on dense features and converting everything to appropriate tensors.
    #
    # Inputs:
    #     X_cat (ndarray): array of integers corresponding to preprocessed
    #                      categorical features
    #     X_int (ndarray): array of integers corresponding to dense features
    #     y (ndarray):     array of bool corresponding to labels
    #     data_split(str): flag for splitting dataset into training/validation/test
    #                      sets
    #     randomize (str): determines randomization scheme
    #         "none": no randomization
    #         "day": randomizes each day"s data (only works if split = True)
    #         "total": randomizes total dataset
    #
    # Outputs:
    #     if split:
    #         X_cat_train (tensor): sparse features for training set
    #         X_int_train (tensor): dense features for training set
    #         y_train (tensor): labels for training set
    #         X_cat_val (tensor): sparse features for validation set
    #         X_int_val (tensor): dense features for validation set
    #         y_val (tensor): labels for validation set
    #         X_cat_test (tensor): sparse features for test set
    #         X_int_test (tensor): dense features for test set
    #         y_test (tensor): labels for test set
    #     else:
    #         X_cat (tensor): sparse features
    #         X_int (tensor): dense features
    #         y (tensor): label

    # define initial set of indices
    indices = np.arange(len(y))

    # create offset per file
    offset_per_file = np.array([0] + [x for x in total_per_file])
    for i in range(days):
        offset_per_file[i + 1] += offset_per_file[i]

    # split dataset
    if data_split == 'train':
        indices = np.array_split(indices, offset_per_file[1:-1])

        # randomize train data (per day)
        if randomize == "day":  # or randomize == "total":
            for i in range(len(indices) - 1):
                indices[i] = np.random.permutation(indices[i])
            print("Randomized indices per day ...")

        train_indices = np.concatenate(indices[:-1])
        test_indices = indices[-1]
        test_indices, val_indices = np.array_split(test_indices, 2)

        print("Defined training and testing indices...")

        # randomize train data (across days)
        if randomize == "total":
            train_indices = np.random.permutation(train_indices)
            print("Randomized indices across days ...")

        # indices = np.concatenate((train_indices, test_indices))

        # create training, validation, and test sets
        X_cat_train = X_cat[train_indices]
        X_int_train = X_int[train_indices]
        y_train = y[train_indices]

        X_cat_val = X_cat[val_indices]
        X_int_val = X_int[val_indices]
        y_val = y[val_indices]

        X_cat_test = X_cat[test_indices]
        X_int_test = X_int[test_indices]
        y_test = y[test_indices]

        print("Split data according to indices...")

        X_cat_train = X_cat_train.astype(np.long)
        X_int_train = np.log(X_int_train.astype(np.float32) + 1)
        y_train = y_train.astype(np.float32)

        X_cat_val = X_cat_val.astype(np.long)
        X_int_val = np.log(X_int_val.astype(np.float32) + 1)
        y_val = y_val.astype(np.float32)

        X_cat_test = X_cat_test.astype(np.long)
        X_int_test = np.log(X_int_test.astype(np.float32) + 1)
        y_test = y_test.astype(np.float32)

        print("Converted to tensors...done!")

        return (
            X_cat_train,
            X_int_train,
            y_train,
            X_cat_val,
            X_int_val,
            y_val,
            X_cat_test,
            X_int_test,
            y_test,
        )

    else:

        # randomize data
        if randomize == "total":
            indices = np.random.permutation(indices)
            print("Randomized indices...")

        X_cat = X_cat[indices].astype(np.long)
        X_int = np.log(X_int[indices].astype(np.float32) + 1)
        y = y[indices].astype(np.float32)

        print("Converted to tensors...done!")

        return (X_cat, X_int, y, [], [], [], [], [], [])


def getCriteoAdData(
        datafile,
        o_filename,
        max_ind_range=-1,
        sub_sample_rate=0.0,
        days=7,
        data_split='train',
        randomize='total',
        criteo_kaggle=True,
        memory_map=False,
        dataset_multiprocessing=False,
):
    # Passes through entire dataset and defines dictionaries for categorical
    # features and determines the number of total categories.
    #
    # Inputs:
    #    datafile : path to downloaded raw data file
    #    o_filename (str): saves results under o_filename if filename is not ""
    #
    # Output:
    #   o_file (str): output file path

    #split the datafile into path and filename
    lstr = datafile.split("/")
    d_path = "/".join(lstr[0:-1]) + "/"
    d_file = lstr[-1].split(".")[0] if criteo_kaggle else lstr[-1]
    npzfile = d_path + ((d_file + "_day") if criteo_kaggle else d_file)
    trafile = d_path + ((d_file + "_fea") if criteo_kaggle else "fea")

    # count number of datapoints in training set
    total_file = d_path + d_file + "_day_count.npz"
    if path.exists(total_file):
        with np.load(total_file) as data:
            total_per_file = list(data["total_per_file"])
        total_count = np.sum(total_per_file)
        print("Skipping counts per file (already exist)")
    else:
        total_count = 0
        total_per_file = []
        if criteo_kaggle:
            # WARNING: The raw data consists of a single train.txt file
            # Each line in the file is a sample, consisting of 13 continuous and
            # 26 categorical features (an extra space indicates that feature is
            # missing and will be interpreted as 0).
            if path.exists(datafile):
                print("Reading data from path=%s" % (datafile))
                with open(str(datafile)) as f:
                    for _ in f:
                        total_count += 1
                total_per_file.append(total_count)
                # reset total per file due to split
                num_data_per_split, extras = divmod(total_count, days)
                total_per_file = [num_data_per_split] * days
                for j in range(extras):
                    total_per_file[j] += 1
                # split into days (simplifies code later on)
                file_id = 0
                boundary = total_per_file[file_id]
                nf = open(npzfile + "_" + str(file_id), "w")
                with open(str(datafile)) as f:
                    for j, line in enumerate(f):
                        if j == boundary:
                            nf.close()
                            file_id += 1
                            nf = open(npzfile + "_" + str(file_id), "w")
                            boundary += total_per_file[file_id]
                        nf.write(line)
                nf.close()
            else:
                sys.exit("ERROR: Criteo Kaggle Display Ad Challenge Dataset path is invalid; please download from https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset")
        else:
            # WARNING: The raw data consist of day_0.gz,... ,day_23.gz text files
            # Each line in the file is a sample, consisting of 13 continuous and
            # 26 categorical features (an extra space indicates that feature is
            # missing and will be interpreted as 0).
            for i in range(days):
                datafile_i = datafile + "_" + str(i)  # + ".gz"
                if path.exists(str(datafile_i)):
                    print("Reading data from path=%s" % (str(datafile_i)))
                    # file day_<number>
                    total_per_file_count = 0
                    with open(str(datafile_i)) as f:
                        for _ in f:
                            total_per_file_count += 1
                    total_per_file.append(total_per_file_count)
                    total_count += total_per_file_count
                else:
                    sys.exit("ERROR: Criteo Terabyte Dataset path is invalid; please download from https://labs.criteo.com/2013/12/download-terabyte-click-logs")

    # process a file worth of data and reinitialize data
    # note that a file main contain a single or multiple splits
    def process_one_file(
            datfile,
            npzfile,
            split,
            num_data_in_split,
            dataset_multiprocessing,
            convertDictsDay=None,
            resultDay=None
    ):
        if dataset_multiprocessing:
            convertDicts_day = [{} for _ in range(26)]

        with open(str(datfile)) as f:
            y = np.zeros(num_data_in_split, dtype="i4")  # 4 byte int
            X_int = np.zeros((num_data_in_split, 13), dtype="i4")  # 4 byte int
            X_cat = np.zeros((num_data_in_split, 26), dtype="i4")  # 4 byte int
            if sub_sample_rate == 0.0:
                rand_u = 1.0
            else:
                rand_u = np.random.uniform(low=0.0, high=1.0, size=num_data_in_split)

            i = 0
            percent = 0
            for k, line in enumerate(f):
                # process a line (data point)
                line = line.split('\t')
                # set missing values to zero
                for j in range(len(line)):
                    if (line[j] == '') or (line[j] == '\n'):
                        line[j] = '0'
                # sub-sample data by dropping zero targets, if needed
                target = np.int32(line[0])
                if target == 0 and \
                   (rand_u if sub_sample_rate == 0.0 else rand_u[k]) < sub_sample_rate:
                    continue

                y[i] = target
                X_int[i] = np.array(line[1:14], dtype=np.int32)
                if max_ind_range > 0:
                    X_cat[i] = np.array(
                        list(map(lambda x: int(x, 16) % max_ind_range, line[14:])),
                        dtype=np.int32
                    )
                else:
                    X_cat[i] = np.array(
                        list(map(lambda x: int(x, 16), line[14:])),
                        dtype=np.int32
                    )

                # count uniques
                if dataset_multiprocessing:
                    for j in range(26):
                        convertDicts_day[j][X_cat[i][j]] = 1
                    # debug prints
                    if float(i)/num_data_in_split*100 > percent+1:
                        percent = int(float(i)/num_data_in_split*100)
                        print(
                            "Load %d/%d (%d%%) Split: %d  Label True: %d  Stored: %d"
                            % (
                                i,
                                num_data_in_split,
                                percent,
                                split,
                                target,
                                y[i],
                            ),
                            end="\n",
                        )
                else:
                    for j in range(26):
                        convertDicts[j][X_cat[i][j]] = 1
                    # debug prints
                    print(
                        "Load %d/%d  Split: %d  Label True: %d  Stored: %d"
                        % (
                            i,
                            num_data_in_split,
                            split,
                            target,
                            y[i],
                        ),
                        end="\r",
                    )
                i += 1

            # store num_data_in_split samples or extras at the end of file
            # count uniques
            # X_cat_t  = np.transpose(X_cat)
            # for j in range(26):
            #     for x in X_cat_t[j,:]:
            #         convertDicts[j][x] = 1
            # store parsed
            filename_s = npzfile + "_{0}.npz".format(split)
            if path.exists(filename_s):
                print("\nSkip existing " + filename_s)
            else:
                np.savez_compressed(
                    filename_s,
                    X_int=X_int[0:i, :],
                    # X_cat=X_cat[0:i, :],
                    X_cat_t=np.transpose(X_cat[0:i, :]),  # transpose of the data
                    y=y[0:i],
                )
                print("\nSaved " + npzfile + "_{0}.npz!".format(split))

        if dataset_multiprocessing:
            resultDay[split] = i
            convertDictsDay[split] = convertDicts_day
            return
        else:
            return i

    # create all splits (reuse existing files if possible)
    recreate_flag = False
    convertDicts = [{} for _ in range(26)]
    # WARNING: to get reproducable sub-sampling results you must reset the seed below
    # np.random.seed(123)
    # in this case there is a single split in each day
    for i in range(days):
        npzfile_i = npzfile + "_{0}.npz".format(i)
        npzfile_p = npzfile + "_{0}_processed.npz".format(i)
        if path.exists(npzfile_i):
            print("Skip existing " + npzfile_i)
        elif path.exists(npzfile_p):
            print("Skip existing " + npzfile_p)
        else:
            recreate_flag = True

    if recreate_flag:
        if dataset_multiprocessing:
            resultDay = Manager().dict()
            convertDictsDay = Manager().dict()
            processes = [Process(target=process_one_file,
                                 name="process_one_file:%i" % i,
                                 args=(npzfile + "_{0}".format(i),
                                       npzfile,
                                       i,
                                       total_per_file[i],
                                       dataset_multiprocessing,
                                       convertDictsDay,
                                       resultDay,
                                       )
                                 ) for i in range(0, days)]
            for process in processes:
                process.start()
            for process in processes:
                process.join()
            for day in range(days):
                total_per_file[day] = resultDay[day]
                print("Constructing convertDicts Split: {}".format(day))
                convertDicts_tmp = convertDictsDay[day]
                for i in range(26):
                    for j in convertDicts_tmp[i]:
                        convertDicts[i][j] = 1
        else:
            for i in range(days):
                total_per_file[i] = process_one_file(
                    npzfile + "_{0}".format(i),
                    npzfile,
                    i,
                    total_per_file[i],
                    dataset_multiprocessing,
                )

    # report and save total into a file
    total_count = np.sum(total_per_file)
    if not path.exists(total_file):
        np.savez_compressed(total_file, total_per_file=total_per_file)
    print("Total number of samples:", total_count)
    print("Divided into days/splits:\n", total_per_file)

    # dictionary files
    counts = np.zeros(26, dtype=np.int32)
    if recreate_flag:
        # create dictionaries
        for j in range(26):
            for i, x in enumerate(convertDicts[j]):
                convertDicts[j][x] = i
            dict_file_j = d_path + d_file + "_fea_dict_{0}.npz".format(j)
            if not path.exists(dict_file_j):
                np.savez_compressed(
                    dict_file_j,
                    unique=np.array(list(convertDicts[j]), dtype=np.int32)
                )
            counts[j] = len(convertDicts[j])
        # store (uniques and) counts
        count_file = d_path + d_file + "_fea_count.npz"
        if not path.exists(count_file):
            np.savez_compressed(count_file, counts=counts)
    else:
        # create dictionaries (from existing files)
        for j in range(26):
            with np.load(d_path + d_file + "_fea_dict_{0}.npz".format(j)) as data:
                unique = data["unique"]
            for i, x in enumerate(unique):
                convertDicts[j][x] = i
        # load (uniques and) counts
        with np.load(d_path + d_file + "_fea_count.npz") as data:
            counts = data["counts"]

    # process all splits
    if dataset_multiprocessing:
        processes = [Process(target=processCriteoAdData,
                           name="processCriteoAdData:%i" % i,
                           args=(d_path,
                                 d_file,
                                 npzfile,
                                 i,
                                 convertDicts,
                                 counts,
                                 )
                           ) for i in range(0, days)]
        for process in processes:
            process.start()
        for process in processes:
            process.join()

    else:
        for i in range(days):
            processCriteoAdData(d_path, d_file, npzfile, i, convertDicts, counts)

    o_file = concatCriteoAdData(
        d_path,
        d_file,
        npzfile,
        trafile,
        days,
        data_split,
        randomize,
        total_per_file,
        total_count,
        memory_map,
        o_filename
    )

    return o_file


def loadDataset(
        dataset,
        max_ind_range,
        sub_sample_rate,
        randomize,
        data_split,
        raw_path="",
        pro_data="",
        memory_map=False
):
    # dataset
    if dataset == "kaggle":
        days = 7
        o_filename = "kaggleAdDisplayChallenge_processed"
    elif dataset == "terabyte":
        days = 24
        o_filename = "terabyte_processed"
    else:
        raise(ValueError("Data set option is not supported"))

    # split the datafile into path and filename
    lstr = raw_path.split("/")
    d_path = "/".join(lstr[0:-1]) + "/"
    d_file = lstr[-1].split(".")[0] if dataset == "kaggle" else lstr[-1]
    npzfile = (d_file + "_day") if dataset == "kaggle" else d_file
    # trafile = d_path + ((d_file + "_fea") if dataset == "kaggle" else "fea")

    # check if pre-processed data is available
    data_ready = True
    if memory_map:
        for i in range(days):
            reo_data = d_path + npzfile + "_{0}_reordered.npz".format(i)
            if not path.exists(str(reo_data)):
                data_ready = False
    else:
        if not path.exists(str(pro_data)):
            data_ready = False

    # pre-process data if needed
    # WARNNING: when memory mapping is used we get a collection of files
    if data_ready:
        print("Reading pre-processed data=%s" % (str(pro_data)))
        file = str(pro_data)
    else:
        print("Reading raw data=%s" % (str(raw_path)))
        file = getCriteoAdData(
            raw_path,
            o_filename,
            max_ind_range,
            sub_sample_rate,
            days,
            data_split,
            randomize,
            dataset == "kaggle",
            memory_map
        )

    return file, days


if __name__ == "__main__":
    ### import packages ###
    import argparse

    ### parse arguments ###
    parser = argparse.ArgumentParser(
        description="Preprocess Criteo dataset"
    )
    # model related parameters
    parser.add_argument("--max-ind-range", type=int, default=-1)
    parser.add_argument("--data-sub-sample-rate", type=float, default=0.0)  # in [0, 1]
    parser.add_argument("--data-randomize", type=str, default="total")  # or day or none
    parser.add_argument("--memory-map", action="store_true", default=False)
    parser.add_argument("--data-set", type=str, default="kaggle")  # or terabyte
    parser.add_argument("--raw-data-file", type=str, default="")
    parser.add_argument("--processed-data-file", type=str, default="")
    args = parser.parse_args()

    loadDataset(
        args.data_set,
        args.max_ind_range,
        args.data_sub_sample_rate,
        args.data_randomize,
        "train",
        args.raw_data_file,
        args.processed_data_file,
        args.memory_map
    )