import os
import pickle
import sys
from concurrent.futures import ProcessPoolExecutor

from jack.util.preprocessing import sort_by_tfidf


def extract_support(triviaqa_question, docs, corpus, max_num_support, max_tokens):
    answers = []
    supports = []
    paragraph_tokens = []
    separator = '$|$'
    for doc in docs:
        doc_tokens = corpus.get_document(doc.doc_id)
        doc_tokens_flat = [t for p in doc_tokens for s in p for t in s]
        doc_paragraph_tokens = [[t for s in p for t in s] for p in doc_tokens]

        # merge many small paragraphs
        if max_tokens > 0:
            new_paragraph_tokens = [[]]
            for s in doc_paragraph_tokens:
                if len(new_paragraph_tokens[-1]) + len(s) >= max_tokens and len(new_paragraph_tokens[-1]) > 0:
                    # start new paragraph
                    if len(s) >= max_tokens:
                        while s:
                            new_paragraph_tokens.append(s[:max_tokens])
                            s = s[max_tokens:]
                    else:
                        new_paragraph_tokens.append(s)
                else:
                    # merge with recent paragraph
                    if len(new_paragraph_tokens[-1]) > 0:
                        new_paragraph_tokens[-1].append(separator)
                    new_paragraph_tokens[-1].extend(s)
        else:
            new_paragraph_tokens = doc_paragraph_tokens
        paragraph_tokens.extend(new_paragraph_tokens)

        p_idx_flat = [i for i, p in enumerate(new_paragraph_tokens) for t in p if t != separator]
        assert len(doc_tokens_flat) == len(p_idx_flat)

        doc_idx_offset = len(supports)
        supports.extend(" ".join(s) for s in new_paragraph_tokens)
        support_offsets = [0]
        for s in supports[doc_idx_offset:]:
            support_offsets.append(support_offsets[-1] + len(s) + 1)
        if doc.answer_spans is not None:
            for flat_s, flat_e in doc.answer_spans:
                p_idx = p_idx_flat[flat_s]
                s = flat_s - sum(1 for p2 in new_paragraph_tokens[:p_idx] for t in p2 if t != separator)
                p = new_paragraph_tokens[p_idx]
                k = 0
                char_s = 0
                while s > k:
                    if p[k] == separator:
                        s += 1
                    char_s += len(p[k]) + 1
                    k += 1
                char_e = char_s + sum(len(t) + 1 for t in doc_tokens_flat[flat_s:flat_e + 1]) - 1
                answers.append({
                    "text": " ".join(doc_tokens_flat[flat_s:flat_e + 1]),
                    "span": [char_s, char_e],
                    "doc_idx": p_idx + doc_idx_offset
                })
        del doc_tokens, p_idx_flat, doc_tokens_flat

    if max_num_support > 0 and len(supports) > max_num_support:
        sorted_supports = sort_by_tfidf(" ".join(triviaqa_question.question), [' '.join(p) for p in paragraph_tokens])
        sorted_supports = [i for i, _ in sorted_supports]
        sorted_supports_rev = {v: k for k, v in enumerate(sorted_supports)}
        if answers:
            min_answer_rev = min(sorted_supports_rev[a['doc_idx']] for a in answers)
            if min_answer_rev >= max_num_support:
                min_answer = sorted_supports[min_answer_rev]
                # force at least one answer by swapping best paragraph with answer to be the n-th paragraph
                old_nth_best = sorted_supports[max_num_support - 1]
                sorted_supports[min_answer_rev] = sorted_supports[max_num_support - 1]
                sorted_supports[max_num_support - 1] = min_answer
                sorted_supports_rev[old_nth_best] = min_answer_rev
                sorted_supports_rev[min_answer] = max_num_support - 1

        sorted_supports_rev = {v: k for k, v in enumerate(sorted_supports)}
        supports = [supports[i] for i in sorted_supports[:max_num_support]]
        is_an_answer = len(answers) > 0
        answers = [a for a in answers if sorted_supports_rev[a['doc_idx']] < max_num_support]
        for a in answers:
            a['doc_idx'] = sorted_supports_rev[a['doc_idx']]
        assert not is_an_answer or len(answers) > 0

    return supports, answers


