import copy
import itertools

import numpy as np
from sortedcontainers import SortedListWithKey

from .tabled_trie import Trie, make_trie

class LevenshteinSearcher:
    """
    Класс для поиска близких слов
    в соответствии с расстоянием Левенштейна

    """
    def __init__(self, alphabet, dictionary, operation_costs=None,
                 allow_spaces=False, euristics='none'):
        self.alphabet = alphabet
        self.allow_spaces = allow_spaces
        if isinstance(euristics, int):
            if euristics < 0:
                raise ValueError("Euristics should be non-negative integer or None")
            else:
                self.euristics = euristics if euristics != 0 else None
        elif euristics in ["none", "None", None]:
            self.euristics = None
        else:
            raise ValueError("Euristics should be non-negative integer or None")
        if isinstance(dictionary, Trie):
            # словарь передан уже в виде бора
            self.dictionary = dictionary
        else:
            self.dictionary = make_trie(alphabet, dictionary, make_cashed=True,
                                        precompute_symbols=self.euristics,
                                        allow_spaces=self.allow_spaces)
        self.transducer = SegmentTransducer(
            alphabet, operation_costs=operation_costs, allow_spaces=allow_spaces)
        self._precompute_euristics()
        self._define_h_function()

    def __contains__(self, word):
        return word in self.dictionary

    def search(self, word, d, allow_spaces=True, return_cost=True):
        """
        Finds all dictionary words in d-window from word
        """
        if not all((c in self.alphabet
                    or (c == " " and self.allow_spaces)) for c in word):
            return []
            # raise ValueError("{0} contains an incorrect symbol".format(word))
        return self._trie_search(
            word, d, allow_spaces=allow_spaces, return_cost=return_cost)

    def _trie_search(self, word, d, transducer=None,
                     allow_spaces=True, return_cost=True):
        """
        Находит все слова в префиксном боре, расстояние до которых
        в соответствии с заданным преобразователем не превышает d
        """
        if transducer is None:
            # разобраться с пробелами
            transducer = self.transducer.inverse()
        allow_spaces &= self.allow_spaces
        trie = self.dictionary
        #  инициализация переменных
        used_agenda_keys = set()
        agenda = SortedListWithKey(key=(lambda x:x[1]))
        h = self.h_func(word, trie.root)
        # agenda[self.agenda_key("", 0, trie.root)] = (0.0, 0.0, h)
        key, value = ("", 0, trie.root), (0.0, 0.0, h)
        agenda.add((key, value))
        answer = dict()
        k = 0
        # очередь с приоритетом с промежуточными результатами
        while len(agenda) > 0:
            key, value = agenda.pop(0)
            if key in used_agenda_keys:
                continue
            used_agenda_keys.add(key)
            low, pos, index = key
            cost, g, h = value
            # g --- текущая стоимость, h --- нижняя оценка будущей стоимости
            # cost = g + h --- нижняя оценка суммарной стоимости
            k += 1
            max_upperside_length = min(len(word) - pos, transducer.max_up_length)
            for upperside_length in range(max_upperside_length + 1):
                new_pos = pos + upperside_length
                curr_up = word[pos: new_pos]
                if curr_up not in transducer.operation_costs:
                    continue
                for curr_low, curr_cost in transducer.operation_costs[curr_up].items():
                    new_g = g + curr_cost
                    if new_g > d:  #если g > d, то h можно не вычислять
                        continue
                    if curr_low == " ":
                        if allow_spaces and trie.is_final(index):
                            new_index = trie.root
                        else:
                            new_index = Trie.NO_NODE
                    else:
                        new_index = trie.descend(index, curr_low)
                    if new_index is Trie.NO_NODE:
                        continue
                    new_low = low + curr_low
                    new_h = self.h_func(word[new_pos: ], new_index)
                    new_cost = new_g + new_h
                    if new_cost > d:
                        continue
                    new_key = (new_low, new_pos, new_index)
                    new_value = (new_cost, new_g, new_h)
                    if new_pos == len(word) and trie.is_final(new_index):
                        old_g = answer.get(new_low, None)
                        if old_g is None or new_g < old_g:
                            answer[new_low] = new_g
                    agenda.add((new_key, new_value))
        answer = sorted(answer.items(), key=(lambda x: x[1]))
        if return_cost:
            return answer
        else:
            return [elem[0] for elem in answer]

    def _precompute_euristics(self):
        """
        Предвычисляет будущие символы и стоимости операций с ними
        для h-эвристики
        """
        if self.euristics is None:
            return
        # вычисление минимальной стоимости операции,
        # приводящей к появлению ('+') или исчезновению ('-') данного символа
        removal_costs = {a : np.inf for a in self.alphabet}
        insertion_costs = {a : np.inf for a in self.alphabet}
        if self.allow_spaces:
            removal_costs[' '] = np.inf
            insertion_costs[' '] = np.inf
        for up, costs in self.transducer.operation_costs.items():
            for low, cost in costs.items():
                if up == low:
                    continue
                if up != '':
                    removal_cost = cost / len(up)
                    for a in up:
                        removal_costs[a] = min(removal_costs[a], removal_cost)
                if low != '':
                    insertion_cost = cost / len(low)
                    for a in low:
                        insertion_costs[a] = min(insertion_costs[a], insertion_cost)
        # предвычисление возможных будущих символов в узлах дерева
        # precompute_future_symbols(self.dictionary, self.euristics, self.allow_spaces)
        # предвычисление стоимостей потери символа в узлах дерева
        self._absense_costs_by_node = _precompute_absense_costs(
            self.dictionary, removal_costs, insertion_costs,
            self.euristics, self.allow_spaces)
        # массив для сохранения эвристик
        self._temporary_euristics = [dict() for i in range(len(self.dictionary))]

    def _define_h_function(self):
        if self.euristics in [None, 0]:
            self.h_func = (lambda *x: 0.0)
        else:
            self.h_func = self._euristic_h_function

    def _euristic_h_function(self, suffix, index):
        """
        Вычисление h-эвристики из работы Hulden,2009 для текущей вершины словаря

        Аргументы:
        ----------
        suffix : string
            непрочитанный суффикс входного слова
        index : int
            индекс текущего узла в словаре

        Возвращает:
        -----------
        cost : float
            оценка снизу для стоимости замены,
            приводящей к входному слову с суффиксом suffix,
            если прочитанный префикс слова без опечатки
            привёл в вершину с номером index
        """
        if self.euristics > 0:
            suffix = suffix[:self.euristics]
        # кэширование результатов
        index_temporary_euristics = self._temporary_euristics[index]
        cost = index_temporary_euristics.get(suffix, None)
        if cost is not None:
            return cost
        # извлечение нужных данных из массивов
        absense_costs = self._absense_costs_by_node[index]
        data = self.dictionary.data[index]
        costs = np.zeros(dtype=np.float64, shape=(self.euristics,))
        # costs[j] --- оценка штрафа при предпросмотре вперёд на j символов
        for i, a in enumerate(suffix):
            costs[i:] += absense_costs[a][i:]
        cost = max(costs)
        index_temporary_euristics[suffix] = cost
        return cost

    def _minimal_replacement_cost(self, first, second):
        first_symbols, second_symbols = set(), set()
        removal_cost, insertion_cost = 0, 0
        for a, b in itertools.zip_longest(first, second, fillvalue=None):
            if a is not None:
                first_symbols.add(a)
            if b is not None:
                second_symbols.add(b)
            removal_cost = max(removal_cost, len(first_symbols - second_symbols))
            insertion_cost = max(insertion_cost, len(second_symbols - first_symbols))
        return min(removal_cost, insertion_cost)

def _precompute_absense_costs(dictionary, removal_costs, insertion_costs, n,
                              allow_spaces=False):
    """
    Вычисляет минимальную стоимость появления нового символа в узлах словаря
    в соответствии со штрафами из costs

    Аргументы:
    ---------------
    dictionary : Trie
        словарь, хранящийся в виде ациклического автомата

    removal_costs : dict
        штрафы за удаление символов

    insertion_costs : dict
        штрафы за вставку символов

    n : int
        глубина ``заглядывания вперёд'' в словаре

    Возвращает
    ---------------
    answer : list of dicts, len(answer)=len(dictionary)
        answer[i][a][j] равно минимальному штрафу за появление символа a
        в j-ой позиции в вершине с номером i
    """
    answer = [dict() for node in dictionary.data]
    if n == 0:
        return answer
    curr_alphabet = copy.copy(dictionary.alphabet)
    if allow_spaces:
        curr_alphabet += [' ']
    for l, (costs_in_node, node) in enumerate(zip(answer, dictionary.data)):
        # определение минимальной стоимости удаления символов
        curr_node_removal_costs = np.empty(dtype=np.float64, shape=(n,))
        if len(node[0]) > 0:
            curr_node_removal_costs[0] = min(removal_costs[symbol] for symbol in node[0])
            for j, symbols in enumerate(node[1:], 1):
                if len(symbols) == 0:
                    curr_node_removal_costs[j:] = curr_node_removal_costs[j-1]
                    break
                curr_cost = min(removal_costs[symbol] for symbol in symbols)
                curr_node_removal_costs[j] = min(curr_node_removal_costs[j-1], curr_cost)
        else:
            curr_node_removal_costs[:] = np.inf
        # определение минимальной стоимости вставки
        for a in curr_alphabet:
            curr_symbol_costs = np.empty(dtype=np.float64, shape=(n,))
            curr_symbol_costs.fill(insertion_costs[a])
            for j, symbols in enumerate(node):
                if a in symbols:
                    curr_symbol_costs[j:] = 0.0
                    break
                curr_symbol_costs[j] = min(curr_symbol_costs[j], curr_node_removal_costs[j])
            costs_in_node[a] = curr_symbol_costs
    return answer

