import csv
import pandas as pd
import regex as re
import nltk
import tqdm
from nltk import trigrams, word_tokenize
from collections import Counter, defaultdict
import string 

nltk.download("punkt")

most_common_en_word = "the:0.4 be:0.2 to:0.1 of:0.05 and:0.025 a:0.0125 :0.2125"
train_count = 125000
# train set
train_data = pd.read_csv("train/in.tsv.xz", sep="\t", on_bad_lines='skip', header=None, quoting=csv.QUOTE_NONE, nrows=train_count)

# training labels
train_labels = pd.read_csv("train/expected.tsv", sep="\t", on_bad_lines='skip', header=None, quoting=csv.QUOTE_NONE,nrows=train_count)

dev_data = pd.read_csv("dev-0/in.tsv.xz", sep="\t", on_bad_lines='skip', header=None, quoting=csv.QUOTE_NONE)

test_data = pd.read_csv("test-A/in.tsv.xz", sep="\t", on_bad_lines='skip', header=None, quoting=csv.QUOTE_NONE)

def prepare_text(text):
    text = text.lower().replace("-\\n", "").replace("\\n", " ")
    text = re.sub(r"\p{P}", "", text)
    return text

def train_trigrams():
    for _, row in tqdm.tqdm(train_data.iterrows()):
        text = prepare_text(str(row["final"]))
        words = word_tokenize(text)
        for w1, w2, w3 in trigrams(words, pad_right=True, pad_left=True):
            if all([w1, w2, w3]):
                model[(w2, w3)][w1] += 1
                model[(w1, w2)][w3] += 1

    for w_pair in model:
        ngram_count = float(sum(model[w_pair].values()))
        for w3 in model[w_pair]:
            model[w_pair][w3] /= ngram_count


def predict_probs(word1, word2):
    raw_prediction = dict(model[word1, word2])
    prediction = dict(Counter(raw_prediction).most_common(6))

    total_prob = 0.0
    str_prediction = ""

    for word, prob in prediction.items():
        total_prob += prob
        str_prediction += f"{word}:{prob} "

    if total_prob == 0.0:
        return most_common_en_word

    remaining_prob = 1 - total_prob

    if remaining_prob < 0.01:
        remaining_prob = 0.01

    str_prediction += f":{remaining_prob}"

    return str_prediction


def write_output():
    with open("dev-0/out.tsv", "w") as file:
        for _, row in dev_data.iterrows():
            text = prepare_text(str(row[7]))
            words = word_tokenize(text)
            if len(words) < 3:
                prediction = most_common_en_word
            else:
                prediction = predict_probs(words[0], words[1])
            file.write(prediction + "\n")

    with open("test-A/out.tsv", "w") as file:
        for _, row in test_data.iterrows():
            text = prepare_text(str(row[7]))
            words = word_tokenize(text)
            if len(words) < 3:
                prediction = most_common_en_word
            else:
                prediction = predict_probs(words[0], words[1])
            file.write(prediction + "\n")


if __name__ == "__main__":
    # Preapare train data
    print("Preparing data...")
    train_data = train_data[[6, 7]]
    train_data = pd.concat([train_data, train_labels], axis=1)
    train_data["final"] = train_data[6] + train_data[0] + train_data[7]

    # declare model
    print("Preparing model...")
    model = defaultdict(lambda: defaultdict(lambda: 0))

    # train model
    print("Model training...")
    train_trigrams()

    # write outputs
    print("Writing outputs...")
    write_output()