import re
import copy
import json
import time

import numpy as np
import xml.etree.ElementTree as ElTree

from datetime import datetime, timezone
from operator import itemgetter


tokenize_regex = re.compile(r'([0-9a-zA-Z]+|[^0-9a-zA-Z])')

def json_to_sent(data):
    '''data: list of json file [{pmid,abstract,title}, ...] '''
    out = dict()
    for paper in data:
        sentences = list()
        
        if len(CoNLL_tokenizer(paper['title'])) < 50:
            title = [paper['title']]
        else:
            title = sentence_split(paper['title'])
        if len(title) != 1 or len(title[0].strip()) > 0:
            sentences.extend(title)

        if len(paper['abstract']) > 0:
            abst = sentence_split(paper['abstract'])
            if len(abst) != 1 or len(abst[0].strip()) > 0:
                sentences.extend(abst)
        out[paper['pmid']] = dict()
        out[paper['pmid']]['sentence'] = sentences
    return out

def input_form(sent_data):
    '''sent_data: dict of sentence, key=pmid {pmid:[sent,sent, ...], pmid: ...}'''
    for pmid in sent_data:
        sent_data[pmid]['words'] = list()
        sent_data[pmid]['wordPos'] = list()
        doc_piv = 0
        for sent in sent_data[pmid]['sentence']:
            wids = list()
            wpos = list()
            sent_piv = 0
            tok = CoNLL_tokenizer(sent)

            for w in tok:
                if len(w) > 20:
                    wids.append(w[:10])
                else:
                    wids.append(w)

                start = doc_piv + sent_piv + sent[sent_piv:].find(w)
                end = start + len(w) - 1
                sent_piv = end - doc_piv + 1
                wpos.append((start, end))
            doc_piv += len(sent)
            sent_data[pmid]['words'].append(wids)
            sent_data[pmid]['wordPos'].append(wpos)

    return sent_data

def softmax(logits):
    out = list()
    for logit in logits:
        temp = np.subtract(logit, np.max(logit))
        p = np.exp(temp) / np.sum(np.exp(temp))
        out.append(np.max(p))
    return out

def CoNLL_tokenizer(text):
    rawTok = [t for t in tokenize_regex.split(text) if t]
    assert ''.join(rawTok) == text
    tok = [t for t in rawTok if t != ' ']
    return tok