class SegmentTransducer:
    """
    Класс, реализующий взвешенный конечный преобразователь,
    осуществляющий замены из заданного списка операций

    Аргументы:
    ----------
    alphabet : list
        алфавит

    operation_costs : dict or None(optional, default=None)
        словарь вида {(up,low) : cost}

    allow_spaces : bool(optional, default=False)
        разрешены ли элементы трансдукции, содержащие пробел
        (используется только если явно не заданы operation costs
        и они равны значению по умолчанию)

    """
    def __init__(self, alphabet, operation_costs=None, allow_spaces=False):
        self.alphabet = alphabet
        if operation_costs is None:
            self._make_default_operation_costs(allow_spaces=allow_spaces)
        elif not isinstance(operation_costs, dict):
            raise TypeError("Operation costs must be a dictionary")
        else:
            self.operation_costs = operation_costs
        self._make_reversed_operation_costs()
        self._make_maximal_key_lengths()
        # self.maximal_value_lengths = {}
        # for up, probs in self.operation_costs.items():
            # СЛИШКОМ МНОГО ВЫЗОВОВ, НАДО КАК-ТО ЗАПОМНИТЬ
            # МАКСИМАЛЬНЫЕ ДЛИНЫ КЛЮЧЕЙ ПРИ ОБРАЩЕНИИ
            # max_low_length = max(len(low) for low in probs) if (len(probs) > 0) else -1
            # self.maximal_value_lengths[up] = self.maximal_key_length

    def get_operation_cost(self, up, low):
        """
        Возвращает стоимость элементарной трансдукции up->low
        или np.inf, если такой элементарной трансдукции нет

        Аргументы:
        ----------
        up, low : string
            элементы элементарной трансдукции

        Возвращает:
        -----------
        cost : float
            стоимость элементарной трансдукции up->low
            (np.inf, если такая трансдукция отсутствует)
        """
        up_costs = self.operation_costs.get(up, None)
        if up_costs is None:
            return np.inf
        cost = up_costs.get(low, np.inf)
        return cost

    def inverse(self):
        """
        Строит пробразователь, задающий обратное конечное преобразование
        """
        # УПРОСТИТЬ ОБРАЩЕНИЕ!!!
        inversed_transducer = SegmentTransducer(self.alphabet, operation_costs=dict())
        inversed_transducer.operation_costs = self._reversed_operation_costs
        inversed_transducer._reversed_operation_costs = self.operation_costs
        inversed_transducer.max_low_length = self.max_up_length
        inversed_transducer.max_up_length = self.max_low_length
        inversed_transducer.max_low_lengths_by_up = self.max_up_lengths_by_low
        inversed_transducer.max_up_lengths_by_low = self.max_low_lengths_by_up
        return inversed_transducer

    def distance(self, first, second, return_transduction = False):
        """
        Вычисляет трансдукцию минимальной стоимости,
        отображающую first в second

        Аргументы:
        -----------
        first : string
        second : string
            Верхний и нижний элементы трансдукции

        return_transduction : bool (optional, default=False)
            следует ли возвращать трансдукцию минимального веса
            (см. возвращаемое значение)

        Возвращает:
        -----------
        (final_cost, transductions) : tuple(float, list)
            если return_transduction=True, то возвращает
            минимальную стоимость трансдукции, переводящей first в second
            и список трансдукций с данной стоимостью

        final_cost : float
            если return_transduction=False, то возвращает
            минимальную стоимость трансдукции, переводящей first в second
        """
        if return_transduction:
            add_pred = (lambda x, y: (y == np.inf or x < y))
        else:
            add_pred = (lambda x, y: (y == np.inf or x <= y))
        clear_pred = (lambda x, y: (y < np.inf and x < y))
        update_func = lambda x, y: min(x, y)
        costs, backtraces = self._fill_levenshtein_table(first, second,
                                                        update_func, add_pred, clear_pred)
        final_cost = costs[-1][-1]
        if final_cost == np.inf:
            transductions = [None]
        elif return_transduction:
            transductions = self._backtraces_to_transductions(first, second, backtraces,
                                                              final_cost, return_cost=False)
        if return_transduction:
            return final_cost, transductions
        else:
            return final_cost

    def transduce(self, first, second, threshold):
        """
        Возвращает все трансдукции, переводящие first в second,
        чья стоимость не превышает threshold

        Возвращает:
        ----------
        result : list
            список вида [(трансдукция, стоимость)]
        """
        add_pred = (lambda x, y: x <= threshold)
        clear_pred =(lambda x, y: False)
        update_func = (lambda x, y: min(x, y))
        costs, backtraces = self._fill_levenshtein_table(first, second,
                                                        update_func, add_pred, clear_pred,
                                                        threshold=threshold)
        result = self._backtraces_to_transductions(first, second,
                                                   backtraces, threshold, return_cost=True)
        return result

    def lower_transductions(self, word, max_cost, return_cost=True):
        """
        Возвращает все трансдукции с верхним элементом word,
        чья стоимость не превышает max_cost

    `   Возвращает:
        ----------
        result : list
            список вида [(трансдукция, стоимость)], если return_cost=True
            список трансдукций, если return_cost=False
            список отсортирован в порядке возрастания стоимости трансдукции
        """
        prefixes = [[] for i in range(len(word) + 1)]
        prefixes[0].append(((), 0.0))
        for pos in range(len(prefixes)):
            # вставки
            prefixes[pos] = self._perform_insertions(prefixes[pos], max_cost)
            max_upperside_length = min(len(word) - pos, self.max_up_length)
            for upperside_length in range(1, max_upperside_length + 1):
                up = word[pos: pos + upperside_length]
                for low, low_cost in self.operation_costs.get(up, dict()).items():
                    for transduction, cost in prefixes[pos]:
                        new_cost = cost + low_cost
                        if new_cost <= max_cost:
                            new_transduction = transduction +(up, low)
                            prefixes[pos + upperside_length].append((new_transduction, new_cost))
        answer = sorted(prefixes[-1], key=(lambda x: x[0]))
        if return_cost:
            return answer
        else:
            return [elem[0] for elem in answer]

    def lower(self, word, max_cost, return_cost=True):
        transductions = self.lower_transductions(word, max_cost, return_cost=True)
        answer = dict()
        for transduction, cost in transductions:
            low = "".join(elem[1] for elem in transductions)
            curr_cost = answer.get(low, None)
            if curr_cost is None or cost < curr_cost:
                answer[low] = cost
        answer = sorted(answer.items(), key=(lambda x: x[1]))
        if return_cost:
            return answer
        else:
            return [elem[0] for elem in answer]

    def upper(self, word, max_cost, return_cost=True):
        inversed_transducer = self.inverse()
        return inversed_transducer.lower(word, max_cost, return_cost)

    def upper_transductions(self, word, max_cost, return_cost=True):
        inversed_transducer = self.inverse()
        return inversed_transducer.lower_transductions(word, max_cost, return_cost)

    def _fill_levenshtein_table(self, first, second, update_func, add_pred, clear_pred,
                               threshold=None):
        """
        Функция, динамически заполняющая таблицу costs стоимости трансдукций,
        costs[i][j] --- минимальная стоимость трансдукции,
        переводящей first[:i] в second[:j]

        Аргументы:
        ----------
        first, second : string
            Верхний и нижний элементы трансдукции
        update_func : callable, float*float -> bool
            update_func(x, y) возвращает новое значение в ячейке таблицы costs,
            если старое значение --- y, а потенциально новое значение --- x
            везде update_func = min
        add_pred : callable : float*float -> bool
            add_pred(x, y) возвращает, производится ли добавление
            нового элемента p стоимости x в ячейку backtraces[i][j]
            в зависимости от значения costs[i][j]=y и текущей стоимости x
        clear_pred : callable : float*float -> bool
            clear_pred(x, y) возвращает, производится ли очистка
            ячейки backtraces[i][j] в зависимости от значения costs[i][j]=y
            и текущей стоимости x элемента p, добавляемого в эту ячейку

        Возвращает:
        -----------
        costs : array, dtype=float, shape=(len(first)+1, len(second)+1)
            массив, в ячейке с индексами i, j которого хранится
            минимальная стоимость трансдукции, переводящей first[:i] в second[:j]
        backtraces : array, dtype=list, shape=(len(first)+1, len(second)+1)
            массив, в ячейке с индексами i, j которого хранятся
            обратные ссылки на предыдущую ячейку в оптимальной трансдукции,
            приводящей в ячейку backtraces[i][j]
        """
        m, n = len(first), len(second)
        # если threshold=None, то в качестве порога берётся удвоенная стоимость
        # трансдукции, отображающей символы на одинаковых позициях друг в друга
        if threshold is None:
            threshold = 0.0
            for a, b in zip(first, second):
                threshold += self.get_operation_cost(a, b)
            if m > n:
                for a in first[n: ]:
                    threshold += self.get_operation_cost(a, '')
            elif m < n:
                for b in second[m: ]:
                    threshold += self.get_operation_cost('', b)
            threshold *= 2
        # инициализация возвращаемых массивов
        costs = np.zeros(shape=(m + 1, n + 1), dtype=np.float64)
        costs[:] = np.inf
        backtraces = [None] * (m + 1)
        for i in range(m + 1):
            backtraces[i] = [[] for j in range(n + 1)]
        costs[0][0] = 0.0
        for i in range(m + 1):
            for i_right in range(i, min(i + self.max_up_length, m) + 1):
                up = first[i: i_right]
                max_low_length = self.max_low_lengths_by_up.get(up, -1)
                if max_low_length == -1: # no up key in transduction
                    continue
                up_costs = self.operation_costs[up]
                for j in range(n + 1):
                    if costs[i][j] > threshold:
                        continue
                    if len(backtraces[i][j]) == 0 and i + j > 0:
                        continue # не нашлось обратных ссылок
                    for j_right in range((j if i_right > i else j + 1),
                                         min(j + max_low_length, n) + 1):
                        low = second[j: j_right]
                        curr_cost = up_costs.get(low, np.inf)
                        old_cost = costs[i_right][j_right]
                        new_cost = costs[i][j] + curr_cost
                        if new_cost > threshold:
                            continue
                        if add_pred(new_cost, old_cost):
                            if clear_pred(new_cost, old_cost):
                                backtraces[i_right][j_right] = []
                            costs[i_right][j_right] = update_func(new_cost, old_cost)
                            backtraces[i_right][j_right].append((i, j))
        return costs, backtraces

    def _make_reversed_operation_costs(self):
        """
        Заполняет массив _reversed_operation_costs
        на основе имеющегося массива operation_costs
        """
        _reversed_operation_costs = dict()
        for up, costs in self.operation_costs.items():
            for low, cost in costs.items():
                if low not in _reversed_operation_costs:
                    _reversed_operation_costs[low] = dict()
                _reversed_operation_costs[low][up] = cost
        self._reversed_operation_costs = _reversed_operation_costs

    def _make_maximal_key_lengths(self):
        """
        Вычисляет максимальную длину элемента low
        в элементарной трансдукции (up, low) для каждого up
        и максимальную длину элемента up
        в элементарной трансдукции (up, low) для каждого low
        """
        self.max_up_length =\
            (max(len(up) for up in self.operation_costs)
             if len(self.operation_costs) > 0 else -1)
        self.max_low_length =\
            (max(len(low) for low in self._reversed_operation_costs)
             if len(self._reversed_operation_costs) > 0 else -1)
        self.max_low_lengths_by_up, self.max_up_lengths_by_low = dict(), dict()
        for up, costs in self.operation_costs.items():
            self.max_low_lengths_by_up[up] =\
                max(len(low) for low in costs) if len(costs) > 0 else -1
        for low, costs in self._reversed_operation_costs.items():
            self.max_up_lengths_by_low[low] =\
                max(len(up) for up in costs) if len(costs) > 0 else -1

    def _backtraces_to_transductions(self, first, second, backtraces, threshold, return_cost=False):
        """
        Восстанавливает трансдукции по таблице обратных ссылок

        Аргументы:
        ----------
        first, second : string
            верхние и нижние элементы трансдукции
        backtraces : array-like, dtype=list, shape=(len(first)+1, len(second)+1)
            таблица обратных ссылок
        threshold : float
            порог для отсева трансдукций,
            возвращаются только трансдукции стоимостью <= threshold
        return_cost : bool (optional, default=False)
            если True, то вместе с трансдукциями возвращается их стоимость

        Возвращает:
        -----------
        result : list
            список вида [(трансдукция, стоимость)], если return_cost=True
            и вида [трансдукция], если return_cost=False,
            содержащий все трансдукции, переводящие first в second,
            чья стоимость не превышает threshold
        """
        m, n = len(first), len(second)
        agenda = [None] * (m + 1)
        for i in range(m + 1):
            agenda[i] = [[] for j in range(n+1)]
        agenda[m][n] = [((), 0.0)]
        for i_right in range(m, -1, -1):
            for j_right in range(n, -1, -1):
                current_agenda = agenda[i_right][j_right]
                if len(current_agenda) == 0:
                    continue
                for (i, j) in backtraces[i_right][j_right]:
                    up, low = first[i:i_right], second[j:j_right]
                    add_cost = self.operation_costs[up][low]
                    for elem, cost in current_agenda:
                        new_cost = cost + add_cost
                        if new_cost <= threshold: # удаление трансдукций большой стоимости
                            agenda[i][j].append((((up, low),) + elem, new_cost))
        if return_cost:
            return agenda[0][0]
        else:
            return [elem[0] for elem in agenda[0][0]]

    def _perform_insertions(self, initial, max_cost):
        """
        возвращает все трансдукции стоимости <= max_cost,
        которые можно получить из элементов initial

        Аргументы:
        ----------
        initial : list of tuples
            список исходных трансдукций вида [(трансдукция, стоимость)]
        max_cost : float
            максимальная стоимость трансдукции

        Возвращает:
        -----------
        final : list of tuples
            финальный список трансдукций вида [(трансдукция, стоимость)]
        """
        queue = list(initial)
        final = initial
        while len(queue) > 0:
            transduction, cost = queue[0]
            queue = queue[1:]
            for string, string_cost in self.operation_costs[""].items():
                new_cost = cost + string_cost
                if new_cost <= max_cost:
                    new_transduction = transduction + ("", string)
                    final.append((new_transduction, new_cost))
                    queue.append((new_transduction, new_cost))
        return final

    def _make_default_operation_costs(self, allow_spaces=False):
        """
        sets 1.0 cost for every replacement, insertion, deletion and transposition
        """
        self.operation_costs = dict()
        self.operation_costs[""] = {c: 1.0 for c in list(self.alphabet) + [' ']}
        for a in self.alphabet:
            current_costs = {c: 1.0 for c in self.alphabet}
            current_costs[a] = 0.0
            current_costs[""] = 1.0
            if allow_spaces:
                current_costs[" "] = 1.0
            self.operation_costs[a] = current_costs
        # транспозиции
        for a, b in itertools.permutations(self.alphabet, 2):
            self.operation_costs[a + b] = {b + a: 1.0}
        # пробелы
        if allow_spaces:
            self.operation_costs[" "] = {c: 1.0 for c in self.alphabet}
            self.operation_costs[" "][""] = 1.0


# Copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from copy import deepcopy
from logging import getLogger
from threading import Timer
from typing import Optional

from deeppavlov.agents.default_agent.default_agent import DefaultAgent
from deeppavlov.core.agent.rich_content import RichMessage

log = getLogger(__name__)

class Conversation:
    """Contains agent (if multi-instanced), receives requests, generates responses.

    Args:
        config: Alexa skill configuration settings.
        agent: DeepPavlov Agent instance.
        conversation_key: Alexa conversation ID.
        self_destruct_callback: Conversation instance deletion callback function.

    Attributes:
        config: Alexa skill configuration settings.
        agent: Alexa skill agent.
        key: Alexa conversation ID.
        stateful: Stateful mode flag.
        timer: Conversation self-destruct timer.
        handled_requests: Mapping of Alexa requests types to requests handlers.
        response_template: Alexa response template.
        """
    def __init__(self, config: dict, agent: DefaultAgent, conversation_key: str,
                 self_destruct_callback: callable) -> None:
        self.config = config
        self.agent = agent
        self.key = conversation_key
        self.self_destruct_callback = self_destruct_callback
        self.stateful: bool = self.config['stateful']
        self.timer: Optional[Timer] = None

        self.handled_requests = {
            'LaunchRequest': self._handle_launch,
            'IntentRequest': self._handle_intent,
            'SessionEndedRequest': self._handle_end,
            '_unsupported': self._handle_unsupported
        }

        self.response_template = {
            'version': '1.0',
            'sessionAttributes': {
                'sessionId': None
            }
        }

        self._start_timer()

    def _start_timer(self) -> None:
        """Initiates self-destruct timer."""
        self.timer = Timer(self.config['conversation_lifetime'], self.self_destruct_callback)
        self.timer.start()

    def _rearm_self_destruct(self) -> None:
        """Rearms self-destruct timer."""
        self.timer.cancel()
        self._start_timer()

    def handle_request(self, request: dict) -> dict:
        """Routes Alexa requests to appropriate handlers.

        Args:
            request: Alexa request.
        Returns:
            response: Response conforming Alexa response specification.
        """
        request_type = request['request']['type']
        request_id = request['request']['requestId']
        log.debug(f'Received request. Type: {request_type}, id: {request_id}')

        if request_type in self.handled_requests.keys():
            response: dict = self.handled_requests[request_type](request)
        else:
            response: dict = self.handled_requests['_unsupported'](request)
            log.warning(f'Unsupported request type: {request_type}, request id: {request_id}')

        self._rearm_self_destruct()

        return response

    def _act(self, utterance: str) -> list:
        """Infers DeepPavlov agent with raw user input extracted from Alexa request.

        Args:
            utterance: Raw user input extracted from Alexa request.
        Returns:
            response: DeepPavlov agent response.
        """
        if self.stateful:
            utterance = [[utterance], [self.key]]
        else:
            utterance = [[utterance]]

        agent_response: list = self.agent(*utterance)

        return agent_response

    def _generate_response(self, response: dict, request: dict) -> dict:
        """Populates generated response with additional data conforming Alexa response specification.

        Args:
            response: Raw user input extracted from Alexa request.
            request: Alexa request.
        Returns:
            response: Response conforming Alexa response specification.
        """
        response_template = deepcopy(self.response_template)
        response_template['sessionAttributes']['sessionId'] = request['session']['sessionId']

        for key, value in response_template.items():
            if key not in response.keys():
                response[key] = value

        return response

    def _handle_intent(self, request: dict) -> dict:
        """Handles IntentRequest Alexa request.

        Args:
            request: Alexa request.
        Returns:
            response: "response" part of response dict conforming Alexa specification.
        """
        intent_name = self.config['intent_name']
        slot_name = self.config['slot_name']

        request_id = request['request']['requestId']
        request_intent: dict = request['request']['intent']

        if intent_name != request_intent['name']:
            log.error(f"Wrong intent name received: {request_intent['name']} in request {request_id}")
            return {'error': 'wrong intent name'}

        if slot_name not in request_intent['slots'].keys():
            log.error(f'No slot named {slot_name} found in request {request_id}')
            return {'error': 'no slot found'}

        utterance = request_intent['slots'][slot_name]['value']
        agent_response = self._act(utterance)

        if not agent_response:
            log.error(f'Some error during response generation for request {request_id}')
            return {'error': 'error during response generation'}

        prediction: RichMessage = agent_response[0]
        prediction: list = prediction.alexa()

        if not prediction:
            log.error(f'Some error during response generation for request {request_id}')
            return {'error': 'error during response generation'}

        response = self._generate_response(prediction[0], request)

        return response

    def _handle_launch(self, request: dict) -> dict:
        """Handles LaunchRequest Alexa request.

        Args:
            request: Alexa request.
        Returns:
            response: "response" part of response dict conforming Alexa specification.
        """
        response = {
            'response': {
                'shouldEndSession': False,
                'outputSpeech': {
                    'type': 'PlainText',
                    'text': self.config['start_message']
                },
                'card': {
                    'type': 'Simple',
                    'content': self.config['start_message']
                }
            }
        }

        response = self._generate_response(response, request)

        return response

    def _handle_end(self, request: dict) -> dict:
        """Handles SessionEndedRequest Alexa request and deletes Conversation instance.

        Args:
            request: Alexa request.
        Returns:
            response: Dummy empty response dict.
        """
        response = {}
        self.self_destruct_callback()
        return response

    def _handle_unsupported(self, request: dict) -> dict:
        """Handles all unsupported types of Alexa requests. Returns standard message.

        Args:
            request: Alexa request.
        Returns:
            response: "response" part of response dict conforming Alexa specification.
        """
        response = {
            'response': {
                'shouldEndSession': False,
                'outputSpeech': {
                    'type': 'PlainText',
                    'text': self.config['unsupported_message']
                },
                'card': {
                    'type': 'Simple',
                    'content': self.config['unsupported_message']
                }
            }
        }

        response = self._generate_response(response, request)

        return response


