# -*- coding:utf-8 -*-
# Author: hankcs
# Date: 2019-11-10 13:19

import math
from typing import Union, Tuple, Any, Iterable
import tensorflow as tf
from hanlp.common.keras_component import KerasComponent
from hanlp_common.structure import SerializableDict
from hanlp.layers.transformers.loader_tf import build_transformer
from hanlp.optimizers.adamw import create_optimizer
from hanlp.transform.table_tf import TableTransform
from hanlp.utils.log_util import logger
from hanlp_common.util import merge_locals_kwargs
from transformers.tokenization_utils import PreTrainedTokenizer

class TransformerTextTransform(TableTransform):

    def __init__(self, config: SerializableDict = None, map_x=False, map_y=True, x_columns=None,
                 y_column=-1, skip_header=True, delimiter='auto', multi_label=False, **kwargs) -> None:
        super().__init__(config, map_x, map_y, x_columns, y_column, multi_label, skip_header, delimiter, **kwargs)
        self.tokenizer: PreTrainedTokenizer = None

    def inputs_to_samples(self, inputs, gold=False):
        tokenizer = self.tokenizer
        max_length = self.config.max_length
        num_features = None
        pad_token = None if self.label_vocab.mutable else tokenizer.convert_tokens_to_ids(['[PAD]'])[0]
        for (X, Y) in super().inputs_to_samples(inputs, gold):
            if self.label_vocab.mutable:
                yield None, Y
                continue
            if isinstance(X, str):
                X = (X,)
            if num_features is None:
                num_features = self.config.num_features
            assert num_features == len(X), f'Numbers of features {num_features} ' \
                                           f'inconsistent with current {len(X)}={X}'
            text_a = X[0]
            text_b = X[1] if len(X) > 1 else None
            tokens_a = self.tokenizer.tokenize(text_a)
            tokens_b = self.tokenizer.tokenize(text_b) if text_b else None
            tokens = ["[CLS]"] + tokens_a + ["[SEP]"]
            segment_ids = [0] * len(tokens)
            if tokens_b:
                tokens += tokens_b
                segment_ids += [1] * len(tokens_b)
            token_ids = self.tokenizer.convert_tokens_to_ids(tokens)
            attention_mask = [1] * len(token_ids)
            diff = max_length - len(token_ids)
            if diff < 0:
                # logger.warning(
                #     f'Input tokens {tokens} exceed the max sequence length of {max_length - 2}. '
                #     f'The exceeded part will be truncated and ignored. '
                #     f'You are recommended to split your long text into several sentences within '
                #     f'{max_length - 2} tokens beforehand.')
                token_ids = token_ids[:max_length]
                attention_mask = attention_mask[:max_length]
                segment_ids = segment_ids[:max_length]
            elif diff > 0:
                token_ids += [pad_token] * diff
                attention_mask += [0] * diff
                segment_ids += [0] * diff

            assert len(token_ids) == max_length, "Error with input length {} vs {}".format(len(token_ids), max_length)
            assert len(attention_mask) == max_length, "Error with input length {} vs {}".format(len(attention_mask),
                                                                                                max_length)
            assert len(segment_ids) == max_length, "Error with input length {} vs {}".format(len(segment_ids),
                                                                                             max_length)

            label = Y
            yield (token_ids, attention_mask, segment_ids), label

    def create_types_shapes_values(self) -> Tuple[Tuple, Tuple, Tuple]:
        max_length = self.config.max_length
        types = (tf.int32, tf.int32, tf.int32), tf.string
        shapes = ([max_length], [max_length], [max_length]), [None, ] if self.config.get('multi_label', None) else []
        values = (0, 0, 0), self.label_vocab.safe_pad_token
        return types, shapes, values

    def x_to_idx(self, x) -> Union[tf.Tensor, Tuple]:
        logger.fatal('map_x should always be set to True')
        exit(1)

    def y_to_idx(self, y) -> tf.Tensor:
        if self.config.get('multi_label', None):
            # need to change index to binary vector
            mapped = tf.map_fn(fn=lambda x: tf.cast(self.label_vocab.lookup(x), tf.int32), elems=y,
                               fn_output_signature=tf.TensorSpec(dtype=tf.dtypes.int32, shape=[None, ]))
            one_hots = tf.one_hot(mapped, len(self.label_vocab))
            idx = tf.reduce_sum(one_hots, -2)
        else:
            idx = self.label_vocab.lookup(y)
        return idx

    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None,
                     batch=None) -> Iterable:
        # Prediction to be Y > 0:
        if self.config.get('multi_label', None):
            preds = Y
        else:
            preds = tf.argmax(Y, axis=-1)
        for y in preds:
            yield self.label_vocab.idx_to_token[y]

    def input_is_single_sample(self, input: Any) -> bool:
        return isinstance(input, (str, tuple))


