#!/usr/bin/env python

from __future__ import division, unicode_literals
import os
import math
import codecs
import torch

import onmt
import onmt.io
import opts
from itertools import takewhile, count

from six.moves import zip_longest
from six.moves import zip

import nltk
import json, argparse, time
import re

from flask import Flask, request
from flask_cors import CORS

global dummy_parser
global dummy_opt
global translator
global pred_score_total
global pred_words_total

##################################################
# API part
##################################################
app = Flask(__name__)
cors = CORS(app)

parser = argparse.ArgumentParser(
    description='translate.py',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
opts.add_md_help_argument(parser)
opts.translate_opts(parser)

opt = parser.parse_args()
opt.model = 'sent2ques/model.pt'
opt.src = 'input.txt'
opt.output = 'pred_docker.txt'
opt.replace_unk = True
opt.verbose = True

def report_score(name, score_total, words_total):
    print("%s AVG SCORE: %.4f, %s PPL: %.4f" % (
        name, score_total / words_total,
        name, math.exp(-score_total/words_total)))

def get_src_words(src_indices, index2str):
    raw_words = (index2str[i] for i in src_indices)
    words = takewhile(lambda w: w != onmt.io.PAD_WORD, raw_words)
    return " ".join(words)

def rem_ne(sent):
    sentences = nltk.sent_tokenize(sent)
    tokenized = [nltk.word_tokenize(sentence) for sentence in sentences]
    pos_tags  = [nltk.pos_tag(sentence) for sentence in tokenized]
    test = list()
    ne_counter = 0
    ne_map = dict()

    new_pos_tags = list()
    i = 0
    pos_tags = pos_tags[0]
    while i < len(pos_tags):
        start_i = i
        new_pos_tag = pos_tags[i]
        if pos_tags[i][1] == 'NNP':
            while i + 1 < len(pos_tags) and pos_tags[i + 1][1] == 'NNP':
                i += 1
                new_pos_tag = (new_pos_tag[0] + ' ' + pos_tags[i][0], 'NNP')
        new_pos_tags.append(new_pos_tag)
        i += 1

    year_counter = 0
    for word in new_pos_tags:
        if len(word) == 2:
            if word[1] == 'CD' and word[0] != 'AAA':
                test.append('<year>')
                ne_map['year ' + str(year_counter)] = word[0]
                year_counter += 1
            elif word[1] == 'NNP' and word[0] != 'AAA':
                ne_map[ne_counter] = word[0]
                test.append('<' + str(ne_counter) + '>')
                ne_counter += 1
            else:
                test.append(word[0])
    test = ' '.join(test)
    return test, ne_map

def put_ne(sent, map):
    if 'AAA' in map:
        sent = sent.replace('AAA', map['AAA'])
    for key, value in map.items():
        if type(key) is int:
            sent = sent.replace('<' + str(key) + '>', value)
    return sent

def replacenth(string, sub, wanted, n):
    where = [m.start() for m in re.finditer(sub, string)][n - 1]
    before = string[:where]
    after = string[where + len(sub):]
    newString = before + wanted + after
    return newString

@app.route('/', methods=['POST'])
def generate():
    global dummy_parser
    global dummy_opt
    global translator
    global pred_score_total
    global pred_words_total

    start = time.time()

    data = request.data.decode("utf-8")

    if data == "":
        params = request.form
        sentences = json.loads(params['sent'])
    else:
        params = json.loads(data)
        sentences = params['sent']
    
    os.write(1, ('\n'.join(sentences)).encode('utf-8'))

    input_file = codecs.open(opt.src, 'w', 'utf-8')
    final_map = list()
    for sent in sentences:
        orig_sent = sent
        sent, ne_map = rem_ne(sent)
        for key, value in ne_map.items():
            new_map = ne_map.copy()
            if type(key) is int:
                ans = '<' + str(key) + '>'
                new_sent = sent.replace(ans, 'AAA')
                new_map['AAA'] = value
            elif 'year' in key:
                new_sent = replacenth(sent, '<year>', 'AAA', int(key[5:]))
                new_map['AAA'] = value
            final_map.append(new_map)

            input_file.write(new_sent + '\n')
    input_file.flush()

    questions = list()

    file_data = onmt.io.build_dataset(translator.fields, opt.data_type,
                                 opt.src, opt.tgt,
                                 src_dir=opt.src_dir,
                                 sample_rate=opt.sample_rate,
                                 window_size=opt.window_size,
                                 window_stride=opt.window_stride,
                                 window=opt.window,
                                 use_filter_pred=False)
    data_type = file_data.data_type

    test_data = onmt.io.OrderedIterator(
        dataset=file_data, device=opt.gpu,
        batch_size=opt.batch_size, train=False, sort=False,
        shuffle=False)

    counter = count(1)
    batch_num = 0
    for batch in test_data:
        pred_batch, gold_batch, pred_scores, gold_scores, attn, src, indices\
            = translator.translate(batch, file_data)
        pred_score_total += sum(score[0] for score in pred_scores)
        pred_words_total += sum(len(x[0]) for x in pred_batch)

        # z_batch: an iterator over the predictions, their scores,
        # the gold sentence, its score, and the source sentence for each
        # sentence in the batch. It has to be zip_longest instead of
        # plain-old zip because the gold_batch has length 0 if the target
        # is not included.
        if data_type == 'text':
            sents = src.split(1, dim=1)
        else:
            sents = [torch.Tensor(1, 1) for i in range(len(pred_scores))]
        z_batch = zip_longest(
                pred_batch, gold_batch,
                pred_scores, gold_scores,
                (sent.squeeze(1) for sent in sents), indices)

        for pred_sents, gold_sent, pred_score, gold_score, src_sent, index\
                in z_batch:
            n_best_preds = [" ".join(pred) for pred in pred_sents[:opt.n_best]]

            sent_number = next(counter)
            if data_type == 'text':
                words = get_src_words(
                    src_sent, translator.fields["src"].vocab.itos)
            else:
                words = test_data.dataset.examples[index].src_path

            output = '\nSENT {}: {}\n'.format(sent_number, words)
            os.write(1, output.encode('utf-8'))

            best_pred = n_best_preds[0]

            sent_words = words.split(' ')
            ques_words = best_pred.split(' ')
            result = ''
            best_pred_attn = attn[(sent_number - 1) % opt.batch_size][0]
            indexes = list()
            for key, value in final_map[sent_number - 1].items():
                if type(key) is str and 'year' in key:
                    indexes.append(int(key[5:]))
            year_best_pred_attn = list()
            for i in range(len(best_pred_attn)):
                year_best_pred_attn.append([best_pred_attn[i][j] for j in indexes])
            for i, word in enumerate(ques_words):
                if word == '<year>':
                    result += ' ' + final_map[sent_number - 1]['year ' + str(year_best_pred_attn[i].index(max(year_best_pred_attn[i])))]
                else: result += ' ' + word

            best_pred = result
            best_pred = put_ne(best_pred, final_map[sent_number - 1])
            # answer = final_map[sent_number - 1][final_map[sent_number - 1]['AAA'][1:-1]]
            questions.append({'question': best_pred, 'answer': final_map[sent_number - 1]['AAA']})

            best_score = pred_score[0]
            output = 'PRED {}: {}\n'.format(sent_number, best_pred)
            os.write(1, output.encode('utf-8'))
            print("PRED SCORE: {:.4f}".format(best_score))

            if len(n_best_preds) > 1:
                print('\nBEST HYP:')
                for score, sent in zip(pred_score, n_best_preds):
                    output = "[{:.4f}] {}\n".format(score, sent)
                    os.write(1, output.encode('utf-8'))

        batch_num += 1

    return json.dumps(questions)


def main():

    global dummy_parser
    global dummy_opt
    global translator
    global pred_score_total
    global pred_words_total
    # global out_file

    nltk.download('punkt')
    nltk.download('averaged_perceptron_tagger')

    dummy_parser = argparse.ArgumentParser(description='train.py')
    opts.model_opts(dummy_parser)
    dummy_opt = dummy_parser.parse_known_args([])[0]

    translator = onmt.Translator(opt, dummy_opt.__dict__)
    # out_file = codecs.open(opt.output, 'w', 'utf-8')
    pred_score_total, pred_words_total = 0, 0

    app.run(host="0.0.0.0", debug=True, port=8886)


if __name__ == "__main__":
    main()
