import csv
import pandas as pd
import regex as re
import nltk

from nltk import trigrams, word_tokenize
from collections import Counter, defaultdict

# nltk.download("punkt")

# train set
train_data = pd.read_csv(
    "train/in.tsv.xz",
    sep="\t",
    error_bad_lines=False,
    warn_bad_lines=False,
    header=None,
    quoting=csv.QUOTE_NONE,
    nrows=100_000
)

# training labels
train_labels = pd.read_csv(
    "train/expected.tsv",
    sep="\t",
    error_bad_lines=False,
    warn_bad_lines=False,
    header=None,
    quoting=csv.QUOTE_NONE,
    nrows=100_000
)

# dev set
dev_data = pd.read_csv(
    "dev-0/in.tsv.xz",
    sep="\t",
    error_bad_lines=False,
    warn_bad_lines=False,
    header=None,
    quoting=csv.QUOTE_NONE,
)

# test set
test_data = pd.read_csv(
    "test-A/in.tsv.xz",
    sep="\t",
    error_bad_lines=False,
    warn_bad_lines=False,
    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 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 "the:0.2 be:0.2 to:0.2 of:0.1 and:0.1 a:0.1 :0.1"

    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 = "the:0.2 be:0.2 to:0.2 of:0.1 and:0.1 a:0.1 :0.1"
            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 = "the:0.2 be:0.2 to:0.2 of:0.1 and:0.1 a:0.1 :0.1"
            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()