from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import json
import logging
import os
import random

import numpy as np
import torch

from torch.utils.data import DataLoader

from model import KGEModel

from dataloader import TrainDataset
from dataloader import TestDataset
from dataloader import BidirectionalOneShotIterator

a_data_path = "../data/FB15k"

def read_triple(file_path, entity2id, relation2id):
    '''
    Read triples and map them into ids.
    '''
    triples = []
    with open(file_path) as fin:
        for line in fin:
            h, r, t = line.strip().split('\t')
            triples.append((entity2id[h], relation2id[r], entity2id[t]))
    return triples


# construct 2 dicts: entity->id, relation->id
with open(os.path.join(a_data_path, 'entities.dict')) as fin:
    entity2id = dict()
    for line in fin:
        eid, entity = line.strip().split('\t')
        entity2id[entity] = int(eid)

with open(os.path.join(a_data_path, 'relations.dict')) as fin:
    relation2id = dict()
    for line in fin:
        rid, relation = line.strip().split('\t')
        relation2id[relation] = int(rid)


nentity = len(entity2id)
nrelation = len(relation2id)

a_nentity = nentity
a_nrelation = nrelation

train_triples = read_triple(os.path.join(a_data_path, 'train.txt'), entity2id, relation2id)
valid_triples = read_triple(os.path.join(a_data_path, 'valid.txt'), entity2id, relation2id)
test_triples = read_triple(os.path.join(a_data_path, 'test.txt'), entity2id, relation2id)

all_true_triples = train_triples + valid_triples + test_triples

a_model = "RotatE"
a_hidden_dim = 1000
a_gamma = 24.0

a_batchsize = 1024
a_neg_sample_size = 256
a_alpha = 1.0
a_lr = 0.0001
a_max_steps = 150000
a_test_batchsize = 16

kge_model = KGEModel(
    model_name=a_model,
    nentity=nentity,
    nrelation=nrelation,
    hidden_dim=a_hidden_dim,
    gamma=a_gamma,
    double_entity_embedding=True,
    double_relation_embedding=False
)

train_dataloader_head = DataLoader(
    TrainDataset(train_triples, nentity, nrelation, a_neg_sample_size, 'head-batch'),
    batch_size=a_batchsize,
    shuffle=True,
    # num_workers=4,
    collate_fn=TrainDataset.collate_fn
)

train_dataloader_tail = DataLoader(
    TrainDataset(train_triples, nentity, nrelation, a_neg_sample_size, 'tail-batch'),
    batch_size=a_batchsize,
    shuffle=True,
    # num_workers=4,
    collate_fn=TrainDataset.collate_fn
)


test_dataloader_head = DataLoader(
                TestDataset(
                    test_triples,
                    all_true_triples,
                    a_nentity,
                    a_nrelation,
                    'head-batch'
                ),
                batch_size=a_test_batchsize,
                # num_workers=max(1, args.cpu_num // 2),
                collate_fn=TestDataset.collate_fn
)

test_dataloader_tail = DataLoader(
    TestDataset(
        test_triples,
        all_true_triples,
        a_nentity,
        a_nrelation,
        'tail-batch'
    ),
    batch_size=a_test_batchsize,
    # num_workers=max(1, args.cpu_num // 2),
    collate_fn=TestDataset.collate_fn
)

train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail)

#############################################

positive_sample, negative_sample, subsampling_weight, mode = next(train_iterator)
print(positive_sample.size())
print(positive_sample[0])
print(positive_sample[1023])
print(negative_sample.size())
print(negative_sample[0])
print(subsampling_weight.size())
print(mode, '\n\n')

for ps, ns, fb, m in test_dataloader_head:
    print(ps.size())
    print(ps[0], '\n')
    print(ns.size())
    print(ns[0], '\n')
    print(fb.data.size)
    print(fb[15][14900:], '\n')
    print(m)
    break




