import argparse
import json
from operator import itemgetter
import os

import stanza

from stanza.utils.default_paths import get_default_paths
from stanza.utils.get_tqdm import get_tqdm
from stanza.utils.datasets.coref.utils import process_document

tqdm = get_tqdm()

def flatten_spans(coref_spans):
    """
    Put span IDs on each span, then flatten them into a single list sorted by first word
    """
    # put span indices on the spans
    #   [[[38, 39], [42, 43], [41, 41], [180, 180], [300, 300]], [[60, 68],
    #   -->
    #   [[[0, 38, 39], [0, 42, 43], [0, 41, 41], [0, 180, 180], [0, 300, 300]], [[1, 60, 68], ...
    coref_spans = [[[span_idx, x, y] for x, y in span] for span_idx, span in enumerate(coref_spans)]
    # flatten list
    #   -->
    #   [[0, 38, 39], [0, 42, 43], [0, 41, 41], [0, 180, 180], [0, 300, 300], [1, 60, 68], ...
    coref_spans = [y for x in coref_spans for y in x]
    # sort by the first word index
    #   -->
    #   [[0, 38, 39], [0, 41, 41], [0, 42, 43], [1, 60, 68], [0, 180, 180], [0, 300, 300], ...
    coref_spans = sorted(coref_spans, key=itemgetter(1))
    return coref_spans

def remove_nulls(coref_spans, sentences):
    """
    Removes the "" and "NULL" words from the sentences

    Also, reindex the spans by the number of words removed.
    So, we might get something like
      [[0, 2], [31, 33], [134, 136], [161, 162]]
      ->
      [[0, 2], [30, 32], [129, 131], [155, 156]]
    """
    word_map = []
    word_idx = 0
    map_idx = 0
    new_sentences = []
    for sentence in sentences:
        new_sentence = []
        for word in sentence:
            word_map.append(map_idx)
            word_idx += 1
            if word != '' and word != 'NULL':
                new_sentence.append(word)
                map_idx += 1
        new_sentences.append(new_sentence)

    new_spans = []
    for mention in coref_spans:
        new_mention = []
        for span in mention:
            span = [word_map[x] for x in span]
            new_mention.append(span)
        new_spans.append(new_mention)
    return new_spans, new_sentences

def arrange_spans_by_sentence(coref_spans, sentences):
    sentence_spans = []

    current_index = 0
    span_idx = 0
    for sentence in sentences:
        current_sentence_spans = []
        end_index = current_index + len(sentence)
        while span_idx < len(coref_spans) and coref_spans[span_idx][1] < end_index:
            new_span = [coref_spans[span_idx][0], coref_spans[span_idx][1] - current_index, coref_spans[span_idx][2] - current_index]
            current_sentence_spans.append(new_span)
            span_idx += 1
        sentence_spans.append(current_sentence_spans)
        current_index = end_index
    return sentence_spans

def convert_dataset_section(pipe, section, use_cconj_heads):
    """
    Reprocess the original data into a format compatible with previous conversion utilities

    - remove blank and NULL words
    - rearrange the spans into spans per sentence instead of a list of indices for each span
    - process the document using a Hindi pipeline
    """
    processed_section = []

    for idx, doc in enumerate(tqdm(section)):
        doc_id = doc['doc_key']
        part_id = ""
        sentences = doc['sentences']
        sentence_speakers = doc['speakers']

        coref_spans = doc['clusters']
        coref_spans, sentences = remove_nulls(coref_spans, sentences)
        coref_spans = flatten_spans(coref_spans)
        coref_spans = arrange_spans_by_sentence(coref_spans, sentences)

        processed = process_document(pipe, doc_id, part_id, sentences, coref_spans, sentence_speakers, use_cconj_heads=use_cconj_heads)
        processed_section.append(processed)
    return processed_section

def remove_nulls_dataset_section(section):
    processed_section = []
    for doc in section:
        sentences = doc['sentences']
        coref_spans = doc['clusters']
        coref_spans, sentences = remove_nulls(coref_spans, sentences)
        doc['sentences'] = sentences
        doc['clusters'] = coref_spans
        processed_section.append(doc)
    return processed_section


def read_json_file(filename):
    with open(filename, encoding="utf-8") as fin:
        dataset = []
        for line in fin:
            line = line.strip()
            if not line:
                continue
            dataset.append(json.loads(line))
    return dataset

def write_json_file(output_filename, converted_section):
    with open(output_filename, "w", encoding="utf-8") as fout:
        json.dump(converted_section, fout, indent=2)

def main():
    parser = argparse.ArgumentParser(
        prog='Convert Hindi Coref Data',
    )
    parser.add_argument('--no_use_cconj_heads', dest='use_cconj_heads', action='store_false', help="Don't use the conjunction-aware transformation")
    parser.add_argument('--remove_nulls', action='store_true', help="The only action is to remove the NULLs and blank tokens")
    args = parser.parse_args()

    paths = get_default_paths()
    coref_input_path = paths["COREF_BASE"]
    hindi_base_path = os.path.join(coref_input_path, "hindi", "dataset")

    sections = ("train", "dev", "test")
    if args.remove_nulls:
        for section in sections:
            input_filename = os.path.join(hindi_base_path, "%s.hindi.jsonlines" % section)
            dataset = read_json_file(input_filename)
            dataset = remove_nulls_dataset_section(dataset)
            output_filename = os.path.join(hindi_base_path, "hi_deeph.%s.nonulls.json" % section)
            with open(output_filename, "w", encoding="utf-8") as fout:
                for doc in dataset:
                    json.dump(doc, fout, ensure_ascii=False)
                    fout.write("\n")
    else:
        pipe = stanza.Pipeline("hi", processors="tokenize,pos,lemma,depparse", package="default_accurate", tokenize_pretokenized=True)

        os.makedirs(paths["COREF_DATA_DIR"], exist_ok=True)

        for section in sections:
            input_filename = os.path.join(hindi_base_path, "%s.hindi.jsonlines" % section)
            dataset = read_json_file(input_filename)

            output_filename = os.path.join(paths["COREF_DATA_DIR"], "hi_deeph.%s.json" % section)
            converted_section = convert_dataset_section(pipe, dataset, use_cconj_heads=args.use_cconj_heads)
            write_json_file(output_filename, converted_section)

if __name__ == '__main__':
    main()
