# -*- coding: utf-8 -*-
# Автор: Гусев Илья
# Описание: Модуль марковских цепей.

import os
import pickle
from collections import Counter, defaultdict
from typing import List, Dict, Tuple

import numpy as np

from rupo.files.reader import Reader, FileType
from rupo.generate.language_model.model_container import ModelContainer
from rupo.main.markup import Markup
from rupo.main.vocabulary import StressVocabulary


class MarkovModelContainer(ModelContainer):
    """
    Марковские цепи.
    """
    def __init__(self, dump_filename: str, vocabulary: StressVocabulary, markup_dump_path: str=None,
                 n_poems: int=None, n_grams: int=2):
        self.n_grams = n_grams
        self.transitions = defaultdict(Counter)  # type: Dict[Tuple, Counter]
        self.vocabulary = vocabulary
        self.dump_filename = dump_filename

        # Делаем дамп модели для ускорения загрузки.
        if os.path.isfile(dump_filename):
            self.load()
        else:
            i = 0
            markups = Reader.read_markups(markup_dump_path, FileType.XML, is_processed=True)
            for markup in markups:
                self.add_markup(markup)
                i += 1
                if n_poems is not None and n_poems == i:
                    break
                if i % 500 == 0:
                    print(i)
            self.save()

    def save(self):
        with open(self.dump_filename, "wb") as f:
            pickle.dump(self.transitions, f, pickle.HIGHEST_PROTOCOL)
        self.vocabulary.save()

    def load(self):
        with open(self.dump_filename, "rb") as f:
            self.transitions = pickle.load(f)
        self.vocabulary.load()

    def generate_chain(self, words: List[int]) -> Dict[Tuple, Counter]:
        """
        Генерация переходов в марковских цепях с учётом частотности.

        :param words: вершины цепи.
        :return: обновленные переходы.
        """
        for i in range(len(words) - self.n_grams + 1):
            current_words = tuple([words[j] for j in range(i, i+self.n_grams-1)])
            next_word = words[i+1]
            self.transitions[current_words][next_word] += 1
        return self.transitions

    def add_markup(self, markup: Markup) -> None:
        """
        Дополнение цепей на основе разметки.

        :param markup: разметка.
        """
        words = []
        for line in markup.lines:
            for word in line.words:
                try:
                    words.append(self.vocabulary.get_word_index(word.to_stressed_word()))
                except IndexError:
                    print("Can't find word in vocabulary.")
                    pass

        # Генерируем переходы.
        self.generate_chain(list(reversed(words)))

    def get_model(self, word_indices: List[int]) -> np.array:
        """
        Получение языковой модели.

        :param word_indices: индексы предыдущих слов.
        :return: языковая модель (распределение вероятностей для следующего слова).
        """
        l = self.vocabulary.size()
        if len(word_indices) < self.n_grams - 1:
            return np.full(self.vocabulary.size(), 1/l, dtype=np.float)
        prev_words = tuple([word_indices[-i] for i in range(1, self.n_grams)])
        if len(self.transitions[prev_words]) == 0:
            return np.full(self.vocabulary.size(), 1/l, dtype=np.float)
        else:
            transition = self.transitions[prev_words]
            s = sum(transition.values())
            model = np.zeros(self.vocabulary.size(), dtype=np.float)
            for index, p in transition.items():
                model[index] = p/s
            return model
