"""
Домашнее задание №3

Реализуйте n-граммную языковую модель, которая сможет отделить
истинные предложения от полученных искусственным способом.

Если вы сомневаетесь, что можно использовать, а что нельзя -- спросите преподавателей.

После того, как побиты бейзлайны, или если подходит к концу срок сдачи,
нужно привести код в порядок и отправить его преподавателям через сайт как файл
(НЕ как ссылку на gist/whatever), убедившись в воспроизводимости результата на лидерборде.

В этом задании можно также уточнить у преподавателей в канале в слаке,
побили ли вы бейзлайны "на привате", если уже побили их "на паблике".

CSC, NLP-2019

"""

from collections import defaultdict

import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer


class LanguageModel(object):
    """
        n-граммная модель.
        Можно расширять и улучшать её любыми трюками, рассказанными на лекции.
    """

    def __init__(self, ngram_size=3):

        if ngram_size < 2:
            raise Exception

        # какого размера нграммы, параметр модели
        self.ngram_size = ngram_size

        self.koefs_ngram = {}
        self.D = 0.75
        self.alpha = 0.045

        # простой способ из коробки построить нграммы
        self.vectorizer = CountVectorizer(ngram_range=(2, ngram_size))
        #print(self.vectorizer)

        # простой способ из коробки построить контексты (то, на что будем делить при оценке вероятности)
        self.context_vectorizer = CountVectorizer(ngram_range=(2 - 1, ngram_size - 1))
        self.ngram_counts = defaultdict(lambda: -np.inf)
        self.words_set_size = None


    def get_koef(self, ngram, curr_word):
        if (self.koefs_ngram.get(ngram)):
            return self.koefs_ngram[ngram];

        ngram_array  = ngram.split(" ")
        if (len(ngram_array)==1):
            return self.sum_context[self.context_vectorizer.vocabulary_[ngram]] / self.words_set_size

        k1 = self.sum_context[self.context_vectorizer.vocabulary_[ngram]]
        k2 = 0

        for ng in self.vectorizer.vocabulary_.keys():
            ng_array = ng.split(" ")

            if (len(ng_array) == len(ngram_array) +1 and ng_array[:-1] == ngram_array):
                k2+=1
        s = " ".join(ngram.split(" ")[1:] + curr_word.split(" "))
        return k2*self.D*self.ngram_counts[s]/k1

    def fit(self, sentences: list):
        """
            Обучение модели на тексте, разбитом на предложения
            :param sentences: список предложений
        """

        print("Fitting sentences")

        # матрица счётчиков нграмм
        counts_matrix = self.vectorizer.fit_transform(sentences)
        #print("counts_matrix",counts_matrix)
        # ...и контекстов
        counts_context_matrix = self.context_vectorizer.fit_transform(sentences)
        #print("counts_context_matrix",counts_context_matrix)

        # сколько всего слов
        self.words_set_size = len(
            set([key for ngram in self.vectorizer.vocabulary_.keys() for key in ngram.split(" ")]))

        print("Summing...")

        # считаем количество нграмм и контекстов -- просто суммируя все столбцы
        # Count(x1,x2,...xn)
        self.sum_ngram = np.sum(counts_matrix, axis=0).A1
        #print(sum_ngram)

        # Count(x1,x2,...x_{n-1})
        self.sum_context = np.sum(counts_context_matrix, axis=0).A1
        #print(sum_context)


        print("shapes: ", self.sum_ngram.shape, self.sum_context.shape)
        print("Iterating through ngrams...")
        count = 0

        for ngram in self.vectorizer.vocabulary_:
            # строим контекст, откидывая последнее слово
            words = " ".join(ngram.split(" ")[:-1])
            #print(ngram)

            # сколько  нграмм = ngram
            # hint: используйте sum_ngram
            #print(ngram, sum_ngram[self.vectorizer.vocabulary_[ngram]])
            # TODO: your code here
            total = self.sum_ngram[self.vectorizer.vocabulary_[ngram]]

            # сколько контекстов = words
            # hint: используйте sum_context
            #print("sum_context",sum_context[words].sum())
            # TODO: your code here
            total_ctx = self.sum_context[self.context_vectorizer.vocabulary_[words]]

            # вычисляем логарифм вероятности нграммы ngram
            # TODO: your code here
            prob = max(total-self.D,0.0)/total_ctx + self.get_koef(words,ngram.split(" ")[-1])
            #print(self.get_koef(words,ngram.split(" ")[-1]))

            # обновляем хранилище
            self.ngram_counts[ngram] = prob
            count += 1
            #print(ngram)
            if (count % 1000 == 0):
                print(count, " of ", self.sum_ngram.shape)

        # логарифмированная вероятность ДЛЯ НЕИЗВЕСТНЫХ НГРАММ
        # TODO: your code here
        self.ngram_counts.default_factory = lambda: 0.000001 #np.log2(0.045 / (0.045*self.words_set_size))

        return self

    def __log_prob_sentence(self, sentence: str):
        """
            Оцениваем логарифм вероятности предложения частотами в рамках данной модели
        :param sentence:
        :return:
        """

        # бьём на токены
        splitted = sentence.split(" ")
        sum_log = 1.0

        for i in range(len(splitted) - self.ngram_size + 1):
            # берём очередную нграмму
            ngram = " ".join(splitted[i: i + self.ngram_size])

            # если не видели, возвращаем "логарифм нуля"
            # иначе складываем логарифмы вероятностей
            sum_log += np.log2(self.ngram_counts[ngram])

        return sum_log

    def log_prob(self, sentence_list: list):
        """
            Логарифм вероятности каждого из предложений в списке
        """
        return list(map(lambda x: self.__log_prob_sentence(x), sentence_list))


df_train = pd.read_csv("train.tsv", sep='\t') #.head(15)
df_test = pd.read_csv("task.tsv", sep='\t')

#print(df_train)

print("Read ", df_train.shape, df_test.shape)

basic_lm = LanguageModel()

sentences_train = df_train["text"].tolist()
basic_lm.fit(sentences=sentences_train)

print("Trained")

test1, test2 = df_test["text1"], df_test["text2"]

logprob1, logprob2 = np.array(basic_lm.log_prob(test1)), np.array(basic_lm.log_prob(test2))

res = pd.DataFrame()
res["id"] = df_test["id"]
res["which"] = 0

res.loc[logprob2 >= logprob1, ["which"]] = 1

res.to_csv("submission.csv", sep=",", index=None, columns=["id", "which"])
print("Done")