# -*- coding: utf-8 -*-
import argparse
import os
import random
from six.moves import urllib
import tarfile
from pathlib import Path
import csv

import nltk
from nltk.stem import SnowballStemmer, WordNetLemmatizer

__author__ = 'rkadlec'

"""
Script for generation of train, test and valid datasets from Ubuntu Corpus 1 on 1 dialogs.
Copyright IBM Corporation 2016
LICENSE: Apache License 2.0  URL: ttp://www.apache.org/licenses/LICENSE-2.0
Contact: Rudolf Kadlec (rudolf_kadlec@cz.ibm.com)
"""

"""
This is a modified version of the script that has a sub parser to determine the desired number of turns.
"""

dialog_end_symbol = "__dialog_end__"
end_of_utterance_symbol = "__eou__"
end_of_turn_symbol = "__eot__"

def translate_dialog_to_lists(dialog_filename):
    """
    Translates the dialog to a list of lists of utterances. In the first
    list each item holds subsequent utterances from the same user. The second level
    list holds the individual utterances.
    :param dialog_filename:
    :return:
    """

    # go through the dialog
    first_turn = True
    dialog = []
    same_user_utterances = []
    #last_user = None
    dialog.append(same_user_utterances)

    with open(dialog_filename, encoding='UTF8') as csvfile:
        dialog_reader = csv.reader(csvfile, delimiter='\t', quotechar=None)
        for dialog_line in dialog_reader:
            if first_turn:
                last_user = dialog_line[1]
                first_turn = False

            if last_user != dialog_line[1]:
                # user has changed
                same_user_utterances = []
                dialog.append(same_user_utterances)

            same_user_utterances.append(dialog_line[3])
            last_user = dialog_line[1]

    dialog.append([dialog_end_symbol])
    return dialog


def get_random_utterances_from_corpus(candidate_dialog_paths,rng,utterances_num=9,min_turn=3,max_turn=20):
    """
    Sample multiple random utterances from the whole corpus.
    :param candidate_dialog_paths:
    :param rng:
    :param utterances_num: number of utterances to generate
    :param min_turn: minimal index of turn that the utterance is selected from
    :return:
    """
    utterances = []
    dialogs_num = len(candidate_dialog_paths)

    for i in range(0,utterances_num):
        # sample random dialog
        dialog_path = candidate_dialog_paths[rng.randint(0,dialogs_num-1)]
        # load the dialog
        dialog = translate_dialog_to_lists(dialog_path)

        # we do not count the last  _dialog_end__ urn
        dialog_len = len(dialog) - 1
        if(dialog_len<min_turn):
            print("Dialog {} was shorter than the minimum required lenght {}".format(dialog_path,dialog_len))
            exit()
        # sample utterance, exclude the last round that is always "dialog end"
        max_ix = min(max_turn, dialog_len) - 1

        # this code deals with corner cases when dialogs are too short
        if min_turn -1 == max_ix:
            turn_index = max_ix
        else:
            turn_index = rng.randint(min_turn,max_ix)

        utterance = singe_user_utterances_to_string(dialog[turn_index])
        utterances.append(utterance)
    return utterances

def singe_user_utterances_to_string(utterances_list):
    """
    Concatenates multiple user's utterances into a single string.
    :param utterances_list:
    :return:
    """
    return " ".join([x+" "+ end_of_utterance_symbol for x in utterances_list])

def dialog_turns_to_string(dialog):
    """
    Translates the whole dialog (list of lists) to a string
    :param dialog:
    :return:
    """
    # join utterances
    turns_as_strings = list(map(singe_user_utterances_to_string,dialog))
    # join turns
    return "".join([x + " " + end_of_turn_symbol + " " for x in turns_as_strings])


#Function that creates a context randomly
#dialog is a list containing the entire conversation session
#rng object used to generate random numbers
#also receives the maximum and minimum size of a context
def create_random_context(dialog,rng,minimum_context_length=2,max_context_length=20, num=1):
    """
    Samples random context from a dialog. Contexts are uniformly sampled from the whole dialog.
    :param dialog:
    :param rng:
    :return: context, index of next utterance that follows the context
    """
    # sample dialog context
    #context_turns = rng.randint(minimum_context_length,len(dialog)-1)
    max_len = min(max_context_length, len(dialog)) - 2
    if max_len <= minimum_context_length:
        context_turns = max_len
    else:
        context_turns = rng.randint(minimum_context_length,max_len)

    # create string
    return dialog_turns_to_string(dialog[:context_turns]),context_turns

