# 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 ssl
from collections import namedtuple
from logging import getLogger
from pathlib import Path
from typing import Union, Optional

from flasgger import Swagger, swag_from
from flask import Flask, request, jsonify, redirect
from flask_cors import CORS

from deeppavlov import build_model
from deeppavlov.agents.default_agent.default_agent import DefaultAgent
from deeppavlov.agents.processors.default_rich_content_processor import DefaultRichContentWrapper
from deeppavlov.core.agent import Agent
from deeppavlov.core.agent.rich_content import RichMessage
from deeppavlov.core.common.paths import get_settings_path
from deeppavlov.skills.default_skill.default_skill import DefaultStatelessSkill
from deeppavlov.utils.server.server import get_server_params

SERVER_CONFIG_FILENAME = 'server_config.json'

log = getLogger(__name__)

app = Flask(__name__)
Swagger(app)
CORS(app)

DialogID = namedtuple('DialogID', ['user_id', 'session_id'])

def interact_alice(agent: Agent):
    """
    Exchange messages between basic pipelines and the Yandex.Dialogs service.
    If the pipeline returns multiple values, only the first one is forwarded to Yandex.
    """
    data = request.get_json()
    text = data['request'].get('command', '').strip()
    payload = data['request'].get('payload')

    session_id = data['session']['session_id']
    user_id = data['session']['user_id']
    message_id = data['session']['message_id']

    dialog_id = DialogID(user_id, session_id)

    response = {
        'response': {
            'end_session': True,
            'text': ''
        },
        "session": {
            'session_id': session_id,
            'message_id': message_id,
            'user_id': user_id
        },
        'version': '1.0'
    }

    agent_response: Union[str, RichMessage] = agent([payload or text], [dialog_id])[0]
    if isinstance(agent_response, RichMessage):
        response['response']['text'] = '\n'.join([j['content']
                                                  for j in agent_response.json()
                                                  if j['type'] == 'plain_text'])
    else:
        response['response']['text'] = str(agent_response)

    return jsonify(response), 200

def start_alice_server(model_config, https=False, ssl_key=None, ssl_cert=None, port=None):
    server_config_path = get_settings_path() / SERVER_CONFIG_FILENAME
    server_params = get_server_params(server_config_path, model_config)

    https = https or server_params['https']

    if not https:
        ssl_key = ssl_cert = None
    else:
        ssh_key = Path(ssl_key or server_params['https_key_path']).resolve()
        if not ssh_key.is_file():
            e = FileNotFoundError('Ssh key file not found: please provide correct path in --key param or '
                                  'https_key_path param in server configuration file')
            log.error(e)
            raise e

        ssh_cert = Path(ssl_cert or server_params['https_cert_path']).resolve()
        if not ssh_cert.is_file():
            e = FileNotFoundError('Ssh certificate file not found: please provide correct path in --cert param or '
                                  'https_cert_path param in server configuration file')
            log.error(e)
            raise e

    host = server_params['host']
    port = port or server_params['port']
    model_endpoint = server_params['model_endpoint']

    model = build_model(model_config)
    skill = DefaultStatelessSkill(model, lang='ru')
    agent = DefaultAgent([skill], skills_processor=DefaultRichContentWrapper())

    start_agent_server(agent, host, port, model_endpoint, ssl_key, ssl_cert)

def start_agent_server(agent: Agent, host: str, port: int, endpoint: str,
                       ssl_key: Optional[Path] = None,
                       ssl_cert: Optional[Path] = None):
    if ssl_key and ssl_cert:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
        ssh_key_path = Path(ssl_key).resolve()
        ssh_cert_path = Path(ssl_cert).resolve()
        ssl_context.load_cert_chain(ssh_cert_path, ssh_key_path)
    else:
        ssl_context = None

    @app.route('/')
    def index():
        return redirect('/apidocs/')

    endpoint_description = {
        'description': 'A model endpoint',
        'parameters': [
            {
                'name': 'data',
                'in': 'body',
                'required': 'true',
                'example': {
                    'meta': {
                        'locale': 'ru-RU',
                        'timezone': 'Europe/Moscow',
                        "client_id": 'ru.yandex.searchplugin/5.80 (Samsung Galaxy; Android 4.4)'
                    },
                    'request': {
                        'command': 'где ближайшее отделение',
                        'original_utterance': 'Алиса спроси у Сбербанка где ближайшее отделение',
                        'type': 'SimpleUtterance',
                        'markup': {
                            'dangerous_context': True
                        },
                        'payload': {}
                    },
                    'session': {
                        'new': True,
                        'message_id': 4,
                        'session_id': '2eac4854-fce721f3-b845abba-20d60',
                        'skill_id': '3ad36498-f5rd-4079-a14b-788652932056',
                        'user_id': 'AC9WC3DF6FCE052E45A4566A48E6B7193774B84814CE49A922E163B8B29881DC'
                    },
                    'version': '1.0'
                }
            }
        ],
        'responses': {
            "200": {
                "description": "A model response"
            }
        }
    }

    @app.route(endpoint, methods=['POST'])
    @swag_from(endpoint_description)
    def answer():
        return interact_alice(agent)

    app.run(host=host, port=port, threaded=False, ssl_context=ssl_context)


# 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 logging import getLogger
from typing import List

import numpy as np

log = getLogger(__name__)

def labels2onehot(labels: [List[str], List[List[str]], np.ndarray], classes:  [list, np.ndarray]) -> np.ndarray:
    """
    Convert labels to one-hot vectors for multi-class multi-label classification

    Args:
        labels: list of samples where each sample is a class or a list of classes which sample belongs with
        classes: array of classes' names

    Returns:
        2d array with one-hot representation of given samples
    """
    n_classes = len(classes)
    y = []
    for sample in labels:
        curr = np.zeros(n_classes)
        if isinstance(sample, list):
            for intent in sample:
                if intent not in classes:
                    log.warning('Unknown intent {} detected. Assigning no class'.format(intent))
                else:
                    curr[np.where(np.array(classes) == intent)[0]] = 1
        else:
            curr[np.where(np.array(classes) == sample)[0]] = 1
        y.append(curr)
    y = np.asarray(y)
    return y

def proba2labels(proba: [list, np.ndarray], confident_threshold: float, classes:  [list, np.ndarray]) -> List[List]:
    """
    Convert vectors of probabilities to labels using confident threshold
    (if probability to belong with the class is bigger than confident_threshold, sample belongs with the class;
    if no probabilities bigger than confident threshold, sample belongs with the class with the biggest probability)

    Args:
        proba: list of samples where each sample is a vector of probabilities to belong with given classes
        confident_threshold (float): boundary of probability to belong with a class
        classes: array of classes' names

    Returns:
        list of lists of labels for each sample
    """
    y = []
    for sample in proba:
        to_add = np.where(sample > confident_threshold)[0]
        if len(to_add) > 0:
            y.append(np.array(classes)[to_add].tolist())
        else:
            y.append(np.array([np.array(classes)[np.argmax(sample)]]).tolist())

    return y

def proba2onehot(proba: [list, np.ndarray], confident_threshold: float, classes:  [list, np.ndarray]) -> np.ndarray:
    """
    Convert vectors of probabilities to one-hot representations using confident threshold

    Args:
        proba: samples where each sample is a vector of probabilities to belong with given classes
        confident_threshold: boundary of probability to belong with a class
        classes: array of classes' names

    Returns:
        2d array with one-hot representation of given samples
    """
    return labels2onehot(proba2labels(proba, confident_threshold, classes), classes)


# 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 inspect
from abc import abstractmethod
from copy import deepcopy
from logging import getLogger
from typing import Optional, List, Union

import numpy as np
import tensorflow as tf
from keras import backend as K
from overrides import overrides

from deeppavlov.core.models.nn_model import NNModel
from deeppavlov.core.models.tf_backend import TfModelMeta
from deeppavlov.core.models.lr_scheduled_model import LRScheduledModel

log = getLogger(__name__)

class KerasModel(NNModel, metaclass=TfModelMeta):
    """
    Builds Keras model with TensorFlow backend.

    Attributes:
        epochs_done: number of epochs that were done
        batches_seen: number of epochs that were seen
        train_examples_seen: number of training samples that were seen
        sess: tf session
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize model using keyword parameters

        Args:
            kwargs (dict): Dictionary with model parameters
        """
        self.epochs_done = 0
        self.batches_seen = 0
        self.train_examples_seen = 0

        super().__init__(save_path=kwargs.get("save_path"),
                         load_path=kwargs.get("load_path"),
                         mode=kwargs.get("mode"))

    @staticmethod
    def _config_session():
        """
        Configure session for particular device

        Returns:
            tensorflow.Session
        """
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        config.gpu_options.visible_device_list = '0'
        return tf.Session(config=config)

    @abstractmethod
    def load(self, *args, **kwargs) -> None:
        pass

    @abstractmethod
    def save(self, *args, **kwargs) -> None:
        pass

    def process_event(self, event_name: str, data: dict) -> None:
        """
        Process event after epoch
        Args:
            event_name: whether event is send after epoch or batch.
                    Set of values: ``"after_epoch", "after_batch"``
            data: event data (dictionary)

        Returns:
            None
        """
        if event_name == "after_epoch":
            self.epochs_done = data["epochs_done"]
            self.batches_seen = data["batches_seen"]
            self.train_examples_seen = data["train_examples_seen"]
        return

