import tensorflow as tf
import numpy as np
from scipy.stats import rankdata
np.random.seed(1234)
import os
import time
import datetime
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from builddata import *
from model import MCapsEED
import config
from  utils import *
# Parameters
# ==================================================
parser = ArgumentParser("MCapsEED", formatter_class=ArgumentDefaultsHelpFormatter, conflict_handler='resolve')


parser.add_argument("--data", default="./data/", help="Data sources.")
parser.add_argument("--run_folder", default="./", help="Data sources.")
parser.add_argument("--name", default="FB15k-237", help="Name of the dataset.")

parser.add_argument("--embedding_dim", default=100, type=int, help="Dimensionality of character embedding (default: 128)")
parser.add_argument("--filter_size", default=1, type=int, help="Comma-separated filter sizes (default: '3,4,5')")
parser.add_argument("--num_filters", default=50, type=int, help="Number of filters per filter size (default: 128)")
parser.add_argument("--learning_rate", default=0.0001, type=float, help="Learning rate")
parser.add_argument("--batch_size", default=20, type=int, help="Batch Size")
parser.add_argument("--neg_ratio", default=1.0, help="Number of negative triples generated by positive (default: 1.0)")
parser.add_argument("--num_epochs", default=100, type=int, help="Number of training epochs")
parser.add_argument("--savedEpochs", default=101, type=int, help="")
parser.add_argument("--allow_soft_placement", default=True, type=bool, help="Allow device soft device placement")
parser.add_argument("--log_device_placement", default=False, type=bool, help="Log placement of ops on devices")
parser.add_argument("--model_name", default='fb15k237', help="")
parser.add_argument("--useConstantInit", default=True)

parser.add_argument('--iter_routing', default=1, type=int, help='number of iterations in routing algorithm')
parser.add_argument('--num_outputs_secondCaps', default=1, type=int, help='')
parser.add_argument('--vec_len_secondCaps', default=10, type=int, help='')
config=config.Config()
args = parser.parse_args()
print(args)
import logging
logging.basicConfig(level=logging.DEBUG,#控制台打印的日志级别
                    filename='new.log',
                    filemode='a',##模式，有w和a，w就是写模式，每次都会重新写日志，覆盖之前的日志
                    #a是追加模式，默认如果不写的话，就是追加模式
                    format=
                    '%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s'
                    #日志格式
                    )
print("Loading data...")
# Load data
print("Loading data...")

train, valid, test, words_indexes, indexes_words, \
headTailSelector, entity2id, id2entity, relation2id, id2relation = build_data(path=args.data, name=args.name)
data_size = len(train)
train_batch = Batch_Loader(train, words_indexes, indexes_words, headTailSelector, \
                           entity2id, id2entity, relation2id, id2relation, batch_size=args.batch_size,
                           neg_ratio=args.neg_ratio)

entity_array = np.array(list(train_batch.indexes_ents.keys()))

lstEmbed = []

#Using the pre-trained embeddings.
print("Using pre-trained model.")
lstEmbed = np.empty([len(words_indexes), args.embedding_dim]).astype(np.float32)
initEnt, initRel = init_norm_Vector(args.data + args.name + '/relation2vec' + str(args.embedding_dim) + '.init',
                                    args.data + args.name + '/entity2vec' + str(args.embedding_dim) + '.init',
                                    args.embedding_dim)
for _word in words_indexes:
    if _word in relation2id:
        index = relation2id[_word]
        _ind = words_indexes[_word]
        lstEmbed[_ind] = initRel[index]
    elif _word in entity2id:
        index = entity2id[_word]
        _ind = words_indexes[_word]
        lstEmbed[_ind] = initEnt[index]
    else:
        print('*****************Error********************!')
        break
lstEmbed = np.array(lstEmbed, dtype=np.float32)

assert len(words_indexes) % (len(entity2id) + len(relation2id)) == 0

wiki_words, entity2wiki, max_len1 = load_description("data/FB15k-237/FB15k_237_description.txt",100)
#{所有单词--->w1,w2}
#{实体:[单词]}
embedding_matrix, word2id = load_pretrain_emb('data/FB15k-237/glove.6B.100d.txt', wiki_words)
entity2wikiID = {}
for entity, words in entity2wiki.items():
    temp = entity2wikiID.get(entity, [])
    for word in words:
        temp.append(word2id[word])
    entity2wikiID[entity] = temp
#映射到Wordindex单词序列
entityindex2wikiID={}
for w,index in words_indexes.items():
    if w in entity2wikiID:
        entityindex2wikiID[index]=np.asarray(entity2wikiID[w],np.int32)

entity_index_word_metrix=np.zeros([len(words_indexes), args.embedding_dim],dtype=np.int32)
for index,embed in entityindex2wikiID.items():
    entity_index_word_metrix[index,:]=embed
entity_index_word_metrix=np.asarray(entity_index_word_metrix,dtype=np.int32)

x_valid = np.array(list(valid.keys())).astype(np.int32)
y_valid = np.array(list(valid.values())).astype(np.float32)

x_test = np.array(list(test.keys())).astype(np.int32)

print("Loading data... finished!")

