from STSB_Dataset import STSB_Dataset
import pandas as pd
import spacy
import string
import random

df = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe

nlp = spacy.load("en_core_web_sm")

sentences1 = df['sentence1'].tolist()
sentences2 = df['sentence2'].tolist()

sentences1_nlp = [nlp(sent) for sent in sentences1]
sentences2_nlp = [nlp(sent) for sent in sentences2]

dependency_hierarchy1 = []
dependency_hierarchy2 = []

sent1roots = []
sent2roots= []


for sent1 in sentences1_nlp:

    sentence_tokens = [token for token in sent1]
    dep_hierarchy1 = []
    dep_hierarchy1.append([token for token in sent1 if token.dep_ == "ROOT"])
    sent1roots.append([token.text for token in sent1 if token.dep_ == "ROOT"][0])
    sentence_tokens.remove(dep_hierarchy1[0][0])

    while len(sentence_tokens) != 0:
        children = []
        for i in range(len(dep_hierarchy1[-1])):
            index_children = [t for t in sent1 if t.head.text == dep_hierarchy1[-1][i].text and t in sentence_tokens]
            for elem in index_children:
                children.append(elem)
                sentence_tokens.remove(elem)
        dep_hierarchy1.append(children)

    dependency_hierarchy1.append(dep_hierarchy1)


for sent2 in sentences2_nlp:

    sentence_tokens = [token for token in sent2]
    dep_hierarchy2 = []
    dep_hierarchy2.append([token for token in sent2 if token.dep_ == "ROOT"])
    sent2roots.append([token.text for token in sent2 if token.dep_ == "ROOT"][0])
    sentence_tokens.remove(dep_hierarchy2[0][0])

    while len(sentence_tokens) != 0:
        children = []
        for i in range(len(dep_hierarchy2[-1])):
            index_children = [t for t in sent2 if t.head.text == dep_hierarchy2[-1][i].text and t in sentence_tokens]
            for elem in index_children:
                children.append(elem)
                sentence_tokens.remove(elem)
        dep_hierarchy2.append(children)

    dependency_hierarchy2.append(dep_hierarchy2)


def get_random_string(length):
    letters = string.ascii_lowercase
    result_str = ''.join(random.choice(letters) for i in range(length))
    return result_str


def mask_token(pos_tag):
    """
    For every sentence in the dataframe, mask that token whose POS matches the formal parameter and is also the highest in the dependency tree's structure. 
    If there's more than one option, take that token which comes first in the sentence.
    Return the changed dataframe.
    """
    df = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe
    
    for e in range(len(dependency_hierarchy1)):
        mask_token = None
        for layer in dependency_hierarchy1[e]:
            if mask_token != None:
                break
            for i in range(len(layer)):
                if layer[i].pos_ == pos_tag:
                    mask_token = layer[i].text
                    df.iloc[e, 0] = df.iloc[e, 0].replace(mask_token, get_random_string(random.randrange(1, 9)))
                    break

    for e in range(len(dependency_hierarchy2)):
        mask_token = None
        for layer in dependency_hierarchy2[e]:
            if mask_token != None:
                break
            for i in range(len(layer)):
                if layer[i].pos_ == pos_tag:
                    mask_token = layer[i].text
                    df.iloc[e, 1] = df.iloc[e, 1].replace(mask_token, get_random_string(random.randrange(1, 9)))
                    break
    return df
            

def mask_first(pos_tag):
    """
    For every sentence in the dataframe, mask the first token whose POS matches the formal parameter.
    Return the changed dataframe.
    """
    df = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe

    i = -1

    for sent1 in sentences1_nlp:
        i += 1
        pos_tags_sent = [token.pos_ for token in sent1]
        pos_tag_tokens = [token.text for token in sent1]

        if pos_tag in pos_tags_sent:
            e = pos_tags_sent.index(pos_tag)
            df.iloc[i, 0] = df.iloc[i, 0].replace(pos_tag_tokens[e], get_random_string(random.randrange(1, 9)))
        else:
            continue
    
    i = -1
    for sent2 in sentences2_nlp:
        i += 1
        pos_tags_sent = [token.pos_ for token in sent2]
        pos_tag_tokens = [token.text for token in sent2]

        if pos_tag in pos_tags_sent:
            e = pos_tags_sent.index(pos_tag)
            df.iloc[i, 1] = df.iloc[i, 1].replace(pos_tag_tokens[e], get_random_string(random.randrange(1, 9)))
        else:
            continue
        
    return df


def mask_first_token():
    """
    For every sentence in the dataframe, mask the first token.
    Return the changed dataframe.
    """

    df = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe

    i = -1

    for sent1 in sentences1_nlp:
        i += 1
        first_token = [token.text for token in sent1][0]
        df.iloc[i, 0] = df.iloc[i, 0].replace(first_token, get_random_string(random.randrange(1, 9)))

    
    i = -1
    for sent2 in sentences2_nlp:
        i += 1
        first_token = [token.text for token in sent1][0]
        df.iloc[i, 0] = df.iloc[i, 0].replace(first_token, get_random_string(random.randrange(1, 9)))
        
    return df
    

def mask_root():
    """
    For every sentence in the dataframe, mask the token which is the root in that sentences dependency tree.
    Return the changed dataframe.
    """
    
    df = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe

    for i in range(len(sent1roots)):
        df.iloc[i, 0] = df.iloc[i, 0].replace(sent1roots[i], get_random_string(random.randrange(1, 9)))
        
    for e in range(len(sent2roots)):
        df.iloc[e, 1] = df.iloc[e, 1].replace(sent2roots[2], get_random_string(random.randrange(1, 9)))
        
    return df