class KerasWrapper(KerasModel):
    """A wrapper over external Keras models. It is used, for example,
    to wrap :class:`~deeppavlov.models.morpho_tagger.network.CharacterTagger`.
    A subclass of :class:`~deeppavlov.core.models.keras_model.KerasModel`

    Attributes:
        cls: the class to be wrapped
        save_path: the path where model is saved
        load_path: the path from where model is loaded
        mode: usage mode
        **kwargs: a dictionary containing model parameters specified in the main part
            of json config that corresponds to the model
    """
    def __init__(self, cls: type, save_path: Optional[str] = None,
                 load_path: Optional[str] = None, mode: str = None,
                 **kwargs) -> None:
        # Calls parent constructor. Results in creation of save_folder if it doesn't exist
        super().__init__(save_path=save_path, load_path=load_path, mode=mode)

        # Dicts are mutable! To prevent changes in config dict outside this class
        # we use deepcopy
        opt = deepcopy(kwargs)

        # Finds all input parameters of the network __init__ to pass them into network later
        network_parameter_names = list(inspect.signature(cls.__init__).parameters)
        # Fills all provided parameters from opt (opt is a dictionary formed from the model
        # json config file, except the "name" field)
        network_parameters = {par: opt[par] for par in network_parameter_names if par in opt}
        self._net = cls(**network_parameters)

        # Finds all parameters for network train to pass them into train method later
        train_parameters_names = list(inspect.signature(self._net.train_on_batch).parameters)

        # Fills all provided parameters from opt
        train_parameters = {par: opt[par] for par in train_parameters_names if par in opt}
        self.train_parameters = train_parameters
        self.opt = opt

        # Tries to load the model from model `load_path`, if it is available
        self.load()

    def load(self) -> None:
        """Checks existence of the model file, loads the model if the file exists"""

        # Checks presence of the model files
        if self.load_path.exists():
            path = str(self.load_path.resolve())
            log.info('[loading model from {}]'.format(path))
            self._net.load(path)

    def save(self) -> None:
        """Saves model to the save_path, provided in config. The directory is
        already created by super().__init__, which is called in __init__ of this class"""
        path = str(self.save_path.absolute())
        log.info('[saving model to {}]'.format(path))
        self._net.save(path)

    def train_on_batch(self, *args) -> None:
        """Trains the model on a single batch.

        Args:
            *args: the list of network inputs.
            Last element of `args` is the batch of targets,
            all previous elements are training data batches
        """
        *data, labels = args
        self._net.train_on_batch(data, labels)

    def __call__(self, *x_batch, **kwargs) -> Union[List, np.ndarray]:
        """
        Predicts answers on batch elements.

        Args:
            instance: a batch to predict answers on
        """
        with self.graph.as_default():
            K.set_session(self.sess)
            return self._net.predict_on_batch(x_batch, **kwargs)