def create_single_dialog_test_example(context_dialog_path, candidate_dialog_paths, rng, distractors_num, max_context_length, num):
    """
    Creates a single example for testing or validation. Each line contains a context, one positive example and N negative examples.
    :param context_dialog_path:
    :param candidate_dialog_paths:
    :param rng:
    :param distractors_num:
    :return: triple (context, positive response, [negative responses])
    """

    dialog = translate_dialog_to_lists(context_dialog_path)

    context_str, next_utterance_ix = create_random_context(dialog, rng, max_context_length=max_context_length, num=num)

    # use the next utterance as positive example
    positive_response = singe_user_utterances_to_string(dialog[next_utterance_ix])

    negative_responses = get_random_utterances_from_corpus(candidate_dialog_paths,rng,distractors_num)
    return context_str, positive_response, negative_responses

#function that creates training examples based on the desired number of turns
#unification of three functions (create_random_context, create_single_dialog_train_example and create_examples) from the original code, plus some modifications
def create_examples_train(candidate_dialog_paths, examples_num, rng, positive_probability, num, minimum_context_length=2,max_context_length=20):
    """
    Creates a list of training examples from a list of dialogs and function that transforms a dialog to an example.
    :param candidate_dialog_paths:
    :param creator_function:
    :return:
    """
    i = 0
    x = 0
    dialogOriginal = ''
    examples = []
    unique_dialogs_num = len(candidate_dialog_paths)
    num = num * 2
    #A turn is a text input (consisting of one or more utterances) accompanied by a response (consisting of one or more utterances)
    while i < examples_num:
        while True:
            context_dialog = candidate_dialog_paths[x % unique_dialogs_num]
            dialogOriginal = translate_dialog_to_lists(context_dialog)
            max_len = min(max_context_length, len(dialogOriginal)) - 2
            x += 1
            if max_len >= num:
                break

        context_turns = rng.randint(num,max_len)
        dialog = dialogOriginal[:context_turns]
        dialog = dialog[-num:]

        context_str = dialog_turns_to_string(dialog)
        next_utterance_ix = context_turns #the next utterence that will be used as a response to the context

        if positive_probability > rng.random():
            response = singe_user_utterances_to_string(dialogOriginal[next_utterance_ix])
            label = 1.0
        else:
            response = get_random_utterances_from_corpus(candidate_dialog_paths,rng,1,
                                                     min_turn=minimum_context_length+1,
                                                     max_turn=max_context_length)[0]
            label = 0.0
        if i % 1000 == 0:
            print(str(i))
        i+=1
        
        result = [context_str, response , label]
        examples.append(result)

    return examples


def convert_csv_with_dialog_paths(csv_file):
    """
    Converts CSV file with comma separated paths to filesystem paths.
    :param csv_file:
    :return:
    """
    def convert_line_to_path(line):
        file, dir = [x.strip() for x in line.split(",")]
        return os.path.join(dir, file)

    return list(map(convert_line_to_path, csv_file))


def prepare_data_maybe_download(directory):
  """
  Download and unpack dialogs if necessary.
  """
  filename = 'ubuntu_dialogs.tgz'
  url = 'http://cs.mcgill.ca/~jpineau/datasets/ubuntu-corpus-1.0/ubuntu_dialogs.tgz'
  dialogs_path = os.path.join(directory, 'dialogs')

  # test it there are some dialogs in the path
  if not os.path.exists(os.path.join(directory,"10","1.tst")):
    # dialogs are missing
    archive_path = os.path.join(directory,filename)
    if not os.path.exists(archive_path):
        # archive missing, download it
        print(("Downloading %s to %s" % (url, archive_path)))
        filepath, _ = urllib.request.urlretrieve(url, archive_path)
        print("Successfully downloaded " + filepath)

    # unpack data
    if not os.path.exists(dialogs_path):
          print("Unpacking dialogs ...")
          with tarfile.open(archive_path) as tar:
                tar.extractall(path=directory)
          print("Archive unpacked.")

    return

#####################################################################################
# Command line script related code
#####################################################################################

