from __future__ import absolute_import, division, print_function

import re
import numpy as np
import spacy

def load_vocab_dict_from_file(dict_file):
    with open(dict_file) as f:
        words = [w.strip() for w in f.readlines()]
    vocab_dict = {words[n]:n for n in range(len(words))}
    return vocab_dict

UNK_IDENTIFIER = '<unk>' # <unk> is the word used to identify unknown words
SENTENCE_SPLIT_REGEX = re.compile(r'(\W+)')
def sentence2vocab_indices(sentence, vocab_dict):
    words = SENTENCE_SPLIT_REGEX.split(sentence.strip())
    words = [w.lower() for w in words if len(w.strip()) > 0]
    # remove .
    if words[-1] == '.':
        words = words[:-1]
    vocab_indices = [(vocab_dict[w] if w in vocab_dict else vocab_dict[UNK_IDENTIFIER])
        for w in words]
    return vocab_indices

PAD_IDENTIFIER = '<pad>'
EOS_IDENTIFIER = '<eos>'
def preprocess_sentence(sentence, vocab_dict, T):
    vocab_indices = sentence2vocab_indices(sentence, vocab_dict)
    # # Append '<eos>' symbol to the end
    # vocab_indices.append(vocab_dict[EOS_IDENTIFIER])
    # Truncate long sentences
    if len(vocab_indices) > T:
        vocab_indices = vocab_indices[:T]
    # Pad short sentences at the beginning with the special symbol '<pad>'
    if len(vocab_indices) < T:
        vocab_indices = [vocab_dict[PAD_IDENTIFIER]] * (T - len(vocab_indices)) + vocab_indices
    return vocab_indices

def get_tree_height(root, depth):
    max_depth = depth
    for c in root.children:
        cur_depth = get_tree_height(c, depth + 1)
        if cur_depth > max_depth:
            max_depth = cur_depth
    return max_depth

def sentence2token_indices(spacy_doc, vocab_dict, T=30):

    tokens = []
    graph = np.zeros((T, T), dtype=np.float32)
    for i in range(T):
        for j in range(T):
            graph[i, j] = 0.1

    root = None
    for token in spacy_doc:
        if token.dep_ == u"ROOT":
            root = token
        try:
            tokens.append(str(token.text))
        except UnicodeEncodeError:
            tokens.append(UNK_IDENTIFIER)
        idx = token.i
        for c in token.children:
            graph[idx, c.i] = 1
            graph[c.i, idx] = 1

    height = get_tree_height(root, 0)

    vocab_indices = [(vocab_dict[w] if w in vocab_dict else vocab_dict[UNK_IDENTIFIER])
        for w in tokens]
    return vocab_indices, graph, height


def preprocess_spacy_sentence(spacy_doc, vocab_dict, T):

    vocab_indices, graph, height = sentence2token_indices(spacy_doc, vocab_dict, T)
    # # Append '<eos>' symbol to the end
    # vocab_indices.append(vocab_dict[EOS_IDENTIFIER])
    # Truncate long sentences
    if len(vocab_indices) > T:
        vocab_indices = vocab_indices[:T]
    # Pad short sentences at the beginning with the special symbol '<pad>'
    if len(vocab_indices) < T:
        vocab_indices = [vocab_dict[PAD_IDENTIFIER]] * (T - len(vocab_indices)) + vocab_indices

    return vocab_indices, graph, height