from pathlib import Path
from typing import Iterator, Dict, Union, Iterable

class Struct:
    def __iter__(self) -> Iterator[str]:
        return iter(self._keys)

    def __len__(self) -> int:
        return len(self._keys)

    def __init__(self, tree: Dict[str, Union[dict, Path]]) -> None:
        self._keys = set()
        for key, value in tree.items():
            key = key.replace('.', '_')
            self._keys.add(key)
            setattr(self, key,
                    Struct(value) if isinstance(value, dict) else value)
        self._keys = frozenset(self._keys)

        self.keys = lambda: self._keys

    def _asdict(self, *, to_string: bool=False) -> dict:
        res = []
        for key in self._keys:
            value = getattr(self, key)
            if isinstance(value, Struct):
                value = value._asdict(to_string=to_string)
            elif to_string:
                value = str(value)
            res.append((key, value))

        return dict(res)

    def __getitem__(self, key: str) -> Union[dict, Path]:
        if key not in self._keys:
            raise KeyError(key)

        item = getattr(self, key)
        if isinstance(item, Struct):
            item = item._asdict()
        return item

    def __dir__(self) -> Iterable:
        return self._keys

    def _ipython_key_completions_(self) -> Iterable:
        return self._keys

    def __str__(self) -> str:
        return str(self._asdict(to_string=True))

    def __repr__(self) -> str:
        return f'Struct({repr(self._asdict())})'

    def _repr_pretty_(self, p, cycle):
        """method that defines ``Struct``'s pretty printing rules for iPython

        Args:
            p (IPython.lib.pretty.RepresentationPrinter): pretty printer object
            cycle (bool): is ``True`` if pretty detected a cycle
        """
        if cycle:
            p.text('Struct(...)')
        else:
            with p.group(7, 'Struct(', ')'):
                p.pretty(self._asdict())

def _build_configs_tree() -> Struct:
    root = Path(__file__).resolve().parent

    tree = {}

    for config in root.glob('**/*.json'):
        leaf = tree
        for part in config.relative_to(root).parent.parts:
            if part not in leaf:
                leaf[part] = {}
            leaf = leaf[part]
        leaf[config.stem] = config

    return Struct(tree)

configs = _build_configs_tree()


# Copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List

import numpy as np

from deeppavlov.core.common.metrics_registry import register_metric

@register_metric('elmo_loss2ppl')
def elmo_loss2ppl(losses: List[np.ndarray]) -> float:
    """ Calculates perplexity by loss

    Args:
        losses: list of numpy arrays of model losses

    Returns:
        perplexity : float
    """
    avg_loss = np.mean(losses)
    return float(np.exp(avg_loss))


# originally based on https://github.com/allenai/bilm-tf/blob/master/bilm/training.py

# Modifications copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import tensorflow as tf

DTYPE = 'float32'
DTYPE_INT = 'int64'

tf.logging.set_verbosity(tf.logging.INFO)