class LRScheduledKerasModel(LRScheduledModel, KerasModel):
    """
    KerasModel enhanced with optimizer, learning rate and momentum
    management and search.
    """
    def __init__(self, **kwargs):
        """
        Initialize model with given parameters

        Args:
            **kwargs: dictionary of parameters
        """
        if isinstance(kwargs.get("learning_rate"), float) and isinstance(kwargs.get("learning_rate_decay"), float):
            KerasModel.__init__(self, **kwargs)
        else:
            KerasModel.__init__(self, **kwargs)
            LRScheduledModel.__init__(self, **kwargs)

    @abstractmethod
    def get_optimizer(self):
        """
        Return instance of keras optimizer

        Args:
            None
        """
        pass

    @overrides
    def _init_learning_rate_variable(self):
        """
        Initialize learning rate

        Returns:
            None
        """
        return None

    @overrides
    def _init_momentum_variable(self):
        """
        Initialize momentum

        Returns:
            None
        """
        return None

    @overrides
    def get_learning_rate_variable(self):
        """
        Extract value of learning rate from optimizer

        Returns:
            learning rate value
        """
        return self.get_optimizer().lr

    @overrides
    def get_momentum_variable(self):
        """
        Extract values of momentum variables from optimizer

        Returns:
            optimizer's `rho` or `beta_1`
        """
        optimizer = self.get_optimizer()
        if hasattr(optimizer, 'rho'):
            return optimizer.rho
        elif hasattr(optimizer, 'beta_1'):
            return optimizer.beta_1
        return None

    @overrides
    def _update_graph_variables(self, learning_rate: float = None, momentum: float = None):
        """
        Update graph variables setting giving `learning_rate` and `momentum`

        Args:
            learning_rate: learning rate value to be set in graph (set if not None)
            momentum: momentum value to be set in graph (set if not None)

        Returns:
            None
        """
        if learning_rate is not None:
            K.set_value(self.get_learning_rate_variable(), learning_rate)
            # log.info(f"Learning rate = {learning_rate}")
        if momentum is not None:
            K.set_value(self.get_momentum_variable(), momentum)
            # log.info(f"Momentum      = {momentum}")

    def process_event(self, event_name: str, data: dict):
        """
        Process event after epoch
        Args:
            event_name: whether event is send after epoch or batch.
                    Set of values: ``"after_epoch", "after_batch"``
            data: event data (dictionary)

        Returns:
            None
        """
        if (isinstance(self.opt.get("learning_rate", None), float) and
                isinstance(self.opt.get("learning_rate_decay", None), float)):
            pass
        else:
            if event_name == 'after_train_log':
                if (self.get_learning_rate_variable() is not None) and ('learning_rate' not in data):
                    data['learning_rate'] = float(K.get_value(self.get_learning_rate_variable()))
                    # data['learning_rate'] = self._lr
                if (self.get_momentum_variable() is not None) and ('momentum' not in data):
                    data['momentum'] = float(K.get_value(self.get_momentum_variable()))
                    # data['momentum'] = self._mom
            else:
                super().process_event(event_name, 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 itertools
from collections import OrderedDict
from itertools import chain
from logging import getLogger

import numpy as np
from sklearn.metrics import f1_score

from deeppavlov.core.common.metrics_registry import register_metric

log = getLogger(__name__)

@register_metric('ner_f1')
def ner_f1(y_true, y_predicted):
    y_true = list(chain(*y_true))
    y_predicted = list(chain(*y_predicted))
    results = precision_recall_f1(y_true,
                                  y_predicted,
                                  print_results=True)
    f1 = results['__total__']['f1']
    return f1

@register_metric('f1')
def round_f1(y_true, y_predicted):
    """
    Calculates F1 (binary) measure.

    Args:
        y_true: list of true values
        y_predicted: list of predicted values

    Returns:
        F1 score
    """
    try:
        predictions = [np.round(x) for x in y_predicted]
    except TypeError:
        predictions = y_predicted

    return f1_score(y_true, predictions)

@register_metric('f1_macro')
def round_f1_macro(y_true, y_predicted):
    """
    Calculates F1 macro measure.

    Args:
        y_true: list of true values
        y_predicted: list of predicted values

    Returns:
        F1 score
    """
    try:
        predictions = [np.round(x) for x in y_predicted]
    except TypeError:
        predictions = y_predicted

    return f1_score(np.array(y_true), np.array(predictions), average="macro")

@register_metric('f1_weighted')
def round_f1_weighted(y_true, y_predicted):
    """
    Calculates F1 weighted measure.

    Args:
        y_true: list of true values
        y_predicted: list of predicted values

    Returns:
        F1 score
    """
    try:
        predictions = [np.round(x) for x in y_predicted]
    except TypeError:
        predictions = y_predicted

    return f1_score(np.array(y_true), np.array(predictions), average="weighted")

def chunk_finder(current_token, previous_token, tag):
    current_tag = current_token.split('-', 1)[-1]
    previous_tag = previous_token.split('-', 1)[-1]
    if previous_tag != tag:
        previous_tag = 'O'
    if current_tag != tag:
        current_tag = 'O'
    if (previous_tag == 'O' and current_token == 'B-' + tag) or \
            (previous_token == 'I-' + tag and current_token == 'B-' + tag) or \
            (previous_token == 'B-' + tag and current_token == 'B-' + tag) or \
            (previous_tag == 'O' and current_token == 'I-' + tag):
        create_chunk = True
    else:
        create_chunk = False

    if (previous_token == 'I-' + tag and current_token == 'B-' + tag) or \
            (previous_token == 'B-' + tag and current_token == 'B-' + tag) or \
            (current_tag == 'O' and previous_token == 'I-' + tag) or \
            (current_tag == 'O' and previous_token == 'B-' + tag):
        pop_out = True
    else:
        pop_out = False
    return create_chunk, pop_out

def precision_recall_f1(y_true, y_pred, print_results=True, short_report=False, entity_of_interest=None):
    # Find all tags
    tags = set()
    for tag in itertools.chain(y_true, y_pred):
        if tag != 'O':
            current_tag = tag[2:]
            tags.add(current_tag)
    tags = sorted(list(tags))

    results = OrderedDict()
    for tag in tags:
        results[tag] = OrderedDict()
    results['__total__'] = OrderedDict()
    n_tokens = len(y_true)
    total_correct = 0
    # Firstly we find all chunks in the ground truth and prediction
    # For each chunk we write starting and ending indices

    for tag in tags:
        count = 0
        true_chunk = []
        pred_chunk = []
        y_true = [str(y) for y in y_true]
        y_pred = [str(y) for y in y_pred]
        prev_tag_true = 'O'
        prev_tag_pred = 'O'
        while count < n_tokens:
            yt = y_true[count]
            yp = y_pred[count]

            create_chunk_true, pop_out_true = chunk_finder(yt, prev_tag_true, tag)
            if pop_out_true:
                true_chunk[-1] = (true_chunk[-1], count - 1)
            if create_chunk_true:
                true_chunk.append(count)

            create_chunk_pred, pop_out_pred = chunk_finder(yp, prev_tag_pred, tag)
            if pop_out_pred:
                pred_chunk[-1] = (pred_chunk[-1], count - 1)
            if create_chunk_pred:
                pred_chunk.append(count)
            prev_tag_true = yt
            prev_tag_pred = yp
            count += 1

        if len(true_chunk) > 0 and not isinstance(true_chunk[-1], tuple):
            true_chunk[-1] = (true_chunk[-1], count - 1)
        if len(pred_chunk) > 0 and not isinstance(pred_chunk[-1], tuple):
            pred_chunk[-1] = (pred_chunk[-1], count - 1)

        # Then we find all correctly classified intervals
        # True positive results
        tp = len(set(pred_chunk).intersection(set(true_chunk)))
        total_correct += tp
        # And then just calculate errors of the first and second kind
        # False negative
        fn = len(true_chunk) - tp
        # False positive
        fp = len(pred_chunk) - tp
        if tp + fp > 0:
            precision = tp / (tp + fp) * 100
        else:
            precision = 0
        if tp + fn > 0:
            recall = tp / (tp + fn) * 100
        else:
            recall = 0
        if precision + recall > 0:
            f1 = 2 * precision * recall / (precision + recall)
        else:
            f1 = 0
        results[tag]['precision'] = precision
        results[tag]['recall'] = recall
        results[tag]['f1'] = f1
        results[tag]['n_predicted_entities'] = len(pred_chunk)
        results[tag]['n_true_entities'] = len(true_chunk)
    total_true_entities = 0
    total_predicted_entities = 0
    total_precision = 0
    total_recall = 0
    total_f1 = 0
    for tag in results:
        if tag == '__total__':
            continue
        n_pred = results[tag]['n_predicted_entities']
        n_true = results[tag]['n_true_entities']
        total_true_entities += n_true
        total_predicted_entities += n_pred
        total_precision += results[tag]['precision'] * n_pred
        total_recall += results[tag]['recall'] * n_true
        total_f1 += results[tag]['f1'] * n_true
    if total_true_entities > 0:
        accuracy = total_correct / total_true_entities * 100
        total_recall = total_recall / total_true_entities
    else:
        accuracy = 0
        total_recall = 0
    if total_predicted_entities > 0:
        total_precision = total_precision / total_predicted_entities
    else:
        total_precision = 0

    if total_precision + total_recall > 0:
        total_f1 = 2 * total_precision * total_recall / (total_precision + total_recall)
    else:
        total_f1 = 0

    results['__total__']['n_predicted_entities'] = total_predicted_entities
    results['__total__']['n_true_entities'] = total_true_entities
    results['__total__']['precision'] = total_precision
    results['__total__']['recall'] = total_recall
    results['__total__']['f1'] = total_f1

    if print_results:
        s = 'processed {len} tokens ' \
            'with {tot_true} phrases; ' \
            'found: {tot_pred} phrases;' \
            ' correct: {tot_cor}.\n\n'.format(len=n_tokens,
                                              tot_true=total_true_entities,
                                              tot_pred=total_predicted_entities,
                                              tot_cor=total_correct)

        s += 'precision:  {tot_prec:.2f}%; ' \
             'recall:  {tot_recall:.2f}%; ' \
             'FB1:  {tot_f1:.2f}\n\n'.format(acc=accuracy,
                                             tot_prec=total_precision,
                                             tot_recall=total_recall,
                                             tot_f1=total_f1)

        if not short_report:
            for tag in tags:
                if entity_of_interest is not None:
                    if entity_of_interest in tag:
                        s += '\t' + tag + ': precision:  {tot_prec:.2f}%; ' \
                                          'recall:  {tot_recall:.2f}%; ' \
                                          'F1:  {tot_f1:.2f} ' \
                                          '{tot_predicted}\n\n'.format(tot_prec=results[tag]['precision'],
                                                                       tot_recall=results[tag]['recall'],
                                                                       tot_f1=results[tag]['f1'],
                                                                       tot_predicted=results[tag]['n_predicted_entities'])
                elif tag != '__total__':
                    s += '\t' + tag + ': precision:  {tot_prec:.2f}%; ' \
                                      'recall:  {tot_recall:.2f}%; ' \
                                      'F1:  {tot_f1:.2f} ' \
                                      '{tot_predicted}\n\n'.format(tot_prec=results[tag]['precision'],
                                                                   tot_recall=results[tag]['recall'],
                                                                   tot_f1=results[tag]['f1'],
                                                                   tot_predicted=results[tag]['n_predicted_entities'])
        elif entity_of_interest is not None:
            s += '\t' + entity_of_interest + ': precision:  {tot_prec:.2f}%; ' \
                              'recall:  {tot_recall:.2f}%; ' \
                              'F1:  {tot_f1:.2f} ' \
                              '{tot_predicted}\n\n'.format(tot_prec=results[entity_of_interest]['precision'],
                                                           tot_recall=results[entity_of_interest]['recall'],
                                                           tot_f1=results[entity_of_interest]['f1'],
                                                           tot_predicted=results[entity_of_interest]['n_predicted_entities'])
        log.debug(s)
    return results


# 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 Tuple, List, Optional

import numpy as np

from deeppavlov.core.common.registry import register
from deeppavlov.core.data.utils import zero_pad
from deeppavlov.core.models.component import Component

@register('capitalization_featurizer')
class CapitalizationPreprocessor(Component):
    """
    Featurizer useful for NER task. It detects following patterns in the words:
    - no capitals
    - single capital single character
    - single capital multiple characters
    - all capitals multiple characters

    Args:
        pad_zeros: whether to pad capitalization features batch with zeros up
            to maximal length or not.

    Attributes:
        dim: dimensionality of the feature vectors, produced by the featurizer
    """
    def __init__(self, pad_zeros: bool = True, *args, **kwargs) -> None:
        self.pad_zeros = pad_zeros
        self._num_of_features = 4

    @property
    def dim(self):
        return self._num_of_features

    def __call__(self, tokens_batch, **kwargs):
        cap_batch = []
        max_batch_len = 0
        for utterance in tokens_batch:
            cap_list = []
            max_batch_len = max(max_batch_len, len(utterance))
            for token in utterance:
                cap = np.zeros(4, np.float32)
                # Check the case and produce corresponding one-hot
                if len(token) > 0:
                    if token[0].islower():
                        cap[0] = 1
                    elif len(token) == 1 and token[0].isupper():
                        cap[1] = 1
                    elif len(token) > 1 and token[0].isupper() and any(ch.islower() for ch in token):
                        cap[2] = 1
                    elif all(ch.isupper() for ch in token):
                        cap[3] = 1
                cap_list.append(cap)
            cap_batch.append(cap_list)
        if self.pad_zeros:
            return zero_pad(cap_batch)
        else:
            return cap_batch

def process_word(word: str, to_lower: bool = False,
                 append_case: Optional[str] = None) -> Tuple[str]:
    """Converts word to a tuple of symbols, optionally converts it to lowercase
    and adds capitalization label.

    Args:
        word: input word
        to_lower: whether to lowercase
        append_case: whether to add case mark
            ('<FIRST_UPPER>' for first capital and '<ALL_UPPER>' for all caps)

    Returns:
        a preprocessed word
    """
    if all(x.isupper() for x in word) and len(word) > 1:
        uppercase = "<ALL_UPPER>"
    elif word[0].isupper():
        uppercase = "<FIRST_UPPER>"
    else:
        uppercase = None
    if to_lower:
        word = word.lower()
    if word.isdigit():
        answer = ["<DIGIT>"]
    elif word.startswith("http://") or word.startswith("www."):
        answer = ["<HTTP>"]
    else:
        answer = list(word)
    if to_lower and uppercase is not None:
        if append_case == "first":
            answer = [uppercase] + answer
        elif append_case == "last":
            answer = answer + [uppercase]
    return tuple(answer)

@register('lowercase_preprocessor')
class LowercasePreprocessor(Component):
    """A callable wrapper over :func:`process_word`.
    Takes as input a batch of tokenized sentences
    and returns a batch of preprocessed sentences.
    """

    def __init__(self, to_lower: bool = True, append_case: str = "first", *args, **kwargs):
        self.to_lower = to_lower
        self.append_case = append_case

    def __call__(self, tokens_batch: List[List[str]], **kwargs) -> List[List[Tuple[str]]]:
        answer = []
        for elem in tokens_batch:
            # if isinstance(elem, str):
            #     elem = NLTKMosesTokenizer()([elem])[0]
            #     # elem = [x for x in re.split("(\w+|[,.])", elem) if x.strip() != ""]
            answer.append([process_word(x, self.to_lower, self.append_case) for x in elem])
        return answer


# 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 logging import getLogger
from typing import List

import numpy as np
import tensorflow as tf

from deeppavlov.core.common.check_gpu import check_gpu_existence

log = getLogger(__name__)

INITIALIZER = tf.orthogonal_initializer
# INITIALIZER = xavier_initializer

def stacked_cnn(units: tf.Tensor,
                n_hidden_list: List,
                filter_width=3,
                use_batch_norm=False,
                use_dilation=False,
                training_ph=None,
                add_l2_losses=False):
    """ Number of convolutional layers stacked on top of each other

    Args:
        units: a tensorflow tensor with dimensionality [None, n_tokens, n_features]
        n_hidden_list: list with number of hidden units at the ouput of each layer
        filter_width: width of the kernel in tokens
        use_batch_norm: whether to use batch normalization between layers
        use_dilation: use power of 2 dilation scheme [1, 2, 4, 8 .. ] for layers 1, 2, 3, 4 ...
        training_ph: boolean placeholder determining whether is training phase now or not.
            It is used only for batch normalization to determine whether to use
            current batch average (std) or memory stored average (std)
        add_l2_losses: whether to add l2 losses on network kernels to
                tf.GraphKeys.REGULARIZATION_LOSSES or not

    Returns:
        units: tensor at the output of the last convolutional layer
    """
    l2_reg = tf.nn.l2_loss if add_l2_losses else None
    for n_layer, n_hidden in enumerate(n_hidden_list):
        if use_dilation:
            dilation_rate = 2 ** n_layer
        else:
            dilation_rate = 1
        units = tf.layers.conv1d(units,
                                 n_hidden,
                                 filter_width,
                                 padding='same',
                                 dilation_rate=dilation_rate,
                                 kernel_initializer=INITIALIZER(),
                                 kernel_regularizer=l2_reg)
        if use_batch_norm:
            assert training_ph is not None
            units = tf.layers.batch_normalization(units, training=training_ph)
        units = tf.nn.relu(units)
    return units

def dense_convolutional_network(units: tf.Tensor,
                                n_hidden_list: List,
                                filter_width=3,
                                use_dilation=False,
                                use_batch_norm=False,
                                training_ph=None):
    """ Densely connected convolutional layers. Based on the paper:
        [Gao 17] https://arxiv.org/abs/1608.06993

        Args:
            units: a tensorflow tensor with dimensionality [None, n_tokens, n_features]
            n_hidden_list: list with number of hidden units at the ouput of each layer
            filter_width: width of the kernel in tokens
            use_batch_norm: whether to use batch normalization between layers
            use_dilation: use power of 2 dilation scheme [1, 2, 4, 8 .. ] for layers 1, 2, 3, 4 ...
            training_ph: boolean placeholder determining whether is training phase now or not.
                It is used only for batch normalization to determine whether to use
                current batch average (std) or memory stored average (std)
        Returns:
            units: tensor at the output of the last convolutional layer
                with dimensionality [None, n_tokens, n_hidden_list[-1]]
        """
    units_list = [units]
    for n_layer, n_filters in enumerate(n_hidden_list):
        total_units = tf.concat(units_list, axis=-1)
        if use_dilation:
            dilation_rate = 2 ** n_layer
        else:
            dilation_rate = 1
        units = tf.layers.conv1d(total_units,
                                 n_filters,
                                 filter_width,
                                 dilation_rate=dilation_rate,
                                 padding='same',
                                 kernel_initializer=INITIALIZER())
        if use_batch_norm:
            units = tf.layers.batch_normalization(units, training=training_ph)
        units = tf.nn.relu(units)
        units_list.append(units)
    return units

def bi_rnn(units: tf.Tensor,
           n_hidden: List,
           cell_type='gru',
           seq_lengths=None,
           trainable_initial_states=False,
           use_peepholes=False,
           name='Bi-'):
    """ Bi directional recurrent neural network. GRU or LSTM

        Args:
            units: a tensorflow tensor with dimensionality [None, n_tokens, n_features]
            n_hidden: list with number of hidden units at the ouput of each layer
            seq_lengths: length of sequences for different length sequences in batch
                can be None for maximum length as a length for every sample in the batch
            cell_type: 'lstm' or 'gru'
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            use_peepholes: whether to use peephole connections (only 'lstm' case affected)
            name: what variable_scope to use for the network parameters

        Returns:
            units: tensor at the output of the last recurrent layer
                with dimensionality [None, n_tokens, n_hidden_list[-1]]
            last_units: tensor of last hidden states for GRU and tuple
                of last hidden stated and last cell states for LSTM
                dimensionality of cell states and hidden states are
                similar and equal to [B x 2 * H], where B - batch
                size and H is number of hidden units
    """

    with tf.variable_scope(name + '_' + cell_type.upper()):
        if cell_type == 'gru':
            forward_cell = tf.nn.rnn_cell.GRUCell(n_hidden, kernel_initializer=INITIALIZER())
            backward_cell = tf.nn.rnn_cell.GRUCell(n_hidden, kernel_initializer=INITIALIZER())
            if trainable_initial_states:
                initial_state_fw = tf.tile(tf.get_variable('init_fw_h', [1, n_hidden]), (tf.shape(units)[0], 1))
                initial_state_bw = tf.tile(tf.get_variable('init_bw_h', [1, n_hidden]), (tf.shape(units)[0], 1))
            else:
                initial_state_fw = initial_state_bw = None
        elif cell_type == 'lstm':
            forward_cell = tf.nn.rnn_cell.LSTMCell(n_hidden, use_peepholes=use_peepholes, initializer=INITIALIZER())
            backward_cell = tf.nn.rnn_cell.LSTMCell(n_hidden, use_peepholes=use_peepholes, initializer=INITIALIZER())
            if trainable_initial_states:
                initial_state_fw = tf.nn.rnn_cell.LSTMStateTuple(
                    tf.tile(tf.get_variable('init_fw_c', [1, n_hidden]), (tf.shape(units)[0], 1)),
                    tf.tile(tf.get_variable('init_fw_h', [1, n_hidden]), (tf.shape(units)[0], 1)))
                initial_state_bw = tf.nn.rnn_cell.LSTMStateTuple(
                    tf.tile(tf.get_variable('init_bw_c', [1, n_hidden]), (tf.shape(units)[0], 1)),
                    tf.tile(tf.get_variable('init_bw_h', [1, n_hidden]), (tf.shape(units)[0], 1)))
            else:
                initial_state_fw = initial_state_bw = None
        else:
            raise RuntimeError('cell_type must be either "gru" or "lstm"s')
        (rnn_output_fw, rnn_output_bw), (fw, bw) = \
            tf.nn.bidirectional_dynamic_rnn(forward_cell,
                                            backward_cell,
                                            units,
                                            dtype=tf.float32,
                                            sequence_length=seq_lengths,
                                            initial_state_fw=initial_state_fw,
                                            initial_state_bw=initial_state_bw)
    kernels = [var for var in forward_cell.trainable_variables +
               backward_cell.trainable_variables if 'kernel' in var.name]
    for kernel in kernels:
        tf.add_to_collection(tf.GraphKeys.REGULARIZATION_LOSSES, tf.nn.l2_loss(kernel))
    return (rnn_output_fw, rnn_output_bw), (fw, bw)

def stacked_bi_rnn(units: tf.Tensor,
                   n_hidden_list: List,
                   cell_type='gru',
                   seq_lengths=None,
                   use_peepholes=False,
                   name='RNN_layer'):
    """ Stackted recurrent neural networks GRU or LSTM

        Args:
            units: a tensorflow tensor with dimensionality [None, n_tokens, n_features]
            n_hidden_list: list with number of hidden units at the ouput of each layer
            seq_lengths: length of sequences for different length sequences in batch
                can be None for maximum length as a length for every sample in the batch
            cell_type: 'lstm' or 'gru'
            use_peepholes: whether to use peephole connections (only 'lstm' case affected)
            name: what variable_scope to use for the network parameters
        Returns:
            units: tensor at the output of the last recurrent layer
                with dimensionality [None, n_tokens, n_hidden_list[-1]]
            last_units: tensor of last hidden states for GRU and tuple
                of last hidden stated and last cell states for LSTM
                dimensionality of cell states and hidden states are
                similar and equal to [B x 2 * H], where B - batch
                size and H is number of hidden units
    """
    for n, n_hidden in enumerate(n_hidden_list):
        with tf.variable_scope(name + '_' + str(n)):
            if cell_type == 'gru':
                forward_cell = tf.nn.rnn_cell.GRUCell(n_hidden)
                backward_cell = tf.nn.rnn_cell.GRUCell(n_hidden)
            elif cell_type == 'lstm':
                forward_cell = tf.nn.rnn_cell.LSTMCell(n_hidden, use_peepholes=use_peepholes)
                backward_cell = tf.nn.rnn_cell.LSTMCell(n_hidden, use_peepholes=use_peepholes)
            else:
                raise RuntimeError('cell_type must be either gru or lstm')

            (rnn_output_fw, rnn_output_bw), (fw, bw) = \
                tf.nn.bidirectional_dynamic_rnn(forward_cell,
                                                backward_cell,
                                                units,
                                                dtype=tf.float32,
                                                sequence_length=seq_lengths)
            units = tf.concat([rnn_output_fw, rnn_output_bw], axis=2)
            if cell_type == 'gru':
                last_units = tf.concat([fw, bw], axis=1)
            else:
                (c_fw, h_fw), (c_bw, h_bw) = fw, bw
                c = tf.concat([c_fw, c_bw], axis=1)
                h = tf.concat([h_fw, h_bw], axis=1)
                last_units = (h, c)
    return units, last_units

def u_shape(units: tf.Tensor,
            n_hidden_list: List,
            filter_width=7,
            use_batch_norm=False,
            training_ph=None):
    """ Network architecture inspired by One Hundred layer Tiramisu.
        https://arxiv.org/abs/1611.09326. U-Net like.

        Args:
            units: a tensorflow tensor with dimensionality [None, n_tokens, n_features]
            n_hidden_list: list with number of hidden units at the ouput of each layer
            filter_width: width of the kernel in tokens
            use_batch_norm: whether to use batch normalization between layers
            training_ph: boolean placeholder determining whether is training phase now or not.
                It is used only for batch normalization to determine whether to use
                current batch average (std) or memory stored average (std)
        Returns:
            units: tensor at the output of the last convolutional layer
                    with dimensionality [None, n_tokens, n_hidden_list[-1]]
    """

    # Bread Crumbs
    units_for_skip_conn = []
    conv_net_params = {'filter_width': filter_width,
                       'use_batch_norm': use_batch_norm,
                       'training_ph': training_ph}

    # Go down the rabbit hole
    for n_hidden in n_hidden_list:
        units = stacked_cnn(units, [n_hidden], **conv_net_params)
        units_for_skip_conn.append(units)
        units = tf.layers.max_pooling1d(units, pool_size=2, strides=2, padding='same')

    units = stacked_cnn(units, [n_hidden], **conv_net_params)

    # Up to the sun light
    for down_step, n_hidden in enumerate(n_hidden_list[::-1]):
        units = tf.expand_dims(units, axis=2)
        units = tf.layers.conv2d_transpose(units, n_hidden, filter_width, strides=(2, 1), padding='same')
        units = tf.squeeze(units, axis=2)

        # Skip connection
        skip_units = units_for_skip_conn[-(down_step + 1)]
        if skip_units.get_shape().as_list()[-1] != n_hidden:
            skip_units = tf.layers.dense(skip_units, n_hidden)
        units = skip_units + units

        units = stacked_cnn(units, [n_hidden], **conv_net_params)
    return units

def stacked_highway_cnn(units: tf.Tensor,
                        n_hidden_list: List,
                        filter_width=3,
                        use_batch_norm=False,
                        use_dilation=False,
                        training_ph=None):
    """ Highway convolutional network. Skip connection with gating
        mechanism.

    Args:
        units: a tensorflow tensor with dimensionality [None, n_tokens, n_features]
        n_hidden_list: list with number of hidden units at the output of each layer
        filter_width: width of the kernel in tokens
        use_batch_norm: whether to use batch normalization between layers
        use_dilation: use power of 2 dilation scheme [1, 2, 4, 8 .. ] for layers 1, 2, 3, 4 ...
        training_ph: boolean placeholder determining whether is training phase now or not.
            It is used only for batch normalization to determine whether to use
            current batch average (std) or memory stored average (std)
    Returns:
        units: tensor at the output of the last convolutional layer
                with dimensionality [None, n_tokens, n_hidden_list[-1]]
    """

    for n_layer, n_hidden in enumerate(n_hidden_list):
        input_units = units
        # Projection if needed
        if input_units.get_shape().as_list()[-1] != n_hidden:
            input_units = tf.layers.dense(input_units, n_hidden)
        if use_dilation:
            dilation_rate = 2 ** n_layer
        else:
            dilation_rate = 1
        units = tf.layers.conv1d(units,
                                 n_hidden,
                                 filter_width,
                                 padding='same',
                                 dilation_rate=dilation_rate,
                                 kernel_initializer=INITIALIZER())
        if use_batch_norm:
            units = tf.layers.batch_normalization(units, training=training_ph)
        sigmoid_gate = tf.layers.dense(input_units, 1, activation=tf.sigmoid, kernel_initializer=INITIALIZER())
        input_units = sigmoid_gate * input_units + (1 - sigmoid_gate) * units
        input_units = tf.nn.relu(input_units)
    units = input_units
    return units

def embedding_layer(token_indices=None,
                    token_embedding_matrix=None,
                    n_tokens=None,
                    token_embedding_dim=None,
                    name: str = None,
                    trainable=True):
    """ Token embedding layer. Create matrix of for token embeddings.
        Can be initialized with given matrix (for example pre-trained
        with word2ve algorithm

    Args:
        token_indices: token indices tensor of type tf.int32
        token_embedding_matrix: matrix of embeddings with dimensionality
            [n_tokens, embeddings_dimension]
        n_tokens: total number of unique tokens
        token_embedding_dim: dimensionality of embeddings, typical 100..300
        name: embedding matrix name (variable name)
        trainable: whether to set the matrix trainable or not

    Returns:
        embedded_tokens: tf tensor of size [B, T, E], where B - batch size
            T - number of tokens, E - token_embedding_dim
    """
    if token_embedding_matrix is not None:
        tok_mat = token_embedding_matrix
        if trainable:
            Warning('Matrix of embeddings is passed to the embedding_layer, '
                    'possibly there is a pre-trained embedding matrix. '
                    'Embeddings paramenters are set to Trainable!')
    else:
        tok_mat = np.random.randn(n_tokens, token_embedding_dim).astype(np.float32) / np.sqrt(token_embedding_dim)
    tok_emb_mat = tf.Variable(tok_mat, name=name, trainable=trainable)
    embedded_tokens = tf.nn.embedding_lookup(tok_emb_mat, token_indices)
    return embedded_tokens

def character_embedding_network(char_placeholder: tf.Tensor,
                                n_characters: int = None,
                                emb_mat: np.array = None,
                                char_embedding_dim: int = None,
                                filter_widths=(3, 4, 5, 7),
                                highway_on_top=False):
    """ Characters to vector. Every sequence of characters (token)
        is embedded to vector space with dimensionality char_embedding_dim
        Convolution plus max_pooling is used to obtain vector representations
        of words.

    Args:
        char_placeholder: placeholder of int32 type with dimensionality [B, T, C]
            B - batch size (can be None)
            T - Number of tokens (can be None)
            C - number of characters (can be None)
        n_characters: total number of unique characters
        emb_mat: if n_characters is not provided the emb_mat should be provided
            it is a numpy array with dimensions [V, E], where V - vocabulary size
            and E - embeddings dimension
        char_embedding_dim: dimensionality of characters embeddings
        filter_widths: array of width of kernel in convolutional embedding network
            used in parallel

    Returns:
        embeddings: tf.Tensor with dimensionality [B, T, F],
            where F is dimensionality of embeddings
    """
    if emb_mat is None:
        emb_mat = np.random.randn(n_characters, char_embedding_dim).astype(np.float32) / np.sqrt(char_embedding_dim)
    else:
        char_embedding_dim = emb_mat.shape[1]
    char_emb_var = tf.Variable(emb_mat, trainable=True)
    with tf.variable_scope('Char_Emb_Network'):
        # Character embedding layer
        c_emb = tf.nn.embedding_lookup(char_emb_var, char_placeholder)

        # Character embedding network
        conv_results_list = []
        for filter_width in filter_widths:
            conv_results_list.append(tf.layers.conv2d(c_emb,
                                                      char_embedding_dim,
                                                      (1, filter_width),
                                                      padding='same',
                                                      kernel_initializer=INITIALIZER))
        units = tf.concat(conv_results_list, axis=3)
        units = tf.reduce_max(units, axis=2)
        if highway_on_top:
            sigmoid_gate = tf.layers.dense(units,
                                           1,
                                           activation=tf.sigmoid,
                                           kernel_initializer=INITIALIZER,
                                           kernel_regularizer=tf.nn.l2_loss)
            deeper_units = tf.layers.dense(units,
                                           tf.shape(units)[-1],
                                           kernel_initializer=INITIALIZER,
                                           kernel_regularizer=tf.nn.l2_loss)
            units = sigmoid_gate * units + (1 - sigmoid_gate) * deeper_units
            units = tf.nn.relu(units)
    return units

def expand_tile(units, axis):
    """Expand and tile tensor along given axis
    Args:
        units: tf tensor with dimensions [batch_size, time_steps, n_input_features]
        axis: axis along which expand and tile. Must be 1 or 2

    """
    assert axis in (1, 2)
    n_time_steps = tf.shape(units)[1]
    repetitions = [1, 1, 1, 1]
    repetitions[axis] = n_time_steps
    return tf.tile(tf.expand_dims(units, axis), repetitions)

def additive_self_attention(units, n_hidden=None, n_output_features=None, activation=None):
    """ Computes additive self attention for time series of vectors (with batch dimension)
        the formula: score(h_i, h_j) = <v, tanh(W_1 h_i + W_2 h_j)>
        v is a learnable vector of n_hidden dimensionality,
        W_1 and W_2 are learnable [n_hidden, n_input_features] matrices

    Args:
        units: tf tensor with dimensionality [batch_size, time_steps, n_input_features]
        n_hidden: number of units in hidden representation of similarity measure
        n_output_features: number of features in output dense layer
        activation: activation at the output

    Returns:
        output: self attended tensor with dimensionality [batch_size, time_steps, n_output_features]
    """
    n_input_features = units.get_shape().as_list()[2]
    if n_hidden is None:
        n_hidden = n_input_features
    if n_output_features is None:
        n_output_features = n_input_features
    units_pairs = tf.concat([expand_tile(units, 1), expand_tile(units, 2)], 3)
    query = tf.layers.dense(units_pairs, n_hidden, activation=tf.tanh, kernel_initializer=INITIALIZER())
    attention = tf.nn.softmax(tf.layers.dense(query, 1), dim=2)
    attended_units = tf.reduce_sum(attention * expand_tile(units, 1), axis=2)
    output = tf.layers.dense(attended_units, n_output_features, activation, kernel_initializer=INITIALIZER())
    return output

def multiplicative_self_attention(units, n_hidden=None, n_output_features=None, activation=None):
    """ Computes multiplicative self attention for time series of vectors (with batch dimension)
        the formula: score(h_i, h_j) = <W_1 h_i,  W_2 h_j>,  W_1 and W_2 are learnable matrices
        with dimensionality [n_hidden, n_input_features], where <a, b> stands for a and b
        dot product

    Args:
        units: tf tensor with dimensionality [batch_size, time_steps, n_input_features]
        n_hidden: number of units in hidden representation of similarity measure
        n_output_features: number of features in output dense layer
        activation: activation at the output

    Returns:
        output: self attended tensor with dimensionality [batch_size, time_steps, n_output_features]
    """
    n_input_features = units.get_shape().as_list()[2]
    if n_hidden is None:
        n_hidden = n_input_features
    if n_output_features is None:
        n_output_features = n_input_features
    queries = tf.layers.dense(expand_tile(units, 1), n_hidden, kernel_initializer=INITIALIZER())
    keys = tf.layers.dense(expand_tile(units, 2), n_hidden, kernel_initializer=INITIALIZER())
    scores = tf.reduce_sum(queries * keys, axis=3, keep_dims=True)
    attention = tf.nn.softmax(scores, dim=2)
    attended_units = tf.reduce_sum(attention * expand_tile(units, 1), axis=2)
    output = tf.layers.dense(attended_units, n_output_features, activation, kernel_initializer=INITIALIZER())
    return output

def cudnn_gru(units, n_hidden, n_layers=1, trainable_initial_states=False,
              seq_lengths=None, input_initial_h=None, name='cudnn_gru', reuse=False):
    """ Fast CuDNN GRU implementation

    Args:
        units: tf.Tensor with dimensions [B x T x F], where
            B - batch size
            T - number of tokens
            F - features

        n_hidden: dimensionality of hidden state
        trainable_initial_states: whether to create a special trainable variable
            to initialize the hidden states of the network or use just zeros
        seq_lengths: tensor of sequence lengths with dimension [B]
        n_layers: number of layers
        input_initial_h: initial hidden state, tensor
        name: name of the variable scope to use
        reuse:whether to reuse already initialized variable

    Returns:
        h - all hidden states along T dimension,
            tf.Tensor with dimensionality [B x T x F]
        h_last - last hidden state, tf.Tensor with dimensionality [B x H]
    """
    with tf.variable_scope(name, reuse=reuse):
        gru = tf.contrib.cudnn_rnn.CudnnGRU(num_layers=n_layers,
                                            num_units=n_hidden)

        if trainable_initial_states:
            init_h = tf.get_variable('init_h', [n_layers, 1, n_hidden])
            init_h = tf.tile(init_h, (1, tf.shape(units)[0], 1))
        else:
            init_h = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = input_initial_h or init_h

        h, h_last = gru(tf.transpose(units, (1, 0, 2)), (initial_h, ))
        h = tf.transpose(h, (1, 0, 2))
        h_last = tf.squeeze(h_last, axis=0)[-1]  # extract last layer state

        # Extract last states if they are provided
        if seq_lengths is not None:
            indices = tf.stack([tf.range(tf.shape(h)[0]), seq_lengths-1], axis=1)
            h_last = tf.gather_nd(h, indices)

        return h, h_last

def cudnn_compatible_gru(units, n_hidden, n_layers=1, trainable_initial_states=False,
                         seq_lengths=None, input_initial_h=None, name='cudnn_gru', reuse=False):
    """ CuDNN Compatible GRU implementation.
        It should be used to load models saved with CudnnGRUCell to run on CPU.

        Args:
            units: tf.Tensor with dimensions [B x T x F], where
                B - batch size
                T - number of tokens
                F - features

            n_hidden: dimensionality of hidden state
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            seq_lengths: tensor of sequence lengths with dimension [B]
            n_layers: number of layers
            input_initial_h: initial hidden state, tensor
            name: name of the variable scope to use
            reuse:whether to reuse already initialized variable

        Returns:
            h - all hidden states along T dimension,
                tf.Tensor with dimensionality [B x T x F]
            h_last - last hidden state, tf.Tensor with dimensionality [B x H]
        """
    with tf.variable_scope(name, reuse=reuse):

        if trainable_initial_states:
            init_h = tf.get_variable('init_h', [n_layers, 1, n_hidden])
            init_h = tf.tile(init_h, (1, tf.shape(units)[0], 1))
        else:
            init_h = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = input_initial_h or init_h

        with tf.variable_scope('cudnn_gru', reuse=reuse):
            def single_cell(): return tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell(n_hidden)
            cell = tf.nn.rnn_cell.MultiRNNCell([single_cell() for _ in range(n_layers)])

            units = tf.transpose(units, (1, 0, 2))

            h, h_last = tf.nn.dynamic_rnn(cell=cell, inputs=units, time_major=True,
                                          initial_state=tuple(tf.unstack(initial_h, axis=0)))
            h = tf.transpose(h, (1, 0, 2))

            h_last = h_last[-1]  # h_last is tuple: n_layers x batch_size x n_hidden

            # Extract last states if they are provided
            if seq_lengths is not None:
                indices = tf.stack([tf.range(tf.shape(h)[0]), seq_lengths-1], axis=1)
                h_last = tf.gather_nd(h, indices)

            return h, h_last

def cudnn_gru_wrapper(units, n_hidden, n_layers=1, trainable_initial_states=False,
                      seq_lengths=None, input_initial_h=None, name='cudnn_gru', reuse=False):

    if check_gpu_existence():
        return cudnn_gru(units, n_hidden, n_layers, trainable_initial_states,
                         seq_lengths, input_initial_h, name, reuse)

    log.info('\nWarning! tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell is used. '
             'It is okay for inference mode, but '
             'if you train your model with this cell it could NOT be used with '
             'tf.contrib.cudnn_rnn.CudnnGRUCell later. '
             )

    return cudnn_compatible_gru(units, n_hidden, n_layers, trainable_initial_states,
                                seq_lengths, input_initial_h, name, reuse)

def cudnn_lstm(units, n_hidden, n_layers=1, trainable_initial_states=None, seq_lengths=None, initial_h=None,
               initial_c=None, name='cudnn_lstm', reuse=False):
    """ Fast CuDNN LSTM implementation

        Args:
            units: tf.Tensor with dimensions [B x T x F], where
                B - batch size
                T - number of tokens
                F - features
            n_hidden: dimensionality of hidden state
            n_layers: number of layers
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            seq_lengths: tensor of sequence lengths with dimension [B]
            initial_h: optional initial hidden state, masks trainable_initial_states
                if provided
            initial_c: optional initial cell state, masks trainable_initial_states
                if provided
            name: name of the variable scope to use
            reuse:whether to reuse already initialized variable

        Returns:
            h - all hidden states along T dimension,
                tf.Tensor with dimensionality [B x T x F]
            h_last - last hidden state, tf.Tensor with dimensionality [B x H]
                where H - number of hidden units
            c_last - last cell state, tf.Tensor with dimensionality [B x H]
                where H - number of hidden units
        """
    with tf.variable_scope(name, reuse=reuse):
        lstm = tf.contrib.cudnn_rnn.CudnnLSTM(num_layers=n_layers,
                                              num_units=n_hidden)
        if trainable_initial_states:
            init_h = tf.get_variable('init_h', [n_layers, 1, n_hidden])
            init_h = tf.tile(init_h, (1, tf.shape(units)[0], 1))
            init_c = tf.get_variable('init_c', [n_layers, 1, n_hidden])
            init_c = tf.tile(init_c, (1, tf.shape(units)[0], 1))
        else:
            init_h = init_c = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = initial_h or init_h
        initial_c = initial_c or init_c

        h, (h_last, c_last) = lstm(tf.transpose(units, (1, 0, 2)), (initial_h, initial_c))
        h = tf.transpose(h, (1, 0, 2))
        h_last = h_last[-1]
        c_last = c_last[-1]

        # Extract last states if they are provided
        if seq_lengths is not None:
            indices = tf.stack([tf.range(tf.shape(h)[0]), seq_lengths-1], axis=1)
            h_last = tf.gather_nd(h, indices)

        return h, (h_last, c_last)

def cudnn_compatible_lstm(units, n_hidden, n_layers=1, trainable_initial_states=None, seq_lengths=None, initial_h=None,
                          initial_c=None, name='cudnn_lstm', reuse=False):
    """ CuDNN Compatible LSTM implementation.
        It should be used to load models saved with CudnnLSTMCell to run on CPU.

        Args:
            units: tf.Tensor with dimensions [B x T x F], where
                B - batch size
                T - number of tokens
                F - features
            n_hidden: dimensionality of hidden state
            n_layers: number of layers
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            seq_lengths: tensor of sequence lengths with dimension [B]
            initial_h: optional initial hidden state, masks trainable_initial_states
                if provided
            initial_c: optional initial cell state, masks trainable_initial_states
                if provided
            name: name of the variable scope to use
            reuse:whether to reuse already initialized variable

        Returns:
            h - all hidden states along T dimension,
                tf.Tensor with dimensionality [B x T x F]
            h_last - last hidden state, tf.Tensor with dimensionality [B x H]
                where H - number of hidden units
            c_last - last cell state, tf.Tensor with dimensionality [B x H]
                where H - number of hidden units
        """

    with tf.variable_scope(name, reuse=reuse):
        if trainable_initial_states:
            init_h = tf.get_variable('init_h', [n_layers, 1, n_hidden])
            init_h = tf.tile(init_h, (1, tf.shape(units)[0], 1))
            init_c = tf.get_variable('init_c', [n_layers, 1, n_hidden])
            init_c = tf.tile(init_c, (1, tf.shape(units)[0], 1))
        else:
            init_h = init_c = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = initial_h or init_h
        initial_c = initial_c or init_c

        with tf.variable_scope('cudnn_lstm', reuse=reuse):
            def single_cell(): return tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell(n_hidden)

            cell = tf.nn.rnn_cell.MultiRNNCell([single_cell() for _ in range(n_layers)])

            units = tf.transpose(units, (1, 0, 2))

            init = tuple([tf.nn.rnn_cell.LSTMStateTuple(ic, ih) for ih, ic in
                          zip(tf.unstack(initial_h, axis=0), tf.unstack(initial_c, axis=0))])

            h, state = tf.nn.dynamic_rnn(cell=cell, inputs=units, time_major=True, initial_state=init)

            h = tf.transpose(h, (1, 0, 2))
            h_last = state[-1].h
            c_last = state[-1].c

            # Extract last states if they are provided
            if seq_lengths is not None:
                indices = tf.stack([tf.range(tf.shape(h)[0]), seq_lengths-1], axis=1)
                h_last = tf.gather_nd(h, indices)

            return h, (h_last, c_last)

def cudnn_lstm_wrapper(units, n_hidden, n_layers=1, trainable_initial_states=None, seq_lengths=None, initial_h=None,
                       initial_c=None, name='cudnn_lstm', reuse=False):

    if check_gpu_existence():
        return cudnn_lstm(units, n_hidden, n_layers, trainable_initial_states,
                          seq_lengths, initial_h, initial_c, name, reuse)

    log.info('\nWarning! tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell is used. '
             'It is okay for inference mode, but '
             'if you train your model with this cell it could NOT be used with '
             'tf.contrib.cudnn_rnn.CudnnLSTMCell later. '
             )

    return cudnn_compatible_lstm(units, n_hidden, n_layers, trainable_initial_states,
                                 seq_lengths, initial_h, initial_c, name, reuse)

def cudnn_bi_gru(units,
                 n_hidden,
                 seq_lengths=None,
                 n_layers=1,
                 trainable_initial_states=False,
                 name='cudnn_bi_gru',
                 reuse=False):
    """ Fast CuDNN Bi-GRU implementation

    Args:
        units: tf.Tensor with dimensions [B x T x F], where
            B - batch size
            T - number of tokens
            F - features
        n_hidden: dimensionality of hidden state
        seq_lengths: number of tokens in each sample in the batch
        n_layers: number of layers
        trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
        name: name of the variable scope to use
        reuse:whether to reuse already initialized variable

    Returns:
        h - all hidden states along T dimension,
            tf.Tensor with dimensionality [B x T x F]
        h_last - last hidden state, tf.Tensor with dimensionality [B x H * 2]
            where H - number of hidden units
    """

    with tf.variable_scope(name, reuse=reuse):
        if seq_lengths is None:
            seq_lengths = tf.ones([tf.shape(units)[0]], dtype=tf.int32) * tf.shape(units)[1]
        with tf.variable_scope('Forward'):
            h_fw, h_last_fw = cudnn_gru_wrapper(units,
                                                n_hidden,
                                                n_layers=n_layers,
                                                trainable_initial_states=trainable_initial_states,
                                                seq_lengths=seq_lengths,
                                                reuse=reuse)

        with tf.variable_scope('Backward'):
            reversed_units = tf.reverse_sequence(units, seq_lengths=seq_lengths, seq_dim=1, batch_dim=0)
            h_bw, h_last_bw = cudnn_gru_wrapper(reversed_units,
                                                n_hidden,
                                                n_layers=n_layers,
                                                trainable_initial_states=trainable_initial_states,
                                                seq_lengths=seq_lengths,
                                                reuse=reuse)
            h_bw = tf.reverse_sequence(h_bw, seq_lengths=seq_lengths, seq_dim=1, batch_dim=0)

    return (h_fw, h_bw), (h_last_fw, h_last_bw)

def cudnn_bi_lstm(units,
                  n_hidden,
                  seq_lengths=None,
                  n_layers=1,
                  trainable_initial_states=False,
                  name='cudnn_bi_gru',
                  reuse=False):
    """ Fast CuDNN Bi-LSTM implementation

    Args:
        units: tf.Tensor with dimensions [B x T x F], where
            B - batch size
            T - number of tokens
            F - features
        n_hidden: dimensionality of hidden state
        seq_lengths: number of tokens in each sample in the batch
        n_layers: number of layers
        trainable_initial_states: whether to create a special trainable variable
            to initialize the hidden states of the network or use just zeros
        name: name of the variable scope to use
        reuse:whether to reuse already initialized variable

    Returns:
        h - all hidden states along T dimension,
            tf.Tensor with dimensionality [B x T x F]
        h_last - last hidden state, tf.Tensor with dimensionality [B x H * 2]
            where H - number of hidden units
        c_last - last cell state, tf.Tensor with dimensionality [B x H * 2]
            where H - number of hidden units
        """
    with tf.variable_scope(name, reuse=reuse):
        if seq_lengths is None:
            seq_lengths = tf.ones([tf.shape(units)[0]], dtype=tf.int32) * tf.shape(units)[1]
        with tf.variable_scope('Forward'):
            h_fw, (h_fw_last, c_fw_last) = cudnn_lstm_wrapper(units,
                                                              n_hidden,
                                                              n_layers=n_layers,
                                                              trainable_initial_states=trainable_initial_states,
                                                              seq_lengths=seq_lengths)

        with tf.variable_scope('Backward'):
            reversed_units = tf.reverse_sequence(units, seq_lengths=seq_lengths, seq_dim=1, batch_dim=0)
            h_bw, (h_bw_last, c_bw_last) = cudnn_lstm_wrapper(reversed_units,
                                                              n_hidden,
                                                              n_layers=n_layers,
                                                              trainable_initial_states=trainable_initial_states,
                                                              seq_lengths=seq_lengths)

            h_bw = tf.reverse_sequence(h_bw, seq_lengths=seq_lengths, seq_dim=1, batch_dim=0)
        return (h_fw, h_bw), ((h_fw_last, c_fw_last), (h_bw_last, c_bw_last))

def cudnn_stacked_bi_gru(units,
                         n_hidden,
                         seq_lengths=None,
                         n_stacks=2,
                         keep_prob=1.0,
                         concat_stacked_outputs=False,
                         trainable_initial_states=False,
                         name='cudnn_stacked_bi_gru',
                         reuse=False):
    """ Fast CuDNN Stacked Bi-GRU implementation

    Args:
        units: tf.Tensor with dimensions [B x T x F], where
            B - batch size
            T - number of tokens
            F - features
        n_hidden: dimensionality of hidden state
        seq_lengths: number of tokens in each sample in the batch
        n_stacks: number of stacked Bi-GRU
        keep_prob: dropout keep_prob between Bi-GRUs (intra-layer dropout)
        concat_stacked_outputs: return last Bi-GRU output or concat outputs from every Bi-GRU,
        trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
        name: name of the variable scope to use
        reuse: whether to reuse already initialized variable

    Returns:
        h - all hidden states along T dimension,
            tf.Tensor with dimensionality [B x T x ((n_hidden * 2) * n_stacks)]
    """
    if seq_lengths is None:
        seq_lengths = tf.ones([tf.shape(units)[0]], dtype=tf.int32) * tf.shape(units)[1]

    outputs = [units]

    with tf.variable_scope(name, reuse=reuse):
        for n in range(n_stacks):

            if n == 0:
                inputs = outputs[-1]
            else:
                inputs = variational_dropout(outputs[-1], keep_prob=keep_prob)

            (h_fw, h_bw), _ = cudnn_bi_gru(inputs, n_hidden, seq_lengths,
                                           n_layers=1,
                                           trainable_initial_states=trainable_initial_states,
                                           name='{}_cudnn_bi_gru'.format(n),
                                           reuse=reuse)

            outputs.append(tf.concat([h_fw, h_bw], axis=2))

    if concat_stacked_outputs:
        return tf.concat(outputs[1:], axis=2)

    return outputs[-1]

def variational_dropout(units, keep_prob, fixed_mask_dims=(1,)):
    """ Dropout with the same drop mask for all fixed_mask_dims

    Args:
        units: a tensor, usually with shapes [B x T x F], where
            B - batch size
            T - tokens dimension
            F - feature dimension
        keep_prob: keep probability
        fixed_mask_dims: in these dimensions the mask will be the same

    Returns:
        dropped units tensor
    """
    units_shape = tf.shape(units)
    noise_shape = [units_shape[n] for n in range(len(units.shape))]
    for dim in fixed_mask_dims:
        noise_shape[dim] = 1
    return tf.nn.dropout(units, keep_prob, noise_shape)


# 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 logging import getLogger
from typing import List, Union, Tuple, Iterable

import keras.layers as kl
import keras.optimizers as ko
import keras.regularizers as kreg
from keras import Model

from deeppavlov.core.common.registry import register
from deeppavlov.core.data.vocab import DefaultVocabulary
from deeppavlov.core.models.keras_model import KerasWrapper
from .cells import Highway
from .common_tagger import *

log = getLogger(__name__)

MAX_WORD_LENGTH = 30

class CharacterTagger:

    """A class for character-based neural morphological tagger

    Parameters:
        symbols: character vocabulary
        tags: morphological tags vocabulary
        word_rnn: the type of character-level network (only `cnn` implemented)
        char_embeddings_size: the size of character embeddings
        char_conv_layers: the number of convolutional layers on character level
        char_window_size: the width of convolutional filter (filters).
            It can be a list if several parallel filters are applied, for example, [2, 3, 4, 5].
        char_filters: the number of convolutional filters for each window width.
            It can be a number, a list (when there are several windows of different width
            on a single convolution layer), a list of lists, if there
            are more than 1 convolution layers, or **None**.
            If **None**, a layer with width **width** contains
            min(**char_filter_multiple** * **width**, 200) filters.

        char_filter_multiple: the ratio between filters number and window width
        char_highway_layers: the number of highway layers on character level
        conv_dropout: the ratio of dropout between convolutional layers
        highway_dropout: the ratio of dropout between highway layers,
        intermediate_dropout: the ratio of dropout between convolutional
            and highway layers on character level
        lstm_dropout: dropout ratio in word-level LSTM
        word_vectorizers: list of parameters for additional word-level vectorizers,
            for each vectorizer it stores a pair of vectorizer dimension and
            the dimension of the corresponding word embedding
        word_lstm_layers: the number of word-level LSTM layers
        word_lstm_units: hidden dimensions of word-level LSTMs
        word_dropout: the ratio of dropout before word level (it is applied to word embeddings)
        regularizer: l2 regularization parameter
        verbose: the level of verbosity
    """
    def __init__(self,
                 symbols: DefaultVocabulary,
                 tags: DefaultVocabulary,
                 word_rnn: str = "cnn",
                 char_embeddings_size: int = 16,
                 char_conv_layers: int = 1,
                 char_window_size: Union[int, List[int]] = 5,
                 char_filters: Union[int, List[int]] = None,
                 char_filter_multiple: int = 25,
                 char_highway_layers: int = 1,
                 conv_dropout: float = 0.0,
                 highway_dropout: float = 0.0,
                 intermediate_dropout: float = 0.0,
                 lstm_dropout: float = 0.0,
                 word_vectorizers: List[Tuple[int, int]] = None,
                 word_lstm_layers: int = 1,
                 word_lstm_units: Union[int, List[int]] = 128,
                 word_dropout: float = 0.0,
                 regularizer: float = None,
                 verbose: int = 1):
        self.symbols = symbols
        self.tags = tags
        self.word_rnn = word_rnn
        self.char_embeddings_size = char_embeddings_size
        self.char_conv_layers = char_conv_layers
        self.char_window_size = char_window_size
        self.char_filters = char_filters
        self.char_filter_multiple = char_filter_multiple
        self.char_highway_layers = char_highway_layers
        self.conv_dropout = conv_dropout
        self.highway_dropout = highway_dropout
        self.intermediate_dropout = intermediate_dropout
        self.lstm_dropout = lstm_dropout
        self.word_dropout = word_dropout
        self.word_vectorizers = word_vectorizers  # a list of additional vectorizer dimensions
        self.word_lstm_layers = word_lstm_layers
        self.word_lstm_units = word_lstm_units
        self.regularizer = regularizer
        self.verbose = verbose
        self._initialize()
        self.build()

    def _initialize(self):
        if isinstance(self.char_window_size, int):
            self.char_window_size = [self.char_window_size]
        if self.char_filters is None or isinstance(self.char_filters, int):
            self.char_filters = [self.char_filters] * len(self.char_window_size)
        if len(self.char_window_size) != len(self.char_filters):
            raise ValueError("There should be the same number of window sizes and filter sizes")
        if isinstance(self.word_lstm_units, int):
            self.word_lstm_units = [self.word_lstm_units] * self.word_lstm_layers
        if len(self.word_lstm_units) != self.word_lstm_layers:
            raise ValueError("There should be the same number of lstm layer units and lstm layers")
        if self.word_vectorizers is None:
            self.word_vectorizers = []
        if self.regularizer is not None:
            self.regularizer = kreg.l2(self.regularizer)
        if self.verbose > 0:
            log.info("{} symbols, {} tags in CharacterTagger".format(self.symbols_number_, self.tags_number_))

    @property
    def symbols_number_(self) -> int:
        """Character vocabulary size
        """
        return len(self.symbols)

    @property
    def tags_number_(self) -> int:
        """Tag vocabulary size
        """
        return len(self.tags)

    def build(self):
        """Builds the network using Keras.
        """
        word_inputs = kl.Input(shape=(None, MAX_WORD_LENGTH+2), dtype="int32")
        inputs = [word_inputs]
        word_outputs = self._build_word_cnn(word_inputs)
        if len(self.word_vectorizers) > 0:
            additional_word_inputs = [kl.Input(shape=(None, input_dim), dtype="float32")
                                      for input_dim, dense_dim in self.word_vectorizers]
            inputs.extend(additional_word_inputs)
            additional_word_embeddings = [kl.Dense(dense_dim)(additional_word_inputs[i])
                                          for i, (_, dense_dim) in enumerate(self.word_vectorizers)]
            word_outputs = kl.Concatenate()([word_outputs] + additional_word_embeddings)
        outputs, lstm_outputs = self._build_basic_network(word_outputs)
        compile_args = {"optimizer": ko.nadam(lr=0.002, clipnorm=5.0),
                        "loss": "categorical_crossentropy", "metrics": ["accuracy"]}
        self.model_ = Model(inputs, outputs)
        self.model_.compile(**compile_args)
        if self.verbose > 0:
            self.model_.summary(print_fn=log.info)
        return self

    def _build_word_cnn(self, inputs):
        """Builds word-level network
        """
        inputs = kl.Lambda(kb.one_hot, arguments={"num_classes": self.symbols_number_},
                           output_shape=lambda x: tuple(x) + (self.symbols_number_,))(inputs)
        char_embeddings = kl.Dense(self.char_embeddings_size, use_bias=False)(inputs)
        conv_outputs = []
        self.char_output_dim_ = 0
        for window_size, filters_number in zip(self.char_window_size, self.char_filters):
            curr_output = char_embeddings
            curr_filters_number = (min(self.char_filter_multiple * window_size, 200)
                                   if filters_number is None else filters_number)
            for _ in range(self.char_conv_layers - 1):
                curr_output = kl.Conv2D(curr_filters_number, (1, window_size),
                                        padding="same", activation="relu",
                                        data_format="channels_last")(curr_output)
                if self.conv_dropout > 0.0:
                    curr_output = kl.Dropout(self.conv_dropout)(curr_output)
            curr_output = kl.Conv2D(curr_filters_number, (1, window_size),
                                    padding="same", activation="relu",
                                    data_format="channels_last")(curr_output)
            conv_outputs.append(curr_output)
            self.char_output_dim_ += curr_filters_number
        if len(conv_outputs) > 1:
            conv_output = kl.Concatenate(axis=-1)(conv_outputs)
        else:
            conv_output = conv_outputs[0]
        highway_input = kl.Lambda(kb.max, arguments={"axis": -2})(conv_output)
        if self.intermediate_dropout > 0.0:
            highway_input = kl.Dropout(self.intermediate_dropout)(highway_input)
        for i in range(self.char_highway_layers - 1):
            highway_input = Highway(activation="relu")(highway_input)
            if self.highway_dropout > 0.0:
                highway_input = kl.Dropout(self.highway_dropout)(highway_input)
        highway_output = Highway(activation="relu")(highway_input)
        return highway_output

    def _build_basic_network(self, word_outputs):
        """
        Creates the basic network architecture,
        transforming word embeddings to intermediate outputs
        """
        if self.word_dropout > 0.0:
            lstm_outputs = kl.Dropout(self.word_dropout)(word_outputs)
        else:
            lstm_outputs = word_outputs
        for j in range(self.word_lstm_layers-1):
            lstm_outputs = kl.Bidirectional(
                kl.LSTM(self.word_lstm_units[j], return_sequences=True,
                        dropout=self.lstm_dropout))(lstm_outputs)
        lstm_outputs = kl.Bidirectional(
                kl.LSTM(self.word_lstm_units[-1], return_sequences=True,
                        dropout=self.lstm_dropout))(lstm_outputs)
        pre_outputs = kl.TimeDistributed(
                kl.Dense(self.tags_number_, activation="softmax",
                         activity_regularizer=self.regularizer),
                name="p")(lstm_outputs)
        return pre_outputs, lstm_outputs

    def _transform_batch(self, data, labels=None, transform_to_one_hot=True):
        data, additional_data = data[0], data[1:]
        L = max(len(x) for x in data)
        X = np.array([self._make_sent_vector(x, L) for x in data])
        X = [X] + [np.array(x) for x in additional_data]
        if labels is not None:
            Y = np.array([self._make_tags_vector(y, L) for y in labels])
            if transform_to_one_hot:
                Y = to_one_hot(Y, len(self.tags))
            return X, Y
        else:
            return X

    def train_on_batch(self, data: List[Iterable], labels: Iterable[list]) -> None:
        """Trains model on a single batch

        Args:
            data: a batch of word sequences
            labels: a batch of correct tag sequences
        Returns:
            the trained model
        """
        X, Y = self._transform_batch(data, labels)
        self.model_.train_on_batch(X, Y)

    def predict_on_batch(self, data: Union[list, tuple],
                         return_indexes: bool = False) -> List[List[str]]:
        """
        Makes predictions on a single batch

        Args:
            data: a batch of word sequences together with additional inputs
            return_indexes: whether to return tag indexes in vocabulary or tags themselves

        Returns:
            a batch of label sequences
        """
        X = self._transform_batch(data)
        objects_number, lengths = len(X[0]), [len(elem) for elem in data[0]]
        Y = self.model_.predict_on_batch(X)
        labels = np.argmax(Y, axis=-1)
        answer: List[List[str]] = [None] * objects_number
        for i, (elem, length) in enumerate(zip(labels, lengths)):
            elem = elem[:length]
            answer[i] = elem if return_indexes else self.tags.idxs2toks(elem)
        return answer

    def _make_sent_vector(self, sent: List, bucket_length: int =None) -> np.ndarray:
        """Transforms a sentence to Numpy array, which will be the network input.

        Args:
            sent: input sentence
            bucket_length: the width of the bucket

        Returns:
            A 3d array, answer[i][j][k] contains the index of k-th letter
            in j-th word of i-th input sentence.
        """
        bucket_length = bucket_length or len(sent)
        answer = np.zeros(shape=(bucket_length, MAX_WORD_LENGTH+2), dtype=np.int32)
        for i, word in enumerate(sent):
            answer[i, 0] = self.tags.tok2idx("BEGIN")
            m = min(len(word), MAX_WORD_LENGTH)
            for j, x in enumerate(word[-m:]):
                answer[i, j+1] = self.symbols.tok2idx(x)
            answer[i, m+1] = self.tags.tok2idx("END")
            answer[i, m+2:] = self.tags.tok2idx("PAD")
        return answer

    def _make_tags_vector(self, tags, bucket_length=None) -> np.ndarray:
        """Transforms a sentence of tags to Numpy array, which will be the network target.

        Args:
            tags: input sentence of tags
            bucket_length: the width of the bucket

        Returns:
            A 2d array, answer[i][j] contains the index of j-th tag in i-th input sentence.
        """
        bucket_length = bucket_length or len(tags)
        answer = np.zeros(shape=(bucket_length,), dtype=np.int32)
        for i, tag in enumerate(tags):
            answer[i] = self.tags.tok2idx(tag)
        return answer

    def save(self, outfile) -> None:
        """Saves model weights to a file

        Args:
            outfile: file with model weights (other model components should be given in config)
        """
        self.model_.save_weights(outfile)

    def load(self, infile) -> None:
        """Loads model weights from a file

        Args:
            infile: file to load model weights from
        """
        self.model_.load_weights(infile)

@register("morpho_tagger")
class MorphoTagger(KerasWrapper):
    """
    A wrapper over :class:`CharacterTagger`.
    It is inherited from :class:`~deeppavlov.core.keras_model.KerasWrapper`.
    It accepts initialization parameters of :class:`CharacterTagger`
    """
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(CharacterTagger, *args, **kwargs)


# 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 itertools
from typing import List, Tuple, Any

from nltk.translate.bleu_score import corpus_bleu, sentence_bleu, SmoothingFunction, brevity_penalty, closest_ref_length

from deeppavlov.core.common.metrics_registry import register_metric
from deeppavlov.metrics.google_bleu import compute_bleu

SMOOTH = SmoothingFunction()

@register_metric('bleu_advanced')
def bleu_advanced(y_true: List[Any], y_predicted: List[Any],
                  weights: Tuple=(1,), smoothing_function=SMOOTH.method1,
                  auto_reweigh=False, penalty=True) -> float:
    """Calculate BLEU score

    Parameters:
        y_true: list of reference tokens
        y_predicted: list of query tokens
        weights: n-gram weights
        smoothing_function: SmoothingFunction
        auto_reweigh: Option to re-normalize the weights uniformly
        penalty: either enable brevity penalty or not

    Return:
        BLEU score
    """

    bleu_measure = sentence_bleu([y_true], y_predicted, weights, smoothing_function, auto_reweigh)

    hyp_len = len(y_predicted)
    hyp_lengths = hyp_len
    ref_lengths = closest_ref_length([y_true], hyp_len)

    bpenalty = brevity_penalty(ref_lengths, hyp_lengths)

    if penalty is True or bpenalty == 0:
        return bleu_measure

    return bleu_measure/bpenalty

@register_metric('bleu')
def bleu(y_true, y_predicted):
    return corpus_bleu([[y_t.lower().split()] for y_t in y_true],
                       [y_p.lower().split() for y_p in y_predicted])

@register_metric('google_bleu')
def google_bleu(y_true, y_predicted):
    return compute_bleu(([y_t.lower().split()] for y_t in y_true),
                        (y_p.lower().split() for y_p in y_predicted))[0]

@register_metric('per_item_bleu')
def per_item_bleu(y_true, y_predicted):
    y_predicted = itertools.chain(*y_predicted)
    return corpus_bleu([[y_t.lower().split()] for y_t in y_true],
                       [y_p.lower().split() for y_p in y_predicted])

@register_metric('per_item_dialog_bleu')
def per_item_dialog_bleu(y_true, y_predicted):
    y_true = (y['text'] for dialog in y_true for y in dialog)
    return corpus_bleu([[y_t.lower().split()] for y_t in y_true],
                       [y_p.lower().split() for y_p in y_predicted])