def convert_triviaqa(triviaqa_question, corpus, max_num_support, max_tokens, is_web):
    question = " ".join(triviaqa_question.question)
    if is_web:
        for doc in triviaqa_question.web_docs:
            supports, answers = extract_support(triviaqa_question, [doc], corpus, max_num_support, max_tokens)
            filename = corpus.file_id_map[doc.doc_id]
            question_id = triviaqa_question.question_id + '--' + filename[4:] + ".txt"
            yield {"questions": [{"answers": answers,
                                  "question": {"text": question, "id": question_id}}],
                   "support": supports}
        for doc in triviaqa_question.entity_docs:
            supports, answers = extract_support(triviaqa_question, [doc], corpus, max_num_support, max_tokens)
            filename = corpus.file_id_map[doc.doc_id]
            question_id = triviaqa_question.question_id + '--' + filename[filename.rindex('/') + 1] + ".txt"
            yield {"questions": [{"answers": answers,
                                  "question": {"text": question, "id": question_id}}],
                   "support": supports}
    else:
        question_id = triviaqa_question.question_id
        supports, answers = extract_support(triviaqa_question, triviaqa_question.entity_docs,
                                            corpus, max_num_support, max_tokens)
        yield {"questions": [{"answers": answers,
                              "question": {"text": question, "id": question_id}}],
               "support": supports}


def process(x, verbose=False):
    dataset, filemap, max_num_support, max_tokens, is_web = x
    instances = []
    corpus = TriviaQaEvidenceCorpusTxt(filemap)
    for i, q in enumerate(dataset):
        if verbose and i % 1000 == 0:
            print("%d/%d done" % (i, len(dataset)))
        instances.extend(x for x in convert_triviaqa(q, corpus, max_num_support, max_tokens, is_web))
    return instances


def convert_dataset(path, filemap, name, num_processes, max_num_support, max_tokens, is_web=True):
    with open(path, 'rb') as f:
        dataset = pickle.load(f)

    if num_processes == 1:
        instances = process((dataset, filemap, max_num_support, max_tokens, is_web), True)
    else:
        chunk_size = 1000
        executor = ProcessPoolExecutor(num_processes)
        instances = []
        i = 0
        for processed in executor.map(
                process, [(dataset[i * chunk_size:(i + 1) * chunk_size], filemap, max_num_support, max_tokens, is_web)
                          for i in range(len(dataset) // chunk_size + 1)]):
            instances.extend(processed)
            i += chunk_size
            print("%d/%d done" % (min(len(dataset), i), len(dataset)))

    return {"meta": {"source": name}, 'instances': instances}


if __name__ == '__main__':
    from docqa.triviaqa.evidence_corpus import TriviaQaEvidenceCorpusTxt
    import json

    dataset = sys.argv[1]

    if len(sys.argv) > 2:
        num_processes = int(sys.argv[2])
    else:
        num_processes = 1

    if len(sys.argv) > 3:
        max_paragraphs = int(sys.argv[3])
    else:
        max_paragraphs = -1

    if len(sys.argv) > 4:
        max_tokens = int(sys.argv[4])
    else:
        max_tokens = -1

    triviaqa_prepro = os.environ['TRIVIAQA_HOME'] + '/preprocessed'

    is_web = dataset.startswith('web')
    dataset, split = dataset.split('-')

    ds = os.path.join(triviaqa_prepro, 'triviaqa/', dataset)
    with open(ds + "/file_map.json") as f:
        filemap = json.load(f)

    fn = '%s-%s.json' % (dataset, split)
    print("Converting %s..." % fn)
    new_ds = convert_dataset(os.path.join(ds, split + '.pkl'), filemap, fn, num_processes,
                             max_paragraphs, max_tokens, is_web)
    with open('data/triviaqa/%s' % fn, 'w') as f:
        json.dump(new_ds, f)
