#encoding=utf-8

"""
The main entry.
You can execute this file to do model training and test.
"""

import numpy
import theano
from theano import tensor

import os

import ConfigParser
import string, os, sys
import interactiveGRU
import interactiveGRUProcessAndAssessBatch
import time
import subprocess

if __name__=='__main__':
    
    cf = ConfigParser.SafeConfigParser()
    cf.read("pythonParamsConfig") # the config file
    
    main_dir=cf.get("param", "root_dir")  # the main dir
    dataset_name=cf.get("param", "dataset_name")  # dataset name
    suffix=cf.get("param", "suffix")  # the suffix of each dataset, such as 10,100 and 1000
    class_name=cf.get("param", "class_name")  # the relation name, such as family, school
    index=cf.get("param", "index") # the index of the dataset file
    
    trainingDataFile=os.path.join(main_dir+'/',dataset_name+'.splits','train.'+suffix,'train_'+class_name+'_'+index) # the full path of training data file. This path will be generated by main_dir, dataset_name, suffix, class_name and index.
    wordsEmbeddings=None  # words的embeddings
    wordsEmbeddings_path=cf.get("param", "wordsEmbeddings_path")  # the file path of embeddings
    typesEmbeddings=None # type info embeddings
    typesEmbeddings_path=cf.get("param", "typesEmbeddings_path")  # the file path of type embeddings
    
    word_dimension=cf.getint("param", "word_dimension") # dimension of node embedding
    type_dimension=cf.getint("param", "type_dimension") # dimension of type embedding
    dimension=cf.getint("param", "dimension") # dimension d
    attention_dimension=cf.getint("param", "attention_dimension") # dimension of attention_dimension, but we do not use this parameter
    
    wordsSize=cf.getint("param", "wordsSize") # the size of vocabulary
    subpaths_map=None # sub-paths map
    subpaths_fileTrain=cf.get("param", "subpaths_file")+index # sub-paths file for training
    subpaths_fileTest=cf.get("param", "subpaths_file")+'Test' # sub-paths file for testing
    sequences_map=None # the map to store sequences
    sequences_fileTrain=main_dir+'/'+dataset_name+'/sequencesSaveFileTrain' # the sequences file for training 
    sequences_fileTest=main_dir+'/'+dataset_name+'/sequencesSaveFileTest'+index# the sequences file for test
    
    maxlen_subpaths=cf.getint("param", "maxlen_subpaths") # the max length for subpath, we will remove the subpaths which longer than this max length
    maxlen=cf.getint("param", "maxlen")  # Sequence longer then this get ignored 
    batch_size=cf.getint("param", "batch_size") # batch size 
    is_shuffle_for_batch=cf.getboolean("param", "is_shuffle_for_batch") # shuffle for batch or not
    alpha=cf.getfloat("param", "alpha") # parameter alpha
    # we split the parameter beta in the paper into two parameters, beta and gamma as following
    beta=cf.getfloat("param", "beta") # parameter beta, which is used when aggregating multiple paths
    gamma=cf.getfloat("param", "gamma") # parameter gamma, which is used when aggregating multiple predecessors
    objective_function_method=cf.get("param", "objective_function_method") # loss function
    objective_function_param=cf.getfloat("param", "objective_function_param") # loss function parameter
    lrate=cf.getfloat("param", "lrate") # learning rate
    max_epochs=cf.getint("param", "max_epochs") # the max epochs in training 
    
    dispFreq=cf.getint("param", "dispFreq") # display frequency
    saveFreq=cf.getint("param", "saveFreq") # save frequency
    saveto=os.path.join(main_dir+'/',dataset_name+'.trainModels','train.'+suffix,'train_'+class_name+'_'+index+'.npz') # parameters save file
    
    decay=cf.getfloat("param", "decay") # the parameter of decay
    
    test_data_file=os.path.join(main_dir+'/',dataset_name+'.splits','test','test_'+class_name+'_'+index) # test data file
    top_num=cf.getint("param", "top_num") # the top num to predict, we set top_num=10
    ideal_data_file=os.path.join(main_dir+'/',dataset_name+'.splits','ideal','ideal_'+class_name+'_'+index) # ground truth file

    
    print 'Start to train the model ...... time ==',time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    interactiveGRU.interactiveGRUTraining(
        trainingDataFile, 
        wordsEmbeddings, 
        wordsEmbeddings_path, 
        typesEmbeddings, 
        typesEmbeddings_path, 
        word_dimension, 
        type_dimension, 
        dimension, 
        attention_dimension, 
        wordsSize, 
        subpaths_map, 
        subpaths_fileTrain, 
        sequences_map, 
        sequences_fileTrain, 
        maxlen_subpaths, 
        maxlen, 
        batch_size, 
        is_shuffle_for_batch, 
        alpha, 
        beta, 
        gamma, 
        objective_function_method, 
        objective_function_param, 
        lrate, 
        max_epochs, 
        dispFreq, 
        saveFreq, 
        saveto, 
        decay)
    
    time.sleep(5) 
    
    print 'Finished training ...... time ==',time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    print 'Start to generate the test model ...... time ==',time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    
        
    
     
    computeFunc=interactiveGRUProcessAndAssessBatch.get_interactiveGRUModel(
        saveto, 
        word_dimension, 
        type_dimension, 
        dimension, 
        attention_dimension, 
        alpha, 
        beta, 
        gamma)
      
    print 'Start to calculate the queries ...... time ==',time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    MAP,MnDCG=interactiveGRUProcessAndAssessBatch.compute_path2vec(
        wordsEmbeddings, 
        wordsEmbeddings_path, 
        typesEmbeddings, 
        typesEmbeddings_path, 
        word_dimension, 
        type_dimension, 
        dimension, 
        attention_dimension, 
        wordsSize, 
        subpaths_map, 
        subpaths_fileTest, 
        sequences_map, 
        sequences_fileTest, 
        maxlen_subpaths, 
        maxlen, 
        alpha, 
        beta, 
        gamma, 
        test_data_file, 
        top_num, 
        ideal_data_file, 
        computeFunc,
        )
      
    print 'Finished the calculation ...... time ==',time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
    print 'MAP==',MAP
    print 'MnDCG==',MnDCG