class LanguageModel(object):
    """
    A class to build the tensorflow computational graph for NLMs

    All hyperparameters and model configuration is specified in a dictionary
    of 'options'.

    is_training is a boolean used to control behavior of dropout layers
        and softmax.  Set to False for testing.

    The LSTM cell is controlled by the 'lstm' key in options
    Here is an example:

     'lstm': {
      'cell_clip': 5,
      'dim': 4096,
      'n_layers': 2,
      'proj_clip': 5,
      'projection_dim': 512,
      'use_skip_connections': True},

        'projection_dim' is assumed token embedding size and LSTM output size.
        'dim' is the hidden state size.
        Set 'dim' == 'projection_dim' to skip a projection layer.
    """
    def __init__(self, options, is_training):
        self.options = options
        self.is_training = is_training
        self.bidirectional = options.get('bidirectional', False)

        # use word or char inputs?
        self.char_inputs = 'char_cnn' in self.options

        # for the loss function
        self.share_embedding_softmax = options.get(
            'share_embedding_softmax', False)
        if self.char_inputs and self.share_embedding_softmax:
            raise ValueError("Sharing softmax and embedding weights requires "
                             "word input")

        self.sample_softmax = options.get('sample_softmax', True)

        self._build()

    def _build_word_embeddings(self):
        n_tokens_vocab = self.options['n_tokens_vocab']
        batch_size = self.options['batch_size']
        unroll_steps = self.options['unroll_steps']

        # LSTM options
        projection_dim = self.options['lstm']['projection_dim']

        # the input token_ids and word embeddings
        self.token_ids = tf.placeholder(DTYPE_INT,
                                        shape=(batch_size, unroll_steps),
                                        name='token_ids')
        # the word embeddings
        with tf.device("/cpu:0"):
            self.embedding_weights = tf.get_variable(
                "embedding", [n_tokens_vocab, projection_dim],
                dtype=DTYPE,
            )
            self.embedding = tf.nn.embedding_lookup(self.embedding_weights,
                                                    self.token_ids)

        # if a bidirectional LM then make placeholders for reverse
        # model and embeddings
        if self.bidirectional:
            self.token_ids_reverse = tf.placeholder(DTYPE_INT,
                                                    shape=(batch_size, unroll_steps),
                                                    name='token_ids_reverse')
            with tf.device("/cpu:0"):
                self.embedding_reverse = tf.nn.embedding_lookup(
                    self.embedding_weights, self.token_ids_reverse)

    def _build_word_char_embeddings(self):
        """
        options contains key 'char_cnn': {

        'n_characters': 262,

        # includes the start / end characters
        'max_characters_per_token': 50,

        'filters': [
            [1, 32],
            [2, 32],
            [3, 64],
            [4, 128],
            [5, 256],
            [6, 512],
            [7, 512]
        ],
        'activation': 'tanh',

        # for the character embedding
        'embedding': {'dim': 16}

        # for highway layers
        # if omitted, then no highway layers
        'n_highway': 2,
        }
        """
        batch_size = self.options['batch_size']
        unroll_steps = self.options['unroll_steps']
        projection_dim = self.options['lstm']['projection_dim']

        cnn_options = self.options['char_cnn']
        filters = cnn_options['filters']
        n_filters = sum(f[1] for f in filters)
        max_chars = cnn_options['max_characters_per_token']
        char_embed_dim = cnn_options['embedding']['dim']
        n_chars = cnn_options['n_characters']
        if n_chars != 261:
            raise Exception("Set n_characters=261 for training see a \
                            https://github.com/allenai/bilm-tf/blob/master/README.md")
        if cnn_options['activation'] == 'tanh':
            activation = tf.nn.tanh
        elif cnn_options['activation'] == 'relu':
            activation = tf.nn.relu

        # the input character ids
        self.tokens_characters = tf.placeholder(DTYPE_INT,
                                                shape=(batch_size, unroll_steps, max_chars),
                                                name='tokens_characters')
        # the character embeddings
        with tf.device("/cpu:0"):
            self.embedding_weights = tf.get_variable("char_embed", [n_chars, char_embed_dim],
                                                     dtype=DTYPE,
                                                     initializer=tf.random_uniform_initializer(-1.0, 1.0))
            # shape (batch_size, unroll_steps, max_chars, embed_dim)
            self.char_embedding = tf.nn.embedding_lookup(self.embedding_weights,
                                                         self.tokens_characters)

            if self.bidirectional:
                self.tokens_characters_reverse = tf.placeholder(DTYPE_INT,
                                                                shape=(batch_size, unroll_steps, max_chars),
                                                                name='tokens_characters_reverse')
                self.char_embedding_reverse = tf.nn.embedding_lookup(
                    self.embedding_weights, self.tokens_characters_reverse)

        # the convolutions
        def make_convolutions(inp, reuse):
            with tf.variable_scope('CNN', reuse=reuse):
                convolutions = []
                for i, (width, num) in enumerate(filters):
                    if cnn_options['activation'] == 'relu':
                        # He initialization for ReLU activation
                        # with char embeddings init between -1 and 1
                        # w_init = tf.random_normal_initializer(
                        #    mean=0.0,
                        #    stddev=np.sqrt(2.0 / (width * char_embed_dim))
                        # )

                        # Kim et al 2015, +/- 0.05
                        w_init = tf.random_uniform_initializer(
                            minval=-0.05, maxval=0.05)
                    elif cnn_options['activation'] == 'tanh':
                        # glorot init
                        w_init = tf.random_normal_initializer(
                            mean=0.0,
                            stddev=np.sqrt(1.0 / (width * char_embed_dim))
                        )
                    w = tf.get_variable(
                        "W_cnn_%s" % i,
                        [1, width, char_embed_dim, num],
                        initializer=w_init,
                        dtype=DTYPE)
                    b = tf.get_variable(
                        "b_cnn_%s" % i, [num], dtype=DTYPE,
                        initializer=tf.constant_initializer(0.0))

                    conv = tf.nn.conv2d(inp, w,
                                        strides=[1, 1, 1, 1],
                                        padding="VALID") + b
                    # now max pool
                    conv = tf.nn.max_pool(conv, [1, 1, max_chars - width + 1, 1],
                                          [1, 1, 1, 1], 'VALID')

                    # activation
                    conv = activation(conv)
                    conv = tf.squeeze(conv, squeeze_dims=[2])

                    convolutions.append(conv)

            return tf.concat(convolutions, 2)

        # for first model, this is False, for others it's True
        reuse = tf.get_variable_scope().reuse
        embedding = make_convolutions(self.char_embedding, reuse)

        self.token_embedding_layers = [embedding]

        if self.bidirectional:
            # re-use the CNN weights from forward pass
            embedding_reverse = make_convolutions(
                self.char_embedding_reverse, True)

        # for highway and projection layers:
        #   reshape from (batch_size, n_tokens, dim) to
        n_highway = cnn_options.get('n_highway')
        use_highway = n_highway is not None and n_highway > 0
        use_proj = n_filters != projection_dim

        if use_highway or use_proj:
            embedding = tf.reshape(embedding, [-1, n_filters])
            if self.bidirectional:
                embedding_reverse = tf.reshape(embedding_reverse,
                                               [-1, n_filters])

        # set up weights for projection
        if use_proj:
            assert n_filters > projection_dim
            with tf.variable_scope('CNN_proj'):
                W_proj_cnn = tf.get_variable(
                    "W_proj", [n_filters, projection_dim],
                    initializer=tf.random_normal_initializer(
                        mean=0.0, stddev=np.sqrt(1.0 / n_filters)),
                    dtype=DTYPE)
                b_proj_cnn = tf.get_variable(
                    "b_proj", [projection_dim],
                    initializer=tf.constant_initializer(0.0),
                    dtype=DTYPE)

        # apply highways layers
        def high(x, ww_carry, bb_carry, ww_tr, bb_tr):
            carry_gate = tf.nn.sigmoid(tf.matmul(x, ww_carry) + bb_carry)
            transform_gate = tf.nn.relu(tf.matmul(x, ww_tr) + bb_tr)
            return carry_gate * transform_gate + (1.0 - carry_gate) * x

        if use_highway:
            highway_dim = n_filters

            for i in range(n_highway):
                with tf.variable_scope('CNN_high_%s' % i):
                    W_carry = tf.get_variable(
                        'W_carry', [highway_dim, highway_dim],
                        # glorit init
                        initializer=tf.random_normal_initializer(
                            mean=0.0, stddev=np.sqrt(1.0 / highway_dim)),
                        dtype=DTYPE)
                    b_carry = tf.get_variable(
                        'b_carry', [highway_dim],
                        initializer=tf.constant_initializer(-2.0),
                        dtype=DTYPE)
                    W_transform = tf.get_variable(
                        'W_transform', [highway_dim, highway_dim],
                        initializer=tf.random_normal_initializer(
                            mean=0.0, stddev=np.sqrt(1.0 / highway_dim)),
                        dtype=DTYPE)
                    b_transform = tf.get_variable(
                        'b_transform', [highway_dim],
                        initializer=tf.constant_initializer(0.0),
                        dtype=DTYPE)

                embedding = high(embedding, W_carry, b_carry,
                                 W_transform, b_transform)
                if self.bidirectional:
                    embedding_reverse = high(embedding_reverse,
                                             W_carry, b_carry,
                                             W_transform, b_transform)
                self.token_embedding_layers.append(tf.reshape(embedding,
                                                   [batch_size, unroll_steps, highway_dim]))

        # finally project down to projection dim if needed
        if use_proj:
            embedding = tf.matmul(embedding, W_proj_cnn) + b_proj_cnn
            if self.bidirectional:
                embedding_reverse = tf.matmul(embedding_reverse, W_proj_cnn) \
                    + b_proj_cnn
            self.token_embedding_layers.append(
                tf.reshape(embedding, [batch_size, unroll_steps, projection_dim])
            )

        # reshape back to (batch_size, tokens, dim)
        if use_highway or use_proj:
            shp = [batch_size, unroll_steps, projection_dim]
            embedding = tf.reshape(embedding, shp)
            if self.bidirectional:
                embedding_reverse = tf.reshape(embedding_reverse, shp)

        # at last assign attributes for remainder of the model
        self.embedding = embedding
        if self.bidirectional:
            self.embedding_reverse = embedding_reverse

    def _build(self):
        # size of input options
        batch_size = self.options['batch_size']

        # LSTM options
        lstm_dim = self.options['lstm']['dim']
        projection_dim = self.options['lstm']['projection_dim']
        n_lstm_layers = self.options['lstm'].get('n_layers', 1)
        dropout = self.options['dropout']
        keep_prob = 1.0 - dropout

        if self.char_inputs:
            self._build_word_char_embeddings()
        else:
            self._build_word_embeddings()

        # now the LSTMs
        # these will collect the initial states for the forward
        #   (and reverse LSTMs if we are doing bidirectional)
        self.init_lstm_state = []
        self.final_lstm_state = []

        # get the LSTM inputs
        if self.bidirectional:
            lstm_inputs = [self.embedding, self.embedding_reverse]
        else:
            lstm_inputs = [self.embedding]

        # now compute the LSTM outputs
        cell_clip = self.options['lstm'].get('cell_clip')
        proj_clip = self.options['lstm'].get('proj_clip')

        use_skip_connections = self.options['lstm'].get('use_skip_connections')

        lstm_outputs = []
        for lstm_num, lstm_input in enumerate(lstm_inputs):
            lstm_cells = []
            for i in range(n_lstm_layers):
                if projection_dim < lstm_dim:
                    # are projecting down output
                    lstm_cell = tf.nn.rnn_cell.LSTMCell(
                        lstm_dim, num_proj=projection_dim,
                        cell_clip=cell_clip, proj_clip=proj_clip)
                else:
                    lstm_cell = tf.nn.rnn_cell.LSTMCell(
                        lstm_dim,
                        cell_clip=cell_clip, proj_clip=proj_clip)

                if use_skip_connections:
                    # ResidualWrapper adds inputs to outputs
                    if i == 0:
                        # don't add skip connection from token embedding to
                        # 1st layer output
                        pass
                    else:
                        # add a skip connection
                        lstm_cell = tf.nn.rnn_cell.ResidualWrapper(lstm_cell)

                # add dropout
                if self.is_training:
                    lstm_cell = tf.nn.rnn_cell.DropoutWrapper(lstm_cell,
                                                              input_keep_prob=keep_prob)

                lstm_cells.append(lstm_cell)

            if n_lstm_layers > 1:
                lstm_cell = tf.nn.rnn_cell.MultiRNNCell(lstm_cells)
            else:
                lstm_cell = lstm_cells[0]

            with tf.control_dependencies([lstm_input]):
                self.init_lstm_state.append(
                    lstm_cell.zero_state(batch_size, DTYPE))
                # NOTE: this variable scope is for backward compatibility
                # with existing models...
                if self.bidirectional:
                    with tf.variable_scope('RNN_%s' % lstm_num):
                        _lstm_output_unpacked, final_state = tf.nn.static_rnn(
                            lstm_cell,
                            tf.unstack(lstm_input, axis=1),
                            initial_state=self.init_lstm_state[-1])
                else:
                    _lstm_output_unpacked, final_state = tf.nn.static_rnn(
                        lstm_cell,
                        tf.unstack(lstm_input, axis=1),
                        initial_state=self.init_lstm_state[-1])
                self.final_lstm_state.append(final_state)

            # (batch_size * unroll_steps, 512)
            lstm_output_flat = tf.reshape(
                tf.stack(_lstm_output_unpacked, axis=1), [-1, projection_dim])
            if self.is_training:
                # add dropout to output
                lstm_output_flat = tf.nn.dropout(lstm_output_flat, keep_prob)
            tf.add_to_collection('lstm_output_embeddings', _lstm_output_unpacked)

            lstm_outputs.append(lstm_output_flat)

        self._build_loss(lstm_outputs)

    def _build_loss(self, lstm_outputs):
        """
        Create:
            self.total_loss: total loss op for training
            self.softmax_W, softmax_b: the softmax variables
            self.next_token_id / _reverse: placeholders for gold input

        """
        batch_size = self.options['batch_size']
        unroll_steps = self.options['unroll_steps']

        n_tokens_vocab = self.options['n_tokens_vocab']

        # DEFINE next_token_id and *_reverse placeholders for the gold input
        def _get_next_token_placeholders(suffix):
            name = 'next_token_id' + suffix
            id_placeholder = tf.placeholder(DTYPE_INT,
                                            shape=(batch_size, unroll_steps),
                                            name=name)
            return id_placeholder

        # get the window and weight placeholders
        self.next_token_id = _get_next_token_placeholders('')
        if self.bidirectional:
            self.next_token_id_reverse = _get_next_token_placeholders(
                '_reverse')

        # DEFINE THE SOFTMAX VARIABLES
        # get the dimension of the softmax weights
        # softmax dimension is the size of the output projection_dim
        softmax_dim = self.options['lstm']['projection_dim']

        # the output softmax variables -- they are shared if bidirectional
        if self.share_embedding_softmax:
            # softmax_W is just the embedding layer
            self.softmax_W = self.embedding_weights

        with tf.variable_scope('softmax'), tf.device('/cpu:0'):
            # Glorit init (std=(1.0 / sqrt(fan_in))
            softmax_init = tf.random_normal_initializer(0.0, 1.0 / np.sqrt(softmax_dim))
            if not self.share_embedding_softmax:
                self.softmax_W = tf.get_variable(
                    'W', [n_tokens_vocab, softmax_dim],
                    dtype=DTYPE,
                    initializer=softmax_init
                )
            self.softmax_b = tf.get_variable(
                'b', [n_tokens_vocab],
                dtype=DTYPE,
                initializer=tf.constant_initializer(0.0))

        # now calculate losses
        # loss for each direction of the LSTM
        self.individual_train_losses = []
        self.individual_eval_losses = []

        if self.bidirectional:
            next_ids = [self.next_token_id, self.next_token_id_reverse]
        else:
            next_ids = [self.next_token_id]

        for id_placeholder, lstm_output_flat in zip(next_ids, lstm_outputs):
            # flatten the LSTM output and next token id gold to shape:
            # (batch_size * unroll_steps, softmax_dim)
            # Flatten and reshape the token_id placeholders
            next_token_id_flat = tf.reshape(id_placeholder, [-1, 1])

            with tf.control_dependencies([lstm_output_flat]):
                sampled_losses = tf.nn.sampled_softmax_loss(self.softmax_W, self.softmax_b,
                                                            next_token_id_flat, lstm_output_flat,
                                                            self.options['n_negative_samples_batch'],
                                                            self.options['n_tokens_vocab'],
                                                            num_true=1)

                # get the full softmax loss
                output_scores = tf.matmul(
                    lstm_output_flat,
                    tf.transpose(self.softmax_W)
                ) + self.softmax_b
                # NOTE: tf.nn.sparse_softmax_cross_entropy_with_logits
                #   expects unnormalized output since it performs the
                #   softmax internally
                losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
                    logits=output_scores,
                    labels=tf.squeeze(next_token_id_flat, squeeze_dims=[1])
                )
            sampled_losses = tf.reshape(sampled_losses, [self.options['batch_size'], -1])
            losses = tf.reshape(losses, [self.options['batch_size'], -1])
            self.individual_train_losses.append(tf.reduce_mean(sampled_losses, axis=1))
            self.individual_eval_losses.append(tf.reduce_mean(losses, axis=1))

        # now make the total loss -- it's the train of the individual losses
        if self.bidirectional:
            self.total_train_loss = 0.5 * (self.individual_train_losses[0] + self.individual_train_losses[1])
            self.total_eval_loss = 0.5 * (self.individual_eval_losses[0] + self.individual_eval_losses[1])
        else:
            self.total_train_loss = self.individual_train_losses[0]
            self.total_eval_loss = self.individual_eval_losses[0]


# Copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import pickle
import sys
from itertools import islice
from logging import getLogger
from pathlib import Path
from typing import Optional, Union

from deeppavlov.core.commands.utils import import_packages, parse_config
from deeppavlov.core.common.chainer import Chainer
from deeppavlov.core.common.params import from_params
from deeppavlov.download import deep_download

log = getLogger(__name__)

def build_model(config: Union[str, Path, dict], mode: str = 'infer',
                load_trained: bool = False, download: bool = False,
                serialized: Optional[bytes] = None) -> Chainer:
    """Build and return the model described in corresponding configuration file."""
    config = parse_config(config)

    if serialized:
        serialized: list = pickle.loads(serialized)

    if download:
        deep_download(config)

    import_packages(config.get('metadata', {}).get('imports', []))

    model_config = config['chainer']

    model = Chainer(model_config['in'], model_config['out'], model_config.get('in_y'))

    for component_config in model_config['pipe']:
        if load_trained and ('fit_on' in component_config or 'in_y' in component_config):
            try:
                component_config['load_path'] = component_config['save_path']
            except KeyError:
                log.warning('No "save_path" parameter for the {} component, so "load_path" will not be renewed'
                            .format(component_config.get('class_name', component_config.get('ref', 'UNKNOWN'))))

        if serialized and 'in' in component_config:
            component_serialized = serialized.pop(0)
        else:
            component_serialized = None

        component = from_params(component_config, mode=mode, serialized=component_serialized)

        if 'in' in component_config:
            c_in = component_config['in']
            c_out = component_config['out']
            in_y = component_config.get('in_y', None)
            main = component_config.get('main', False)
            model.append(component, c_in, c_out, in_y, main)

    return model

def interact_model(config: Union[str, Path, dict]) -> None:
    """Start interaction with the model described in corresponding configuration file."""
    model = build_model(config)

    while True:
        args = []
        for in_x in model.in_x:
            args.append((input('{}::'.format(in_x)),))
            # check for exit command
            if args[-1][0] in {'exit', 'stop', 'quit', 'q'}:
                return

        pred = model(*args)
        if len(model.out_params) > 1:
            pred = zip(*pred)

        print('>>', *pred)

def predict_on_stream(config: Union[str, Path, dict], batch_size: int = 1, file_path: Optional[str] = None) -> None:
    """Make a prediction with the component described in corresponding configuration file."""
    if file_path is None or file_path == '-':
        if sys.stdin.isatty():
            raise RuntimeError('To process data from terminal please use interact mode')
        f = sys.stdin
    else:
        f = open(file_path, encoding='utf8')

    model: Chainer = build_model(config)

    args_count = len(model.in_x)
    while True:
        batch = list((l.strip() for l in islice(f, batch_size * args_count)))

        if not batch:
            break

        args = []
        for i in range(args_count):
            args.append(batch[i::args_count])

        res = model(*args)
        if len(model.out_params) == 1:
            res = [res]
        for res in zip(*res):
            res = json.dumps(res, ensure_ascii=False)
            print(res, flush=True)

    if f is not sys.stdin:
        f.close()


# Copyright 2018 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from logging import getLogger
from pathlib import Path
from typing import Dict, List, Union, Tuple, Optional

from deeppavlov.core.common.registry import register
from deeppavlov.core.data.dataset_reader import DatasetReader
from deeppavlov.core.data.utils import download_decompress, mark_done

WORD_COLUMN, POS_COLUMN, TAG_COLUMN = 1, 3, 5

log = getLogger(__name__)

def get_language(filepath: str) -> str:
    """Extracts language from typical UD filename
    """
    return filepath.split("-")[0]

def read_infile(infile: Union[Path, str], from_words=False,
                word_column: int = WORD_COLUMN, pos_column: int = POS_COLUMN,
                tag_column: int = TAG_COLUMN, max_sents: int = -1,
                read_only_words: bool = False) -> List[Tuple[List, Union[List, None]]]:
    """Reads input file in CONLL-U format

    Args:
        infile: a path to a file
        word_column: column containing words (default=1)
        pos_column: column containing part-of-speech labels (default=3)
        tag_column: column containing fine-grained tags (default=5)
        max_sents: maximal number of sents to read
        read_only_words: whether to read only words

    Returns:
        a list of sentences. Each item contains a word sequence and a tag sequence, which is ``None``
        in case ``read_only_words = True``
    """
    answer, curr_word_sent, curr_tag_sent = [], [], []
    if from_words:
        word_column, read_only_words = 0, True
    with open(infile, "r", encoding="utf8") as fin:
        for line in fin:
            line = line.strip()
            if line.startswith("#"):
                continue
            if line == "":
                if len(curr_word_sent) > 0:
                    if read_only_words:
                        curr_tag_sent = None
                    answer.append((curr_word_sent, curr_tag_sent))
                curr_tag_sent, curr_word_sent = [], []
                if len(answer) == max_sents:
                    break
                continue
            splitted = line.split("\t")
            index = splitted[0]
            if not from_words and not index.isdigit():
                continue
            curr_word_sent.append(splitted[word_column])
            if not read_only_words:
                pos, tag = splitted[pos_column], splitted[tag_column]
                tag = pos if tag == "_" else "{},{}".format(pos, tag)
                curr_tag_sent.append(tag)
        if len(curr_word_sent) > 0:
            if read_only_words:
                curr_tag_sent = None
            answer.append((curr_word_sent, curr_tag_sent))
    return answer

@register('morphotagger_dataset_reader')
class MorphotaggerDatasetReader(DatasetReader):
    """Class to read training datasets in UD format"""

    URL = 'http://files.deeppavlov.ai/datasets/UD2.0_source/'

    def read(self, data_path: Union[List, str],
             language: Optional[None] = None,
             data_types: Optional[List[str]] = None,
             **kwargs) -> Dict[str, List]:
        """Reads UD dataset from data_path.

        Args:
            data_path: can be either
                1. a directory containing files. The file for data_type 'mode'
                is then data_path / {language}-ud-{mode}.conllu
                2. a list of files, containing the same number of items as data_types
            language: a language to detect filename when it is not given
            data_types: which dataset parts among 'train', 'dev', 'test' are returned

        Returns:
            a dictionary containing dataset fragments (see ``read_infile``) for given data types
        """
        if data_types is None:
            data_types = ["train", "dev"]
        elif isinstance(data_types, str):
            data_types = list(data_types)
        for data_type in data_types:
            if data_type not in ["train", "dev", "test"]:
                raise ValueError("Unknown data_type: {}, only train, dev and test "
                                 "datatypes are allowed".format(data_type))
        if isinstance(data_path, str):
            data_path = Path(data_path)
        if isinstance(data_path, Path):
            if data_path.exists():
                is_file = data_path.is_file()
            else:
                is_file = (len(data_types) == 1)
            if is_file:
                # path to a single file
                data_path, reserve_data_path = [data_path], None
            else:
                # path to data directory
                if language is None:
                    raise ValueError("You must implicitly provide language "
                                     "when providing data directory as source")
                reserve_data_path = data_path
                data_path = [data_path / "{}-ud-{}.conllu".format(language, mode)
                             for mode in data_types]
                reserve_data_path = [
                    reserve_data_path / language / "{}-ud-{}.conllu".format(language, mode)
                    for mode in data_types]
        else:
            data_path = [Path(data_path) for data_path in data_path]
            reserve_data_path = None
        if len(data_path) != len(data_types):
            raise ValueError("The number of input files in data_path and data types "
                             "in data_types must be equal")
        has_missing_files = any(not filepath.exists() for filepath in data_path)
        if has_missing_files and reserve_data_path is not None:
            has_missing_files = any(not filepath.exists() for filepath in reserve_data_path)
            if not has_missing_files:
                data_path = reserve_data_path
        if has_missing_files:
            # Files are downloaded from the Web repository
            dir_path = data_path[0].parent
            language = language or get_language(data_path[0].parts[-1])
            url = self.URL + "{}.tar.gz".format(language)
            log.info('[downloading data from {} to {}]'.format(url, dir_path))
            dir_path.mkdir(exist_ok=True, parents=True)
            download_decompress(url, dir_path)
            mark_done(dir_path)
        data = {}
        for mode, filepath in zip(data_types, data_path):
            if mode == "dev":
                mode = "valid"
#             if mode == "test":
#                 kwargs["read_only_words"] = True
            data[mode] = read_infile(filepath, **kwargs)
        return data


# Copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import random
from typing import Tuple, List, Dict, Any, Iterator

import numpy as np

from deeppavlov.core.common.registry import register
from deeppavlov.core.data.data_learning_iterator import DataLearningIterator
from deeppavlov.models.preprocessors.capitalization import process_word

def preprocess_data(data: List[Tuple[List[str], List[str]]], to_lower: bool = True,
                    append_case: str = "first") -> List[Tuple[List[Tuple[str]], List[str]]]:
    """Processes all words in data using
    :func:`~deeppavlov.dataset_iterators.morphotagger_iterator.process_word`.

    Args:
        data: a list of pairs (words, tags), each pair corresponds to a single sentence
        to_lower: whether to lowercase
        append_case: whether to add case mark

    Returns:
        a list of preprocessed sentences
    """
    new_data = []
    for words, tags in data:
        new_words = [process_word(word, to_lower=to_lower, append_case=append_case)
                     for word in words]
        # tags could also be processed in future
        new_tags = tags
        new_data.append((new_words, new_tags))
    return new_data