with tf.Graph().as_default():
    session_conf = tf.ConfigProto(allow_soft_placement=args.allow_soft_placement,
                                  log_device_placement=args.log_device_placement)
    session_conf.gpu_options.allow_growth = True
    sess = tf.Session(config=session_conf)
    with sess.as_default():
        global_step = tf.Variable(0, name="global_step", trainable=False)
        mcapseed=MCapsEED(
            sequence_length=x_valid.shape[1],
            pre_trained=lstEmbed,
            embedding_size=args.embedding_dim,
            filter_size=args.filter_size,
            num_filters=args.num_filters,
            vocab_size=len(words_indexes),
            iter_routing=args.iter_routing,
            batch_size=args.batch_size*2,
            num_outputs_secondCaps=args.num_outputs_secondCaps,
            vec_len_secondCaps=args.vec_len_secondCaps,
            useConstantInit=args.useConstantInit,
            entity_index_word_metrix=entity_index_word_metrix,
            embedding_matrix=embedding_matrix,
            config=config
        )
        # optimizer = tf.contrib.opt.NadamOptimizer(1e-3)
        optimizer = tf.train.AdamOptimizer(learning_rate=args.learning_rate)
        # optimizer = tf.train.RMSPropOptimizer(learning_rate=args.learning_rate)
        # optimizer = tf.train.GradientDescentOptimizer(learning_rate=args.learning_rate)
        grads_and_vars = optimizer.compute_gradients(mcapseed.total_loss)
        train_op = optimizer.apply_gradients(grads_and_vars, global_step=global_step)

        out_dir = os.path.abspath(os.path.join(args.run_folder, "runs_CapsE", args.model_name))
        print("Writing to {}\n".format(out_dir))

        checkpoint_dir = os.path.abspath(os.path.join(out_dir, "checkpoints"))
        print('checkpoint_dir', checkpoint_dir)
        checkpoint_prefix = os.path.join(checkpoint_dir, "model")
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        print('success!')

        # Initialize all variables
        sess.run(tf.global_variables_initializer())


        def predict(x_batch, writer=None):
            feed_dict = {
                mcapseed.input_x: x_batch
            }
            scores = sess.run([mcapseed.predictions], feed_dict)
            return scores


        def test_prediction(x_batch, head_or_tail='head'):

            hits10 = 0.0
            mrr = 0.0
            mr = 0.0
            hits1 = 0.0
            print(len(x_batch))
            for i in range(len(x_batch)):
                new_x_batch = np.tile(x_batch[i], (len(entity2id), 1))
                if head_or_tail == 'head':
                    new_x_batch[:, 0] = entity_array
                else:  # 'tail'
                    new_x_batch[:, 2] = entity_array

                lstIdx = []
                for tmpIdxTriple in range(len(new_x_batch)):
                    tmpTriple = (new_x_batch[tmpIdxTriple][0], new_x_batch[tmpIdxTriple][1],
                                 new_x_batch[tmpIdxTriple][2])
                    if (tmpTriple in train) or (tmpTriple in valid) or (
                            tmpTriple in test):  # also remove the valid test triple
                        lstIdx.append(tmpIdxTriple)
                new_x_batch = np.delete(new_x_batch, lstIdx, axis=0)
                # thus, insert the valid test triple again, to the beginning of the array
                new_x_batch = np.insert(new_x_batch, 0, x_batch[i],
                                        axis=0)  # thus, the index of the valid test triple is equal to 0

                # for running with a batch size
                while len(new_x_batch) % ((int(args.neg_ratio) + 1) * args.batch_size) != 0:
                    new_x_batch = np.append(new_x_batch, [x_batch[i]], axis=0)

                results = []
                listIndexes = range(0, len(new_x_batch), (int(args.neg_ratio) + 1) * args.batch_size)
                for tmpIndex in range(len(listIndexes) - 1):
                    results = np.append(results, predict(
                        new_x_batch[listIndexes[tmpIndex]:listIndexes[tmpIndex + 1]]))
                results = np.append(results,
                                    predict(new_x_batch[listIndexes[-1]:]
                                            ))

                results = np.reshape(results, -1)
                results_with_id = rankdata(results, method='ordinal')
                _filter = results_with_id[0]

                mr += _filter
                mrr += 1.0 / _filter
                if _filter <= 10:
                    hits10 += 1

                if _filter == 1:
                    hits1 += 1

            return np.array([mr, mrr, hits1, hits10])


        def train_step(x_batch, y_batch):
            """
            A single training step
            """
            feed_dict = {
                mcapseed.input_x: x_batch,
                mcapseed.input_y: y_batch
            }
            _, step, loss = sess.run([train_op, global_step, mcapseed.total_loss], feed_dict)
            return loss


        num_batches_per_epoch = int((data_size - 1) / args.batch_size) + 1
        for epoch in range(args.num_epochs):
            loss = 0
            for batch_num in range(num_batches_per_epoch):
                x_batch, y_batch = train_batch()
                loss = train_step(x_batch, y_batch)
                current_step = tf.train.global_step(sess, global_step)
                loss += loss
            print(f'{epoch}_mean_loss:{loss / num_batches_per_epoch}')
            logging.info(f'{epoch}_mean_loss:{loss / num_batches_per_epoch}')
            path = mcapseed.saver.save(sess, checkpoint_prefix, global_step=epoch)
            print("Saved model checkpoint to {}\n".format(path))
            logging.info("Saved model checkpoint to {}\n".format(path))
            if epoch > 0:
                if epoch % 50 == 0:
                    head_results = test_prediction(x_test,
                                                   head_or_tail='head')
                    tail_results = test_prediction(x_test,
                                                   head_or_tail='tail')
                    head_results, tail_results = np.asarray(head_results), np.asarray(tail_results)
                    result = list(np.sum([head_results, tail_results], axis=0) / (2 * len(x_test)))
                    print('head_tail_result: mr, mrr, hits@1, hits@10 --> ', result[0], result[1], result[2], result[3])
                    logging.info(
                        f'head_tail_result: mr, mrr, hits@1, hits@10 -->{result[0]},{result[1]},{result[2]},{result[3]}')
                    path = mcapseed.saver.save(sess, checkpoint_prefix, global_step=epoch)
                    print("Saved model checkpoint to {}\n".format(path))
                    logging.info("Saved model checkpoint to {}\n".format(path))