from io import open
import sys
import re
from conllu import parse_incr
from conllu import parse

base_filename = sys.argv[1]
conll_filename = base_filename + ".conllu"
pattern_filename = base_filename + ".patterns"
print(conll_filename, file=sys.stderr)
print(pattern_filename, file=sys.stderr)

def read_pattern_file(filename):
    heads = {}
    labels = {}
    with open(filename) as f:
        n = 0
        for line in f:
            if line[0] == '#':
                continue
            line = line.strip()
            tokens = line.split()
            if n % 3 == 0: #pos
                key = re.sub(" +", " ", line)

#                print("key = ", key, file=sys.stderr)
                heads[key] = ""
                labels[key] = ""
            elif n % 3 == 1 : # heads
                heads[key] = tokens
#                print("heads = ", heads[key], file=sys.stderr)    
            elif n % 3 == 2 : # heads
                labels[key] = tokens
#                print("labels = ", labels[key], file=sys.stderr)
            n+= 1

    return heads, labels  


heads, labels = read_pattern_file(pattern_filename)


with open(conll_filename, "r", encoding="utf-8") as data_file :
    for tokenlist in parse_incr(data_file):
        print(tokenlist.serialize(), file=sys.stderr)
        pos_list = ""
        l = 0
        for token in tokenlist:
            l += 1
            pos_list += str(l)
            pos_list += ':'
            if token["upos"] == "gogo":# or token["upos"] == "ADP":
                pos_list += token["form"]
            else :
                pos_list += token["upos"]
            pos_list += " "
        pos_list = pos_list.rstrip()
        print(pos_list, file=sys.stderr)

        if pos_list not in heads :
            print("cannot find an entry in heads for ", pos_list, file=sys.stderr)

        if pos_list not in labels :
            print("cannot find an entry in labels for ", pos_list, file=sys.stderr)

        head_list = heads[pos_list]
        label_list = labels[pos_list]

        for index, token in enumerate(tokenlist):
            if token["head"] != int(head_list[index]) and (tokenlist[index-1]["upos"] == "ADP" or tokenlist[index-2]["upos"]):
                print("PP ATTACH token", index + 1, "actual head : ", token["head"], "correct head : ", head_list[index], file=sys.stderr)
            else:
                if token["head"] != int(head_list[index]):
                    print("ERROR token", index + 1, "actual head : ", token["head"], "correct head : ", head_list[index], file=sys.stderr)
                    token["head"] = head_list[index]

                if token["deprel"] != label_list[index]:
                    print("ERROR token ", index + 1, "actual label : ", token["deprel"], "correct label : ", label_list[index], file=sys.stderr)
                    token["deprel"] = label_list[index]
        print(tokenlist.serialize(), end ='')