if __name__ == '__main__':


    def create_eval_dataset(args, file_list_csv):
        rng = random.Random(args.seed)
        # training dataset
        f = open(os.path.join("meta", file_list_csv), 'r')
        dialog_paths = [os.path.join(args.data_root, "dialogs", path) for path in convert_csv_with_dialog_paths(f)]

        data_set = create_examples(dialog_paths,
                                   len(dialog_paths),
                                   lambda context_dialog, candidates : create_single_dialog_test_example(context_dialog, candidates, rng,
                                                                     args.n, args.max_context_length, args.turn))
        # output the dataset
        with open(args.output, 'w', encoding='UTF8') as f:
            w = csv.writer(f)
            header = ["Context", "Ground Truth Utterance"]
            header.extend(["Distractor_{}".format(x) for x in range(args.n)])
            w.writerow(header)

            stemmer = SnowballStemmer("english")
            lemmatizer = WordNetLemmatizer()

            for row in data_set:
                translated_row = [row[0], row[1]]
                translated_row.extend(row[2])
                
                if args.tokenize:
                    translated_row = list(map(nltk.word_tokenize, translated_row))
                    if args.stem:
                        translated_row = [list(map(stemmer.stem, sub)) for sub in translated_row]
                    if args.lemmatize:
                        translated_row = [[lemmatizer.lemmatize(tok, pos='v') for tok in sub] for sub in translated_row]
                        
                    translated_row = [" ".join(x) for x in translated_row]

                w.writerow(translated_row)
        print(("Dataset stored in: {}".format(args.output)))


    def train_cmd(args):
        rng = random.Random(args.seed)

        f = open(Path(os.path.join("meta", "trainfiles.csv")))
        dialog_paths = [os.path.join(args.data_root, "dialogs", path) for path in convert_csv_with_dialog_paths(f)]

        train_set = create_examples_train(dialog_paths,args.e, rng, args.p,  args.turn, max_context_length=args.max_context_length)

        stemmer = SnowballStemmer("english")
        lemmatizer = WordNetLemmatizer()

        # output the dataset
        with open(args.output, 'w', encoding='UTF8') as f:
            w = csv.writer(f)
            w.writerow(["Context", "Utterance", "Label"])
            for row in train_set:
                translated_row = row

                if args.tokenize:
                    translated_row = [nltk.word_tokenize(row[i]) for i in [0,1]]

                    if args.stem:
                        translated_row = [list(map(stemmer.stem, sub)) for sub in translated_row]
                    if args.lemmatize:
                        translated_row = [[lemmatizer.lemmatize(tok, pos='v') for tok in sub] for sub in translated_row]

                    translated_row = [" ".join(x) for x in translated_row]
                    translated_row.append(int(float(row[2])))

                w.writerow(translated_row)
        print(("Train dataset stored in: {}".format(args.output)))

    def valid_cmd(args):
        create_eval_dataset(args, "valfiles.csv")

    def test_cmd(args):
        create_eval_dataset(args, "testfiles.csv")


    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                     description="Script that creates train, valid and test set from 1 on 1 dialogs in Ubuntu Corpus. " +
                                                 "The script downloads 1on1 dialogs from internet and then it randomly samples all the datasets with positive and negative examples.")

    parser.add_argument('--data_root', default='.',
                        help='directory where 1on1 dialogs will be downloaded and extracted, the data will be downloaded from cs.mcgill.ca/~jpineau/datasets/ubuntu-corpus-1.0/ubuntu_dialogs.tgz')

    parser.add_argument('--seed', type=int, default=1234,
                        help='seed for random number generator')

    parser.add_argument('--max_context_length', type=int, default=20,
                        help='maximum number of dialog turns in the context')

    parser.add_argument('-o', '--output', default=None,
                        help='output csv')

    parser.add_argument('-t', '--tokenize', action='store_true',
                        help='tokenize the output')

    parser.add_argument('-l', '--lemmatize', action='store_true',
                        help='lemmatize the output by nltk.stem.WorldNetLemmatizer (applied only when -t flag is present)')

    parser.add_argument('-s', '--stem', action='store_true',
                        help='stem the output by nltk.stem.SnowballStemmer (applied only when -t flag is present)')

    subparsers = parser.add_subparsers(help='sub-command help')

    parser_train = subparsers.add_parser('train', help='trainset generator')
    parser_train.add_argument('-p', type=float, default=0.5, help='positive example probability')
    parser_train.add_argument('-e', type=int, default=1000000, help='number of examples to generate')
    parser_train.add_argument('-t', '--turn', type=int, default=1, help='number of turns per examples to generate')
    parser_train.set_defaults(func=train_cmd)

    parser_test = subparsers.add_parser('test', help='testset generator')
    parser_test.add_argument('-n', type=int, default=9, help='number of distractor examples for each context')
    parser_test.set_defaults(func=test_cmd)

    parser_valid = subparsers.add_parser('valid', help='validset generator')
    parser_valid.add_argument('-n', type=int, default=9, help='number of distractor examples for each context')
    parser_valid.set_defaults(func=valid_cmd)

    args = parser.parse_args()

    # download and unpack data if necessary
    prepare_data_maybe_download(args.data_root)

    # create dataset
    args.func(args)
