# Copyright (c) 2018, deepakn94, codyaustun, robieta. All rights reserved.
#
# 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.
#
# -----------------------------------------------------------------------
#
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
#
# 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 time
import torch
import tqdm

class _TestNegSampler:
    def __init__(self, train_ratings, nb_neg):
        self.nb_neg = nb_neg
        self.nb_users = int(train_ratings[:, 0].max()) + 1
        self.nb_items = int(train_ratings[:, 1].max()) + 1

        # compute unique ids for quickly created hash set and fast lookup
        ids = (train_ratings[:, 0] * self.nb_items) + train_ratings[:, 1]
        self.set = set(ids)

    def generate(self, batch_size=128*1024):
        users = torch.arange(0, self.nb_users).reshape([1, -1]).repeat([self.nb_neg, 1]).transpose(0, 1).reshape(-1)

        items = [-1] * len(users)

        random_items = torch.LongTensor(batch_size).random_(0, self.nb_items).tolist()
        print('Generating validation negatives...')
        for idx, u in enumerate(tqdm.tqdm(users.tolist())):
            if not random_items:
                random_items = torch.LongTensor(batch_size).random_(0, self.nb_items).tolist()
            j = random_items.pop()
            while u * self.nb_items + j in self.set:
                if not random_items:
                    random_items = torch.LongTensor(batch_size).random_(0, self.nb_items).tolist()
                j = random_items.pop()

            items[idx] = j
        items = torch.LongTensor(items)
        return items


def create_test_data(train_ratings, test_ratings, args):
    test_users = test_ratings[:,0]
    test_pos = test_ratings[:,1].reshape(-1,1)

    begin = time.time()
    sampler = _TestNegSampler(train_ratings.cpu().numpy(), args.valid_negative)
    test_negs = sampler.generate().cuda()
    end = time.time()
    print('Generating validation negatives took: ', end - begin)
    del train_ratings

    # create items with real sample at last position
    test_users = test_users.reshape(-1,1).repeat(1, 1 + args.valid_negative)
    test_items = torch.cat((test_negs.reshape(-1, args.valid_negative), test_pos), dim=1)
    del test_ratings, test_negs

    # generate dup mask and real indices for exact same behavior on duplication compare to reference
    # here we need a sort that is stable(keep order of duplicates)
    sorted_items, indices = torch.sort(test_items) # [1,1,1,2], [3,1,0,2]
    sum_item_indices = sorted_items.float()+indices.float()/len(indices[0]) #[1.75,1.25,1.0,2.5]
    indices_order = torch.sort(sum_item_indices)[1] #[2,1,0,3]
    stable_indices = torch.gather(indices, 1, indices_order) #[0,1,3,2]
    # produce -1 mask
    dup_mask = (sorted_items[:,0:-1] == sorted_items[:,1:])
    dup_mask = torch.cat((torch.zeros_like(test_pos, dtype=torch.uint8), dup_mask),dim=1)
    dup_mask = torch.gather(dup_mask,1,stable_indices.sort()[1])
    # produce real sample indices to later check in topk
    sorted_items, indices = (test_items != test_pos).sort()
    sum_item_indices = sorted_items.float()+indices.float()/len(indices[0])
    indices_order = torch.sort(sum_item_indices)[1]
    stable_indices = torch.gather(indices, 1, indices_order)
    real_indices = stable_indices[:,0]

    if args.distributed:
        test_users = torch.chunk(test_users, args.world_size)[args.local_rank]
        test_items = torch.chunk(test_items, args.world_size)[args.local_rank]
        dup_mask = torch.chunk(dup_mask, args.world_size)[args.local_rank]
        real_indices = torch.chunk(real_indices, args.world_size)[args.local_rank]

    test_users = test_users.view(-1).split(args.valid_batch_size)
    test_items = test_items.view(-1).split(args.valid_batch_size)

    return test_users, test_items, dup_mask, real_indices


def prepare_epoch_train_data(train_ratings, nb_items, args):
    # create label
    train_label = torch.ones_like(train_ratings[:,0], dtype=torch.float32)
    neg_label = torch.zeros_like(train_label, dtype=torch.float32)
    neg_label = neg_label.repeat(args.negative_samples)
    train_label = torch.cat((train_label,neg_label))
    del neg_label

    train_users = train_ratings[:,0]
    train_items = train_ratings[:,1]

    train_users_per_worker = len(train_label) / args.world_size
    train_users_begin = int(train_users_per_worker * args.local_rank)
    train_users_end = int(train_users_per_worker * (args.local_rank + 1))

    # prepare data for epoch
    neg_users = train_users.repeat(args.negative_samples)
    neg_items = torch.empty_like(neg_users, dtype=torch.int64).random_(0, nb_items)

    epoch_users = torch.cat((train_users, neg_users))
    epoch_items = torch.cat((train_items, neg_items))

    del neg_users, neg_items

    # shuffle prepared data and split into batches
    epoch_indices = torch.randperm(train_users_end - train_users_begin, device='cuda:{}'.format(args.local_rank))
    epoch_indices += train_users_begin

    epoch_users = epoch_users[epoch_indices]
    epoch_items = epoch_items[epoch_indices]
    epoch_label = train_label[epoch_indices]

    if args.distributed:
        local_batch = args.batch_size // args.world_size
    else:
        local_batch = args.batch_size

    epoch_users = epoch_users.split(local_batch)
    epoch_items = epoch_items.split(local_batch)
    epoch_label = epoch_label.split(local_batch)

    # the last batch will almost certainly be smaller, drop it
    epoch_users = epoch_users[:-1]
    epoch_items = epoch_items[:-1]
    epoch_label = epoch_label[:-1]

    return epoch_users, epoch_items, epoch_label