@register('morphotagger_dataset')
class MorphoTaggerDatasetIterator(DataLearningIterator):
    """
    Iterates over data for Morphological Tagging.
    A subclass of :class:`~deeppavlov.core.data.data_learning_iterator.DataLearningIterator`.

    Args:
        seed: random seed for data shuffling
        shuffle: whether to shuffle data during batching
        validation_split: the fraction of validation data
            (is used only if there is no `valid` subset in `data`)
        min_train_fraction: minimal fraction of train data in train+dev dataset,
                For fair comparison with UD Pipe it is set to 0.9 for UD experiments.
                It is actually used only for Turkish data.
    """
    def __init__(self, data: Dict[str, List[Tuple[Any, Any]]], seed: int = None,
                 shuffle: bool = True,  min_train_fraction: float = 0.0,
                 validation_split: float = 0.2) -> None:
        self.validation_split = validation_split
        self.min_train_fraction = min_train_fraction
        super().__init__(data, seed, shuffle)

    def split(self) -> None:
        """
        Splits the `train` part to `train` and `valid`, if no `valid` part is specified.
        Moves deficient data from `valid` to `train` if both parts are given,
        but `train` subset is too small.
        """
        if len(self.valid) == 0:
            if self.shuffle:
                random.shuffle(self.train)
            L = int(len(self.train) * (1.0 - self.validation_split))
            self.train, self.valid = self.train[:L], self.valid[L:]
        elif self.min_train_fraction > 0.0:
            train_length = len(self.train)
            valid_length = len(self.valid)
            gap = int(self.min_train_fraction * (train_length + valid_length)) - train_length
            if gap > 0:
                self.train.extend(self.valid[:gap])
                self.valid = self.valid[gap:]
        return

    def gen_batches(self, batch_size: int, data_type: str = 'train',
                    shuffle: bool = None, return_indexes: bool = False) -> Iterator[tuple]:
        """Generate batches of inputs and expected output to train neural networks

        Args:
            batch_size: number of samples in batch
            data_type: can be either 'train', 'test', or 'valid'
            shuffle: whether to shuffle dataset before batching
            return_indexes: whether to return indexes of batch elements in initial dataset

        Yields:
            a tuple of a batch of inputs and a batch of expected outputs.
            If `return_indexes` is True, also yields indexes of batch elements.
        """
        if shuffle is None:
            shuffle = self.shuffle
        data = self.data[data_type]
        if shuffle:
            random.shuffle(data)
        lengths = [len(x[0]) for x in data]
        indexes = np.argsort(lengths)
        L = len(data)
        if batch_size < 0:
            batch_size = L
        for start in range(0, L, batch_size):
            indexes_to_yield = indexes[start:start+batch_size]
            data_to_yield = tuple(list(x) for x in zip(*([data[i] for i in indexes_to_yield])))
            if return_indexes:
                yield indexes_to_yield, data_to_yield
            else:
                yield data_to_yield


import importlib
import json
from logging import getLogger
from pathlib import Path
from typing import Callable, Any

from deeppavlov.core.common.errors import ConfigError

log = getLogger(__name__)

_registry_path = Path(__file__).parent / 'metrics_registry.json'
if _registry_path.exists():
    with _registry_path.open(encoding='utf-8') as f:
        _REGISTRY = json.load(f)
else:
    _REGISTRY = {}

def fn_from_str(name: str) -> Callable[..., Any]:
    """Returns a function object with the name given in string."""
    try:
        module_name, fn_name = name.split(':')
    except ValueError:
        raise ConfigError('Expected function description in a `module.submodules:function_name` form, but got `{}`'
                          .format(name))

    return getattr(importlib.import_module(module_name), fn_name)

def register_metric(metric_name: str) -> Callable[..., Any]:
    """Decorator for metric registration."""
    def decorate(fn):
        fn_name = fn.__module__ + ':' + fn.__name__
        if metric_name in _REGISTRY and _REGISTRY[metric_name] != fn_name:
            log.warning('"{}" is already registered as a metric name, the old function will be ignored'
                        .format(metric_name))
        _REGISTRY[metric_name] = fn_name
        return fn
    return decorate

def get_metric_by_name(name: str) -> Callable[..., Any]:
    """Returns a metric callable with a corresponding name."""
    if name not in _REGISTRY:
        raise ConfigError(f'"{name}" is not registered as a metric')
    return fn_from_str(_REGISTRY[name])


# Copyright 2017 Neural Networks and Deep Learning lab, MIPT
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Union, Tuple, List, Optional
from logging import getLogger
from abc import abstractmethod
import math
from enum import IntEnum

import numpy as np

from deeppavlov.core.common.errors import ConfigError

log = getLogger(__name__)

class DecayType(IntEnum):
    """ Data class, each decay type is assigned a number. """
    NO = 1
    LINEAR = 2
    COSINE = 3
    EXPONENTIAL = 4
    POLYNOMIAL = 5
    ONECYCLE = 6
    TRAPEZOID = 7

    @classmethod
    def from_str(cls, label: str) -> int:
        """
        Convert given string label of decay type to special index

        Args:
            label: name of decay type.
                Set of values: `"linear"`, `"cosine"`, `"exponential"`,
                 `"onecycle"`, `"trapezoid"`, `["polynomial", K]`, where K is a polynomial power

        Returns:
            index of decay type
        """
        label_norm = label.replace('1', 'one').upper()
        if label_norm in cls.__members__:
            return DecayType[label_norm]
        else:
            raise NotImplementedError

class DecayScheduler:
    """
    Given initial and endvalue, this class generates the next value
    depending on decay type and number of iterations. (by calling next_val().)
    """

    def __init__(self, dec_type: Union[str, DecayType], start_val: float,
                 num_it: int = 0, end_val: float = None, extra: float = None) -> None:
        if isinstance(dec_type, DecayType):
            self.dec_type = dec_type
        else:
            self.dec_type = DecayType.from_str(dec_type)
        self.nb, self.extra = num_it, extra
        self.start_val, self.end_val = start_val, end_val
        self.iters = 0
        if self.end_val is None and not (self.dec_type in [1, 4]):
            self.end_val = 0
        if self.dec_type == DecayType.ONECYCLE:
            self.cycle_nb = math.ceil(self.nb / 2)
            self.div = 1.0 if not self.start_val else self.end_val / self.start_val
        if self.dec_type == DecayType.TRAPEZOID:
            self.div = 1.0 if not self.start_val else self.end_val / self.start_val

    def __str__(self):
        return f"DecayScheduler(start_val={self.start_val}, end_val={self.end_val}"\
            f", dec_type={self.dec_type.name}, num_it={self.nb}, extra={self.extra})"

    def next_val(self) -> float:
        self.iters = min(self.iters + 1, self.nb)
        # print(f"iters = {self.iters}/{self.nb}")
        if self.dec_type == DecayType.NO:
            return self.start_val
        elif self.dec_type == DecayType.LINEAR:
            pct = self.iters / self.nb
            return self.start_val + pct * (self.end_val - self.start_val)
        elif self.dec_type == DecayType.COSINE:
            cos_out = math.cos(math.pi * self.iters / self.nb) + 1
            return self.end_val + (self.start_val - self.end_val) / 2 * cos_out
        elif self.dec_type == DecayType.EXPONENTIAL:
            ratio = self.end_val / self.start_val
            return self.start_val * (ratio ** (self.iters / self.nb))
        elif self.dec_type == DecayType.POLYNOMIAL:
            delta_val = self.start_val - self.end_val
            return self.end_val + delta_val * (1 - self.iters / self.nb) ** self.extra
        elif self.dec_type == DecayType.ONECYCLE:
            if self.iters > self.cycle_nb:
                # decaying from end_val to start_val for cycle_nb steps
                pct = 1 - (self.iters - self.cycle_nb) / self.cycle_nb
                return self.start_val * (1 + pct * (self.div - 1))
            else:
                # raising from start_val to end_val for cycle_nb steps
                pct = self.iters / self.cycle_nb
                return self.start_val * (1 + pct * (self.div - 1))
        elif self.dec_type == DecayType.TRAPEZOID:
            if self.iters > 0.6 * self.nb:
                # decaying from end_val to start_val for 4/10 * nb steps
                pct = 2.5 * (self.nb - self.iters) / self.nb
                return self.start_val * (1 + pct * (self.div - 1))
            elif self.iters > 0.1 * self.nb:
                # constant end_val
                return self.end_val
            else:
                # raising from start_val to end_val for 1/10 * nb steps
                pct = 10.0 * self.iters / self.nb
                return self.start_val * (1 + pct * (self.div - 1))

DType = Union[str, DecayType]