def sentence_split(text):
    sentences = list()
    sent = ''
    piv = 0
    for idx, char in enumerate(text):
        if char in "?!":
            if idx > len(text) - 3:
                sent = text[piv:]
                piv = -1
            else:
                sent = text[piv:idx + 1]
                piv = idx + 1

        elif char == '.':
            if idx > len(text) - 3:
                sent = text[piv:]
                piv = -1
            elif (text[idx + 1] == ' ') and (
                    text[idx + 2] in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ-"' + "'"):
                sent = text[piv:idx + 1]
                piv = idx + 1

        if sent != '':
            toks = CoNLL_tokenizer(sent)
            if len(toks) > 100:
                while True:
                    rawTok = [t for t in tokenize_regex.split(sent) if t]
                    cut = ''.join(rawTok[:200])
                    sent = ''.join(rawTok[200:])
                    sentences.append(cut)

                    if len(CoNLL_tokenizer(sent)) < 100:
                        if sent.strip() == '':
                            sent = ''
                            break
                        else:
                            sentences.append(sent)
                            sent = ''
                            break
            else:
                sentences.append(sent)
                sent = ''

            if piv == -1:
                break

    if piv != -1:
        sent = text[piv:]
        toks = CoNLL_tokenizer(sent)
        if len(toks) > 100:
            while True:
                rawTok = [t for t in tokenize_regex.split(sent) if t]
                cut = ''.join(rawTok[:200])
                sent = ''.join(rawTok[200:])
                sentences.append(cut)

                if len(CoNLL_tokenizer(sent)) < 100:
                    if sent.strip() == '':
                        sent = ''
                        break
                    else:
                        sentences.append(sent)
                        sent = ''
                        break
        else:
            sentences.append(sent)
            sent = ''

    return sentences

def get_prob(data, sent_data, predicDict, logitsDict, entity_types=None):
    for idx, paper in enumerate(data):
        pmid = paper['pmid']
        
        if len(paper['abstract']) > 0:
            content = paper['title'] + ' ' + paper['abstract']
        else:
            content = paper['title']

        for ent_type in entity_types:
            paper['entities'][ent_type] = []
        paper['prob'] = dict()

        for dtype in entity_types:
            for sentidx, tags in enumerate(predicDict[dtype][pmid]):
                B_flag = False
                # get position of entity corresponding to types
                for widx, tag in enumerate(tags):
                    if tag == 'O':
                        if B_flag:
                            tmpSE["end"] = \
                            sent_data[pmid]['wordPos'][sentidx][widx - 1][1]
                            paper['entities'][dtype].append(tmpSE)
                        B_flag = False
                        continue
                    elif tag == 'B':
                        if B_flag:
                            tmpSE["end"] = \
                            sent_data[pmid]['wordPos'][sentidx][widx - 1][1]
                            paper['entities'][dtype].append(tmpSE)
                        tmpSE = {
                            "start": sent_data[pmid]['wordPos'][sentidx][widx][
                                0]}
                        B_flag = True
                    elif tag == "I":
                        continue
                if B_flag:
                    tmpSE["end"] = sent_data[pmid]['wordPos'][sentidx][-1][1]
                    paper['entities'][dtype].append(tmpSE)

            # get prob. of entity logits corresponding to types
            logs = list()
            for t_sent in logitsDict[dtype][pmid]:
                logs.extend(t_sent)
            paper['prob'][dtype] = list()
            for pos in paper['entities'][dtype]:
                if pos['start'] == pos['end']:
                    soft = softmax(logs[len(
                        CoNLL_tokenizer(content[:pos['start']])):len(
                        CoNLL_tokenizer(content[:pos['end']])) + 1])
                    paper['prob'][dtype].append(
                        (pos, float(np.average(soft))))
                else:
                    soft = softmax(logs[len(
                        CoNLL_tokenizer(content[:pos['start']])):len(
                        CoNLL_tokenizer(content[:pos['end']]))])
                    paper['prob'][dtype].append(
                        (pos, float(np.average(soft))))

    return data

def detokenize(tokens, predicts, logits):
    pred = dict({
        'toks': tokens[:],
        'labels': predicts[:],
        'logit': logits[:]
    })  # dictionary for predicted tokens and labels.

    bert_toks = list()
    bert_labels = list()
    bert_logits = list()
    tmp_p = list()
    tmp_l = list()
    tmp_s = list()
    for t, l, s in zip(pred['toks'], pred['labels'], pred['logit']):
        if t == '[CLS]' or t == '<s>':  # non-text tokens will not be evaluated.
            continue
        elif t == '[SEP]' or t == '</s>':  # newline
            bert_toks.append(tmp_p)
            bert_labels.append(tmp_l)
            bert_logits.append(tmp_s)
            tmp_p = list()
            tmp_l = list()
            tmp_s = list()
            continue
        elif t[:2] == '##':  # if it is a piece of a word (broken by Word Piece tokenizer)
            tmp_p[-1] = tmp_p[-1] + t[2:]  # append pieces
        elif t.startswith('Ġ'): # roberta tokenizer
            t = t.replace('Ġ', ' ')
            tmp_p[-1] = tmp_p[-1] + t
        else:
            tmp_p.append(t)
            tmp_l.append(l)
            tmp_s.append(s)
    return bert_toks, bert_labels, bert_logits

# https://stackoverflow.com/a/3620972
PROF_DATA = {}

class Profile(object):
    def __init__(self, prefix):
        self.prefix = prefix

    def __call__(self, fn):
        def with_profiling(*args, **kwargs):
            global PROF_DATA
            start_time = time.time()
            ret = fn(*args, **kwargs)

            elapsed_time = time.time() - start_time
            key = '[' + self.prefix + '].' + fn.__name__

            if key not in PROF_DATA:
                PROF_DATA[key] = [0, list()]
            PROF_DATA[key][0] += 1
            PROF_DATA[key][1].append(elapsed_time)

            return ret

        return with_profiling

def show_prof_data():
    for fname, data in sorted(PROF_DATA.items()):
        max_time = max(data[1])
        avg_time = sum(data[1]) / len(data[1])
        total_time = sum(data[1])
        print("\n{} -> called {} times".format(fname, data[0]))
        print("Time total: {:.3f}, max: {:.3f}, avg: {:.3f}".format(
            total_time, max_time, avg_time))

def clear_prof_data():
    global PROF_DATA
    PROF_DATA = {}

# Ref. dict of SR4GN
species_human_excl_homo_sapiens = \
    'person|infant|Child|people|participants|woman|' \
    'Girls|Man|Peoples|Men|Participant|Patients|' \
    'humans|Persons|mans|participant|Infants|Boys|' \
    'Human|Humans|Women|children|Mans|child|Participants|Girl|' \
    'Infant|girl|patient|patients|boys|men|infants|' \
    'man|girls|Children|Boy|women|persons|human|Woman|' \
    'peoples|Patient|People|boy|Person'.split('|')
    
def filter_entities(ner_results):
    num_filtered_species_per_doc = list()

    for idx, paper in enumerate(ner_results):

        if len(paper['abstract']) > 0:
            content = paper['title'] + ' ' + paper['abstract']
        else:
            content = paper['title']

        valid_species = list()
        species = paper['entities']['species']
        for spcs in species:
            entity_mention = content[spcs['start']:spcs['end']+1]
            if entity_mention in species_human_excl_homo_sapiens:
                spcs['end'] += 1
                continue
            valid_species.append(spcs)

        num_filtered_species = len(species) - len(valid_species)
        if num_filtered_species > 0:
            paper['entities']['species'] = valid_species

        num_filtered_species_per_doc.append((paper['pmid'],
                                             num_filtered_species))

    return num_filtered_species_per_doc

# from convert.py
def pubtator2dict_list(pubtator_file_path):
    dict_list = list()

    title_pmid = ''
    # abstract_pmid = ''
    title = ''
    abstract_text = ''
    doc_line_num = 0

    with open(pubtator_file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.rstrip()
            if len(line) == 0:
                               
                doc_dict = {
                    'pmid': title_pmid,
                    'entities': {},
                }
                doc_dict['title'] = title
                doc_dict['abstract'] = abstract_text

                dict_list.append(doc_dict)

                doc_line_num = 0
                continue

            if doc_line_num == 0:
                title_cols = line.split('|t|')

                if len(title_cols) != 2:
                    return '{"error": "wrong #title_cols {}"}'\
                        .format(len(title_cols))

                title_pmid = title_cols[0]

                if '- No text -' == title_cols[1]:
                    # make tmvar2 results empty
                    title = ''
                else:
                    title = title_cols[1]
            elif doc_line_num == 1:
                abstract_cols = line.split('|a|')

                if len(abstract_cols) != 2:
                    if len(abstract_cols) > 2:
                        abstract_text = "|a|".join(abstract_cols[1:])
                    else:
                        return '{"error": "wrong #abstract_cols {}"}'.format(len(abstract_cols))
                else:
                    if '- No text -' == abstract_cols[1]:
                        # make tmvar2 results empty
                        abstract_text = ''
                    else:
                        abstract_text = abstract_cols[1]

            doc_line_num += 1
    return dict_list

def preprocess(text):
    text = text.replace('\r ', ' ')

    text = text.replace('\u2028', ' ')
    text = text.replace('\u2029', ' ')

    # HAIR SPACE
    # https://www.fileformat.info/info/unicode/char/200a/index.htm
    text = text.replace('\u200A', ' ')

    # THIN SPACE
    # https://www.fileformat.info/info/unicode/char/2009/index.htm
    text = text.replace('\u2009', ' ')
    text = text.replace('\u2008', ' ')

    # FOUR-PER-EM SPACE
    # https://www.fileformat.info/info/unicode/char/2005/index.htm
    text = text.replace('\u2005', ' ')
    text = text.replace('\u2004', ' ')
    text = text.replace('\u2003', ' ')

    # EN SPACE
    # https://www.fileformat.info/info/unicode/char/2002/index.htm
    text = text.replace('\u2002', ' ')

    # NO-BREAK SPACE
    # https://www.fileformat.info/info/unicode/char/00a0/index.htm
    text = text.replace('\u00A0', ' ')

    # https://www.fileformat.info/info/unicode/char/f8ff/index.htm
    text = text.replace('\uF8FF', ' ')

    # https://www.fileformat.info/info/unicode/char/202f/index.htm
    text = text.replace('\u202F', ' ')

    text = text.replace('\uFEFF', ' ')
    text = text.replace('\uF044', ' ')
    text = text.replace('\uF02D', ' ')
    text = text.replace('\uF0BB', ' ')

    text = text.replace('\uF048', 'Η')
    text = text.replace('\uF0B0', '°')

    # MIDLINE HORIZONTAL ELLIPSIS: ⋯
    # https://www.fileformat.info/info/unicode/char/22ef/index.htm
    # text = text.replace('\u22EF', '...')

    return text