import json
import os
import argparse
import numpy as np
import tensorflow as tf
from tqdm import tqdm
import random

from retrieval.score.score import fever_score
from retrieval.sentences.deep_models.BiLSTM_RANKING import BiLSTM_RANKING
from retrieval.sentences.deep_models.Decomposable_Atten import Decomposable_Atten
from retrieval.sentences.deep_models.ESIM import ESIM
from retrieval.sentences.deep_models.USE_RANKING import USERANKING
from common.dataset.reader import JSONLineReader
from drqascripts.retriever.build_tfidf_lines import OnlineTfidfDocRanker

def write_predictions(final_predictions, write_path):
    with open(write_path, "w+") as f:
        for prediction in final_predictions:
            f.write(json.dumps(prediction) + "\n")

def post_processing(clf, X, indexes, k=50):
    """
    predict scores for each claim and sentences in the predicted pages,
    get the order of score in descending format, and reorder (doc_id,sent_id) pair with the order,and extract first 5 most similar sentences
    :param clf:
    :param X:
    :param indexes:C
    :param k:
    :return:
    """

    predictions = []
    all_scores = []
    for idx, line_input in enumerate(X):
        scores = clf.predict(line_input)
        scores = np.asarray(np.reshape(scores, newshape=(-1,)))
        orders = np.argsort(scores, axis=-1)[::-1]
        scores = scores[orders]
        sents_indexes = np.asarray(indexes[idx])
        predictions.append(sents_indexes[orders][:k])
        all_scores.append(scores[:k])
    return predictions, all_scores


def post_processing_baseline(X, indexes, k=50):
    predictions = []
    all_scores = []
    for idx, line_input in enumerate(X):
        scores = np.random.uniform(low=-1, high=1, size=(len(line_input),))
        scores = np.asarray(np.reshape(scores, newshape=(-1,)))
        orders = np.argsort(scores, axis=-1)[::-1]
        scores = scores[orders]
        sents_indexes = np.asarray(indexes[idx])
        predictions.append(sents_indexes[orders][:k])
        all_scores.append(scores[:k])
    return predictions, all_scores


def get_whole_evidence(evidence_set, db):
    pos_sents = []
    for evidence in evidence_set:
        page = evidence[0]
        doc_lines = db[page]['lines']
        # doc_lines = db[page]
        for i in range(len(doc_lines)):
            if i == evidence[1]:
                pos_sents.append(doc_lines[i])
    pos_sent = ' '.join(pos_sents)
    return pos_sent
    
    
def prediction_processing(dataset_path, predictions, db_filename):
    """
    process the predicted (doc_id,sent_id) pairs to the score system desired format
    :param dataset_path:
    :param predictions:
    :return:
    """

    final_predictions = []
    jsr = JSONLineReader()
    with open(db_filename) as f: 
        db = json.load(f)
        
    out_error_ana = []   
    with open(dataset_path, "r") as f:
        lines = jsr.process(f)

        cnt = 0
        for line in lines:
            
            pos_set_ref = line['evidence']
            if len(pos_set_ref) == 0 or not pos_set_ref[0]:
                continue

            line['predicted_evidence'] = [[prediction[0], int(prediction[1])] for prediction in predictions[cnt]]
            cnt_gold = 0
            out_error_ana.append("Claim: "+str(cnt))
            out_error_ana.append(line['claim'])
            out_error_ana.append("Gold evidence:")
            for evidence_set in line['evidence']:
                for evidence_sentence in evidence_set:
                    out_error_ana.append(get_whole_evidence([evidence_sentence], db))
                    cnt_gold += 1
                    
            out_error_ana.append("Predicted evidence:")
            for evidence_set in line['predicted_evidence'][:cnt_gold]:
                out_error_ana.append(get_whole_evidence([evidence_set], db))
            out_error_ana.append("")

            line['predicted_label'] = "refutes"
            final_predictions.append(line)
            cnt += 1
            if cnt == len(predictions):
                break

    return final_predictions, out_error_ana


def post_processing_tfidf(X, indexes, k=50):
  
    predictions = []
    all_scores = []
    args.tokenizer = 'simple'
    args.ngram = 2
    args.max_sent = k
    args.hash_size = int(np.math.pow(2, 24))
    for idx, line_input in tqdm(enumerate(X)):
  
        sentences = [line_tmp[1] for line_tmp in line_input]
        claim = line_input[0][0]
        tfidf = OnlineTfidfDocRanker(args, sentences, None)
        line_ids,scores = tfidf.closest_docs(claim, args.max_sent)
        sents_indexes = np.asarray(indexes[idx])
  
        predictions.append(sents_indexes[line_ids][:k])
        all_scores.append(scores[:k])
    return predictions, all_scores