class LRScheduledModel:
    """
    Abstract model enhanced with optimizer, learning rate and momentum
    management and search.

    Args:
        learning_rate: learning rate value or ranges
        learning_rate_decay: learning rate decay type
                Set of values: `"linear"`, `"onecycle"`, `"trapezoid"`,
                               `"exponential"`, `"cosine"`, `["polynomial", K]`, where K is a polynomial power
        learning_rate_decay_epochs: number of epochs for learning rate decay process
        learning_rate_decay_batches: number of batches for learning rate decay process
        learning_rate_drop_div: division coefficient for learning rate in case of
                exceeding patience `learning_rate_drop_patience`
        learning_rate_drop_patience: patience limit of loss increase
        momentum: range of momentum values
        momentum_decay: momentum decay type.
                Set of values: `"linear"`, `"onecycle"`, `"trapezoid"`,
                               `"exponential"`, `"cosine"`, `["polynomial", K]`, where K is a polynomial power
        momentum_decay_epochs: number of epochs for momentum decay process
        momentum_decay_batches: number of batches for momentum decay process
        fit_batch_size: batch size when fitting learning rate
        fit_learning_rate: range of learning rate values to explore
        fit_learning_rate_div: division coefficient for best learning rate obtained from fitting,
                divided learning rate value will be used when training model
        fit_beta: smoothing coefficient for loss calculation when fitting learning rate
        fit_min_batches: number of batches to train model on before fitting learning rate
        fit_max_batches: number of batches to train model on when fitting learning rate
        *args: other parameters
        **kwargs: other parameters
    """

    @abstractmethod
    def _init_learning_rate_variable(self):
        pass

    @abstractmethod
    def _init_momentum_variable(self):
        pass

    @abstractmethod
    def _update_graph_variables(self, learning_rate=None, momentum=None):
        """
        Update learning rate in graph if `learning_rate` is not None,
        Update momentum in graph if `momentum` is not None
        """
        if learning_rate is not None:
            # update learning rate
            pass
        if momentum is not None:
            # update momentum
            pass

    def __init__(self,
                 learning_rate: Union[None, float, Tuple[float, float]] = None,
                 learning_rate_decay: Union[DType, Tuple[DType, float]] = DecayType.NO,
                 learning_rate_decay_epochs: int = 0,
                 learning_rate_decay_batches: int = 0,
                 learning_rate_drop_div: float = 2.0,
                 learning_rate_drop_patience: Optional[int] = None,
                 momentum: Union[None, float, Tuple[float, float]] = None,
                 momentum_decay: Union[DType, Tuple[DType, float]] = DecayType.NO,
                 momentum_decay_epochs: int = 0,
                 momentum_decay_batches: int = 0,
                 fit_batch_size: Union[None, int, str] = None,
                 fit_learning_rate: Tuple[float, float] = (1e-7, 100),
                 fit_learning_rate_div: float = 10.,
                 fit_beta: float = 0.98,
                 fit_min_batches: int = 10,
                 fit_max_batches: Optional[int] = None,
                 *args, **kwargs) -> None:
        """
        Initialize learning rate scheduler
        """
        if learning_rate_decay_epochs and learning_rate_decay_batches:
            raise ConfigError("isn't able to update learning rate every batch"
                              " and every epoch simultaneously")
        if momentum_decay_epochs and momentum_decay_batches:
            raise ConfigError("isn't able to update momentum every batch"
                              " and every epoch simultaneously")

        start_val, end_val = learning_rate, None
        if isinstance(learning_rate, (tuple, list)):
            start_val, end_val = learning_rate
        dec_type, extra = learning_rate_decay, None
        if isinstance(learning_rate_decay, (tuple, list)):
            dec_type, extra = learning_rate_decay

        self._lr = start_val
        num_it, self._lr_update_on_batch = learning_rate_decay_epochs, False
        if learning_rate_decay_batches > 0:
            num_it, self._lr_update_on_batch = learning_rate_decay_batches, True

        self._lr_schedule = DecayScheduler(start_val=start_val, end_val=end_val,
                                           num_it=num_it, dec_type=dec_type, extra=extra)
        self._lr_var = self._init_learning_rate_variable()

        start_val, end_val = momentum, None
        if isinstance(momentum, (tuple, list)):
            start_val, end_val = momentum
        dec_type, extra = momentum_decay, None
        if isinstance(momentum_decay, (tuple, list)):
            dec_type, extra = momentum_decay

        self._mom = start_val
        num_it, self._mom_update_on_batch = momentum_decay_epochs, False
        self._mom_update_on_batch = momentum_decay_batches > 0
        num_it = momentum_decay_epochs if self._mom_update_on_batch else momentum_decay_batches

        self._mom_schedule = DecayScheduler(start_val=start_val, end_val=end_val,
                                            num_it=num_it, dec_type=dec_type,
                                            extra=extra)
        self._mom_var = self._init_momentum_variable()

        self._learning_rate_drop_patience = learning_rate_drop_patience
        self._learning_rate_drop_div = learning_rate_drop_div
        self._learning_rate_cur_impatience = 0.
        self._learning_rate_last_impatience = 0.
        self._learning_rate_cur_div = 1.
        self._fit_batch_size = fit_batch_size
        self._fit_learning_rate = fit_learning_rate
        self._fit_learning_rate_div = fit_learning_rate_div
        self._fit_beta = fit_beta
        self._fit_min_batches = fit_min_batches
        self._fit_max_batches = fit_max_batches

    def get_learning_rate(self):
        """
        Return current learning rate value

        Returns:
            learning rate
        """
        if self._lr is None:
            raise ConfigError("Please specify `learning_rate` parameter"
                              " before training")
        return self._lr

    def get_learning_rate_variable(self):
        """
        Return current learning rate variable

        Returns:
            learning rate variable
        """
        return self._lr_var

    def get_momentum(self):
        """
        Return current momentum value

        Returns:
            momentum
        """
        return self._mom

    def get_momentum_variable(self):
        """
        Return current momentum variable

        Returns:
            momentum variable
        """
        return self._mom_var

    def fit(self, *args):
        """
        Find the best learning rate schedule, and set obtained values of learning rate
        and momentum for further model training. Best learning rate will be divided
        by `fit_learning_rate_div` for further training model.

        Args:
            *args: arguments

        Returns:

        """
        data = list(zip(*args))
        self.save()
        if self._fit_batch_size is None:
            raise ConfigError("in order to use fit() method"
                              " set `fit_batch_size` parameter")
        bs = int(self._fit_batch_size)
        data_len = len(data)
        num_batches = self._fit_max_batches or ((data_len - 1) // bs + 1)

        avg_loss = 0.
        best_loss = float('inf')
        lrs, losses = [], []
        _lr_find_schedule = DecayScheduler(start_val=self._fit_learning_rate[0],
                                           end_val=self._fit_learning_rate[1],
                                           dec_type="exponential",
                                           num_it=num_batches)
        self._lr = _lr_find_schedule.start_val
        self._mom = 0.
        self._update_graph_variables(learning_rate=self._lr, momentum=self._mom)
        best_lr = _lr_find_schedule.start_val
        for i in range(num_batches):
            batch_start = (i * bs) % data_len
            batch_end = batch_start + bs
            report = self.train_on_batch(*zip(*data[batch_start:batch_end]))
            if not isinstance(report, dict):
                report = {'loss': report}
            # Calculating smoothed loss
            avg_loss = self._fit_beta*avg_loss + (1 - self._fit_beta)*report['loss']
            smoothed_loss = avg_loss / (1 - self._fit_beta**(i + 1))
            lrs.append(self._lr)
            losses.append(smoothed_loss)
            log.info(f"Batch {i}/{num_batches}: smooth_loss = {smoothed_loss}"
                     f", lr = {self._lr}, best_lr = {best_lr}")
            if math.isnan(smoothed_loss) or (smoothed_loss > 4 * best_loss):
                break
            if (smoothed_loss < best_loss) and (i >= self._fit_min_batches):
                best_loss = smoothed_loss
                best_lr = self._lr
            self._lr = _lr_find_schedule.next_val()
            self._update_graph_variables(learning_rate=self._lr)

            if i >= num_batches:
                break
        # best_lr /= 10
        end_val = self._get_best(lrs, losses)

        start_val = end_val
        if self._lr_schedule.dec_type in (DecayType.ONECYCLE, DecayType.TRAPEZOID):
            start_val = end_val / self._fit_learning_rate_div
        elif self._lr_schedule.dec_type in (DecayType.POLYNOMIAL, DecayType.EXPONENTIAL,
                                            DecayType.LINEAR, DecayType.COSINE):
            start_val = end_val
            end_val = end_val / self._fit_learning_rate_div
        self._lr_schedule = DecayScheduler(start_val=start_val,
                                           end_val=end_val,
                                           num_it=self._lr_schedule.nb,
                                           dec_type=self._lr_schedule.dec_type,
                                           extra=self._lr_schedule.extra)
        log.info(f"Found best learning rate value = {best_lr}"
                 f", setting new learning rate schedule with {self._lr_schedule}.")

        self.load()
        self._lr = self._lr_schedule.start_val
        self._mom = self._mom_schedule.start_val
        self._update_graph_variables(learning_rate=self._lr, momentum=self._mom)
        return {'smoothed_loss': losses, 'learning_rate': lrs}

    @staticmethod
    def _get_best(values: List[float], losses: List[float],
                  max_loss_div: float = 0.9, min_val_div: float = 10.0) -> float:
        """
        Find the best value according to given losses

        Args:
            values: list of considered values
            losses: list of obtained loss values corresponding to `values`
            max_loss_div: maximal divergence of loss to be considered significant
            min_val_div: minimum divergence of loss to be considered significant

        Returns:
            best value divided by `min_val_div`
        """
        assert len(values) == len(losses), "lengths of values and losses should be equal"
        min_ind = np.argmin(losses)
        for i in range(min_ind - 1, 0, -1):
            if (losses[i] * max_loss_div > losses[min_ind]) or\
                    (values[i] * min_val_div < values[min_ind]):
                return values[i + 1]
        return values[min_ind] / min_val_div

    def process_event(self, event_name: str, data: dict) -> None:
        """
        Update learning rate and momentum variables after event (given by `event_name`)

        Args:
            event_name: name of event after which the method was called.
                    Set of values: `"after_validation"`, `"after_batch"`, `"after_epoch"`, `"after_train_log"`
            data: dictionary with parameters values

        Returns:
            None
        """
        if event_name == "after_validation":
            if data['impatience'] > self._learning_rate_last_impatience:
                self._learning_rate_cur_impatience += 1
            else:
                self._learning_rate_cur_impatience = 0

            self._learning_rate_last_impatience = data['impatience']

            if (self._learning_rate_drop_patience is not None) and\
                    (self._learning_rate_cur_impatience >=
                     self._learning_rate_drop_patience):
                self._learning_rate_cur_impatience = 0
                self._learning_rate_cur_div *= self._learning_rate_drop_div
                self._lr /= self._learning_rate_drop_div
                self._update_graph_variables(learning_rate=self._lr)
                log.info(f"New learning rate dividor = {self._learning_rate_cur_div}")
        if event_name == 'after_batch':
            if (self._lr is not None) and self._lr_update_on_batch:
                self._lr = self._lr_schedule.next_val() / self._learning_rate_cur_div
                self._update_graph_variables(learning_rate=self._lr)
            if (self._mom is not None) and self._mom_update_on_batch:
                self._mom = min(1., max(0., self._mom_schedule.next_val()))
                self._update_graph_variables(momentum=self._mom)
        if event_name == 'after_epoch':
            if (self._lr is not None) and not self._lr_update_on_batch:
                self._lr = self._lr_schedule.next_val() / self._learning_rate_cur_div
                self._update_graph_variables(learning_rate=self._lr)
            if (self._mom is not None) and not self._mom_update_on_batch:
                self._mom = min(1., max(0., self._mom_schedule.next_val()))
                self._update_graph_variables(momentum=self._mom)
        if event_name == 'after_train_log':
            if (self._lr is not None) and ('learning_rate' not in data):
                data['learning_rate'] = self._lr
            if (self._mom is not None) and ('momentum' not in data):
                data['momentum'] = self._mom