# current_learning_rate = a_lr
# optimizer = torch.optim.Adam(
#     filter(lambda p: p.requires_grad, kge_model.parameters()),
#     lr=current_learning_rate
# )
#
# warm_up_steps = a_max_steps // 2
#
# init_step = 0
# step = 0
#
#
# # # Training Loop
# # for step in range(init_step, a_max_steps):
# #
# #     log = kge_model.train_step(kge_model, optimizer, train_iterator, args=None)
# #
# #     if step >= warm_up_steps:
# #         current_learning_rate = current_learning_rate / 10
# #         logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step))
# #         optimizer = torch.optim.Adam(
# #             filter(lambda p: p.requires_grad, kge_model.parameters()),
# #             lr=current_learning_rate
# #         )
# #         warm_up_steps = warm_up_steps * 3
#
#
# def train_step(model, optimizer, train_iterator, args):
#     '''
#     A single train step. Apply back-propation and return the loss
#     '''
#
#     model.train()
#
#     optimizer.zero_grad()
#
#     positive_sample, negative_sample, subsampling_weight, mode = next(train_iterator)
#
#     if args.cuda:
#         positive_sample = positive_sample.cuda()
#         negative_sample = negative_sample.cuda()
#         subsampling_weight = subsampling_weight.cuda()
#
#     negative_score = model((positive_sample, negative_sample), mode=mode)
#
#     if args.negative_adversarial_sampling:
#         # In self-adversarial sampling, we do not apply back-propagation on the sampling weight
#         negative_score = (F.softmax(negative_score * args.adversarial_temperature, dim=1).detach()
#                           * F.logsigmoid(-negative_score)).sum(dim=1)
#     else:
#         negative_score = F.logsigmoid(-negative_score).mean(dim=1)
#
#     positive_score = model(positive_sample)
#
#     positive_score = F.logsigmoid(positive_score).squeeze(dim=1)
#
#     if args.uni_weight:
#         positive_sample_loss = - positive_score.mean()
#         negative_sample_loss = - negative_score.mean()
#     else:
#         positive_sample_loss = - (subsampling_weight * positive_score).sum() / subsampling_weight.sum()
#         negative_sample_loss = - (subsampling_weight * negative_score).sum() / subsampling_weight.sum()
#
#     loss = (positive_sample_loss + negative_sample_loss) / 2
#
#     if args.regularization != 0.0:
#         # Use L3 regularization for ComplEx and DistMult
#         regularization = args.regularization * (
#                 model.entity_embedding.norm(p=3) ** 3 +
#                 model.relation_embedding.norm(p=3).norm(p=3) ** 3
#         )
#         loss = loss + regularization
#         regularization_log = {'regularization': regularization.item()}
#     else:
#         regularization_log = {}
#
#     loss.backward()
#
#     optimizer.step()
#
#     log = {
#         **regularization_log,
#         'positive_sample_loss': positive_sample_loss.item(),
#         'negative_sample_loss': negative_sample_loss.item(),
#         'loss': loss.item()
#     }
#
#     return log
#
#
# @staticmethod
# def test_step(model, test_triples, all_true_triples, args):
#     '''
#     Evaluate the model on test or valid datasets
#     '''
#
#     model.eval()
#
#     if args.countries:
#         # Countries S* datasets are evaluated on AUC-PR
#         # Process test data for AUC-PR evaluation
#         sample = list()
#         y_true = list()
#         for head, relation, tail in test_triples:
#             for candidate_region in args.regions:
#                 y_true.append(1 if candidate_region == tail else 0)
#                 sample.append((head, relation, candidate_region))
#
#         sample = torch.LongTensor(sample)
#         if args.cuda:
#             sample = sample.cuda()
#
#         with torch.no_grad():
#             y_score = model(sample).squeeze(1).cpu().numpy()
#
#         y_true = np.array(y_true)
#
#         # average_precision_score is the same as auc_pr
#         auc_pr = average_precision_score(y_true, y_score)
#
#         metrics = {'auc_pr': auc_pr}
#
#     else:
#         # Otherwise use standard (filtered) MRR, MR, HITS@1, HITS@3, and HITS@10 metrics
#         # Prepare dataloader for evaluation
#         test_dataloader_head = DataLoader(
#             TestDataset(
#                 test_triples,
#                 all_true_triples,
#                 args.nentity,
#                 args.nrelation,
#                 'head-batch'
#             ),
#             batch_size=args.test_batch_size,
#             num_workers=max(1, args.cpu_num // 2),
#             collate_fn=TestDataset.collate_fn
#         )
#
#         test_dataloader_tail = DataLoader(
#             TestDataset(
#                 test_triples,
#                 all_true_triples,
#                 args.nentity,
#                 args.nrelation,
#                 'tail-batch'
#             ),
#             batch_size=args.test_batch_size,
#             num_workers=max(1, args.cpu_num // 2),
#             collate_fn=TestDataset.collate_fn
#         )
#
#         test_dataset_list = [test_dataloader_head, test_dataloader_tail]
#
#         logs = []
#
#         step = 0
#         total_steps = sum([len(dataset) for dataset in test_dataset_list])
#
#         with torch.no_grad():
#             for test_dataset in test_dataset_list:
#                 for positive_sample, negative_sample, filter_bias, mode in test_dataset:
#                     if args.cuda:
#                         positive_sample = positive_sample.cuda()
#                         negative_sample = negative_sample.cuda()
#                         filter_bias = filter_bias.cuda()
#
#                     batch_size = positive_sample.size(0)
#
#                     score = model((positive_sample, negative_sample), mode)
#                     score += filter_bias
#
#                     # Explicitly sort all the entities to ensure that there is no test exposure bias
#                     argsort = torch.argsort(score, dim=1, descending=True)
#
#                     if mode == 'head-batch':
#                         positive_arg = positive_sample[:, 0]
#                     elif mode == 'tail-batch':
#                         positive_arg = positive_sample[:, 2]
#                     else:
#                         raise ValueError('mode %s not supported' % mode)
#
#                     for i in range(batch_size):
#                         # Notice that argsort is not ranking
#                         ranking = (argsort[i, :] == positive_arg[i]).nonzero()
#                         assert ranking.size(0) == 1
#
#                         # ranking + 1 is the true ranking used in evaluation metrics
#                         ranking = 1 + ranking.item()
#                         logs.append({
#                             'MRR': 1.0 / ranking,
#                             'MR': float(ranking),
#                             'HITS@1': 1.0 if ranking <= 1 else 0.0,
#                             'HITS@3': 1.0 if ranking <= 3 else 0.0,
#                             'HITS@10': 1.0 if ranking <= 10 else 0.0,
#                         })
#
#                     if step % args.test_log_steps == 0:
#                         logging.info('Evaluating the model... (%d/%d)' % (step, total_steps))
#
#                     step += 1
#
#         metrics = {}
#         for metric in logs[0].keys():
#             metrics[metric] = sum([log[metric] for log in logs]) / len(logs)
#
#     return metrics