def pipeline(model="bilstm_ranking"):
    """
    the pipeline to store different kind of models, and their different kind of inputs
    :param model:
    :return:
    """
    from retrieval.sentences.data_processing.data import Data

    path = os.getcwd()
    path = path.replace("/src/retrieval/sentences", "")
    dev_path = os.path.join(path,        "data/datasets/snopes.evidence.dev.jsonl") 
    test_path = os.path.join(path,       "data/datasets/snopes.evidence.test.jsonl") 
    new_train_path = os.path.join(path,  "data/datasets/snopes.evidence.train.jsonl") 
    db_filename = os.path.join(path,     "data/datasets/snopes.page.json") 
    fasttext_path = os.path.join(path, "data/fasttext/wiki.en.bin")
    # fasttext_path = os.path.join(path, "data/fasttext/wiki-news-300d-1M.vec")
    
    random_state = 55
    tf.set_random_seed(random_state)
    np.random.seed(random_state)
    random.seed(random_state)
    
    print(fasttext_path)

    seed = 55
    tf.set_random_seed(seed)
    np.random.seed(seed)

    if model == "use_ranking": 
        s_max_length = 60
        h_max_length = 20
        num_negatives = 2
        data = Data(path, new_train_path, dev_path, test_path, fasttext_path, num_negatives=num_negatives, 
            h_max_length=h_max_length, s_max_length=s_max_length, random_seed=102, db_filepath=db_filename)
        
        clf = USERANKING(random_state=100, dropout_rate=0.1, learning_rate=0.001, num_epoch=1, batch_size=256, word_dict=data.iword_dict,
                 num_negative=num_negatives, activation=tf.nn.relu, h_max_length=h_max_length, s_max_length=s_max_length)
         
        clf.fit(data.X_train, data.X_dev, data.dev_indexes)
        predictions, scores = post_processing(clf, data.X_test, data.test_location_indexes)

        
    elif model == "random_baseline": 
        s_max_length = 60
        h_max_length = 20
        num_negatives = 1
        data = Data(path, new_train_path, dev_path, test_path, fasttext_path, num_negatives=num_negatives, 
            h_max_length=h_max_length, s_max_length=s_max_length, random_seed=102, db_filepath=db_filename)
        predictions, _ = post_processing_baseline(data.X_test, data.test_location_indexes)
 
 
    elif model == "tf_idf": 
        s_max_length = 20
        h_max_length = 20
        num_negatives = 1
        data = Data(path, new_train_path, dev_path, test_path, fasttext_path, num_negatives=num_negatives, 
            h_max_length=h_max_length, s_max_length=s_max_length, random_seed=140, db_filepath=db_filename)
        predictions, _ = post_processing_tfidf(data.X_test, data.test_location_indexes)

    elif model == "bilstm_ranking":
        s_max_length = 60
        h_max_length = 20
        data = Data(path, new_train_path, dev_path, test_path, fasttext_path, num_negatives=5, h_max_length=h_max_length,
                    s_max_length=s_max_length, random_seed=100, db_filepath=db_filename)
                                     

        clf = BiLSTM_RANKING(h_max_length=h_max_length, lstm_neurons=[64, 64, 64], s_max_length=s_max_length, learning_rate=0.001,
                             batch_size=256, num_epoch=2,
                             embedding=data.embed, word_dict=data.iword_dict, dropout_rate=0.1, random_state=55)


        clf.fit(data.X_train_indexes, data.dev_indexes, data.dev_labels)

        predictions, _ = post_processing(clf, data.test_indexes, data.test_location_indexes)


    elif model == "decompos_att":
        s_max_length = 60
        h_max_length = 20
        data = Data(path, new_train_path, dev_path, test_path, fasttext_path, num_negatives=5, h_max_length=h_max_length,
                    s_max_length=s_max_length, random_seed=100, db_filepath=db_filename)
        
        clf = Decomposable_Atten(h_max_length=h_max_length, s_max_length=s_max_length, learning_rate=0.0001, batch_size=128, num_epoch=10,
                                 embedding=data.embed, word_dict=data.iword_dict, dropout_rate=0.1, random_state=55)
        clf.fit(data.X_train_indexes, data.dev_indexes, data.dev_labels)
        predictions, _ = post_processing(clf, data.test_indexes, data.test_location_indexes)


    elif model == "esim":
        s_max_length = 60
        h_max_length = 20
        data = Data(path, new_train_path, dev_path, test_path, fasttext_path, num_negatives=5, h_max_length=h_max_length,
                    s_max_length=s_max_length, random_seed=131, db_filepath=db_filename)
    
        model_store_dir = "model/experiment3/"
        clf = ESIM(h_max_length=h_max_length, s_max_length=s_max_length, learning_rate=0.001, batch_size=256, num_epoch=2,
                   model_store_dir=model_store_dir,
                   embedding=data.embed, word_dict=data.iword_dict, dropout_rate=0.2, random_state=88, num_units=128,
                   activation=tf.nn.relu, share_rnn=True)
        
        clf.fit(data.X_train_indexes, data.dev_indexes, data.dev_labels)
        predictions, _ = post_processing(clf, data.test_indexes, data.test_location_indexes)

    final_predictions, _ = prediction_processing(test_path, predictions, db_filename)
    write_path = os.path.join(path, "data/predictions/model_"+model+".jsonl")
    write_predictions(final_predictions, write_path)

    strict_score, label_accuracy, precision, recall, f1 = fever_score(final_predictions, actual=None, max_evidence=5)
    print("strict_score: {} label_accuracy: {} precision: {} recall: {} f1: {} ".format(strict_score,
                                                                                                      label_accuracy,
                                                                                                      precision, recall,                                                                            f1))
if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', 
                        help='\'bilstm_ranking\', \'esim\' , \'decompos_att\', \'tf_idf\' or \'random_baseline\'', 
                        required=True)

    args = parser.parse_args()
    pipeline(model=args.model)
    



   
