# -*- coding:utf-8 -*-
# Author: hankcs
# Date: 2019-09-14 20:33
from abc import ABC
from typing import Union, Any, Tuple, Iterable

import tensorflow as tf
from hanlp.components.taggers.transformers.transformer_transform_tf import TransformerTransform

from hanlp.common.transform_tf import Transform

from hanlp.common.keras_component import KerasComponent
from hanlp.components.taggers.ngram_conv.ngram_conv_tagger import NgramConvTaggerTF
from hanlp.components.taggers.rnn_tagger_tf import RNNTaggerTF
from hanlp.components.taggers.transformers.transformer_tagger_tf import TransformerTaggerTF
from hanlp.metrics.chunking.sequence_labeling import iobes_to_span
from hanlp_common.util import merge_locals_kwargs


class IOBES_NamedEntityRecognizer(KerasComponent, ABC):

    def predict_batch(self, batch, inputs=None):
        for words, tags in zip(inputs, super().predict_batch(batch, inputs)):
            yield from iobes_to_span(words, tags)


class IOBES_Transform(Transform):

    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None,
                     batch=None) -> Iterable:
        for words, tags in zip(inputs, super().Y_to_outputs(Y, gold, inputs=inputs, X=X, batch=batch)):
            yield from iobes_to_span(words, tags)


class RNNNamedEntityRecognizerTF(RNNTaggerTF, IOBES_NamedEntityRecognizer):

    def fit(self, trn_data: str, dev_data: str = None, save_dir: str = None, embeddings=100, embedding_trainable=False,
            rnn_input_dropout=0.2, rnn_units=100, rnn_output_dropout=0.2, epochs=20, logger=None,
            loss: Union[tf.keras.losses.Loss, str] = None,
            optimizer: Union[str, tf.keras.optimizers.Optimizer] = 'adam', metrics='f1', batch_size=32,
            dev_batch_size=32, lr_decay_per_epoch=None,
            run_eagerly=False,
            verbose=True, **kwargs):
        # assert kwargs.get('run_eagerly', True), 'This component can only run eagerly'
        # kwargs['run_eagerly'] = True
        return super().fit(**merge_locals_kwargs(locals(), kwargs))

    def build_loss(self, loss, **kwargs):
        if not loss:
            loss = tf.keras.losses.SparseCategoricalCrossentropy(
                reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE,
                from_logits=True)
        return super().build_loss(loss, **kwargs)


class NgramConvNamedEntityRecognizerTF(NgramConvTaggerTF, IOBES_NamedEntityRecognizer):

    def fit(self, trn_data: Any, dev_data: Any, save_dir: str, word_embed: Union[str, int, dict] = 200,
            ngram_embed: Union[str, int, dict] = 50, embedding_trainable=True, window_size=4, kernel_size=3,
            filters=(200, 200, 200, 200, 200), dropout_embed=0.2, dropout_hidden=0.2, weight_norm=True,
            loss: Union[tf.keras.losses.Loss, str] = None,
            optimizer: Union[str, tf.keras.optimizers.Optimizer] = 'adam', metrics='f1', batch_size=100,
            epochs=100, logger=None, verbose=True, **kwargs):
        return super().fit(trn_data, dev_data, save_dir, word_embed, ngram_embed, embedding_trainable, window_size,
                           kernel_size, filters, dropout_embed, dropout_hidden, weight_norm, loss, optimizer, metrics,
                           batch_size, epochs, logger, verbose, **kwargs)


class IOBES_TransformerTransform(IOBES_Transform, TransformerTransform):
    pass


class TransformerNamedEntityRecognizerTF(TransformerTaggerTF):

    def __init__(self, transform: TransformerTransform = None) -> None:
        if not transform:
            transform = IOBES_TransformerTransform()
        super().__init__(transform)

    def fit(self, trn_data, dev_data, save_dir, transformer, optimizer='adamw', learning_rate=5e-5, weight_decay_rate=0,
            epsilon=1e-8, clipnorm=1.0, warmup_steps_ratio=0, use_amp=False, max_seq_length=128, batch_size=32,
            epochs=3, metrics='f1', run_eagerly=False, logger=None, verbose=True, **kwargs):
        return super().fit(**merge_locals_kwargs(locals(), kwargs))