class TransformerClassifierTF(KerasComponent):

    def __init__(self, bert_text_transform=None) -> None:
        if not bert_text_transform:
            bert_text_transform = TransformerTextTransform()
        super().__init__(bert_text_transform)
        self.model: tf.keras.Model
        self.transform: TransformerTextTransform = bert_text_transform

    # noinspection PyMethodOverriding
    def fit(self, trn_data: Any, dev_data: Any, save_dir: str, transformer: str, max_length: int = 128,
            optimizer='adamw', warmup_steps_ratio=0.1, use_amp=False, batch_size=32,
            epochs=3, logger=None, verbose=1, **kwargs):
        return super().fit(**merge_locals_kwargs(locals(), kwargs))

    def evaluate_output(self, tst_data, out, num_batches, metric):
        out.write('sentence\tpred\tgold\n')
        total, correct, score = 0, 0, 0
        for idx, batch in enumerate(tst_data):
            outputs = self.model.predict_on_batch(batch[0])
            outputs = tf.argmax(outputs, axis=1)
            for X, Y_pred, Y_gold, in zip(batch[0][0], outputs, batch[1]):
                feature = ' '.join(self.transform.tokenizer.convert_ids_to_tokens(X.numpy()))
                feature = feature.replace(' ##', '')  # fix sub-word generated by BERT tagger
                out.write('{}\t{}\t{}\n'.format(feature,
                                                self._y_id_to_str(Y_pred),
                                                self._y_id_to_str(Y_gold)))
                total += 1
                correct += int(tf.equal(Y_pred, Y_gold).numpy())
            score = correct / total
            print('\r{}/{} {}: {:.2f}'.format(idx + 1, num_batches, metric, score * 100), end='')
        print()
        return score

    def _y_id_to_str(self, Y_pred) -> str:
        return self.transform.label_vocab.idx_to_token[Y_pred.numpy()]

    def build_loss(self, loss, **kwargs):
        if loss:
            assert isinstance(loss, tf.keras.losses.loss), 'Must specify loss as an instance in tf.keras.losses'
            return loss
        elif self.config.get('multi_label', None):
            # Loss to be BinaryCrossentropy for multi-label:
            loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
        else:
            loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        return loss

    # noinspection PyMethodOverriding
    def build_optimizer(self, optimizer, use_amp, train_steps, warmup_steps, **kwargs):
        if optimizer == 'adamw':
            opt = create_optimizer(init_lr=5e-5, num_train_steps=train_steps, num_warmup_steps=warmup_steps)
            # opt = tfa.optimizers.AdamW(learning_rate=3e-5, epsilon=1e-08, weight_decay=0.01)
            # opt = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
            self.config.optimizer = tf.keras.utils.serialize_keras_object(opt)
            lr_config = self.config.optimizer['config']['learning_rate']['config']
            if hasattr(lr_config['decay_schedule_fn'], 'get_config'):
                lr_config['decay_schedule_fn'] = dict(
                    (k, v) for k, v in lr_config['decay_schedule_fn'].config().items() if not k.startswith('_'))
        else:
            opt = super().build_optimizer(optimizer)
        if use_amp:
            # loss scaling is currently required when using mixed precision
            opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')
        return opt

    # noinspection PyMethodOverriding
    def build_model(self, transformer, max_length, **kwargs):
        model, self.transform.tokenizer = build_transformer(transformer, max_length, len(self.transform.label_vocab),
                                                            tagging=False)
        return model

    def build_vocab(self, trn_data, logger):
        train_examples = super().build_vocab(trn_data, logger)
        warmup_steps_per_epoch = math.ceil(train_examples * self.config.warmup_steps_ratio / self.config.batch_size)
        self.config.warmup_steps = warmup_steps_per_epoch * self.config.epochs
        return train_examples

    def build_metrics(self, metrics, logger, **kwargs):
        if self.config.get('multi_label', None):
            metric = tf.keras.metrics.BinaryAccuracy('binary_accuracy')
        else:
            metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
        return [metric]
