# -*- coding:utf-8 -*-
# Author: hankcs
# Date: 2019-10-27 14:22
import inspect
from abc import ABC, abstractmethod
from typing import Generator, Tuple, Union, Iterable, Any

import tensorflow as tf

from hanlp_common.structure import SerializableDict
from hanlp.common.vocab_tf import VocabTF
from hanlp.utils.io_util import get_resource
from hanlp.utils.log_util import logger


class Transform(ABC):

    def __init__(self, config: SerializableDict = None, map_x=True, map_y=True, **kwargs) -> None:
        super().__init__()
        self.map_y = map_y
        self.map_x = map_x
        if kwargs:
            if not config:
                config = SerializableDict()
            for k, v in kwargs.items():
                config[k] = v
        self.config = config
        self.output_types = None
        self.output_shapes = None
        self.padding_values = None
        # Fix tf memory leak: https://github.com/tensorflow/tensorflow/issues/37653#issuecomment-1000517720
        self.py_func_set_to_cleanup = set()

    @abstractmethod
    def fit(self, trn_path: str, **kwargs) -> int:
        """
        Build the vocabulary from training file

        Parameters
        ----------
        trn_path : path to training set
        kwargs

        Returns
        -------
        int
            How many samples in the training set
        """
        raise NotImplementedError('%s.%s()' % (self.__class__.__name__, inspect.stack()[0][3]))

    def build_config(self):
        """
        By default, call build_types_shapes_values, usually called in component's build method.
        You can perform other building task here. Remember to call super().build_config
        """
        self.output_types, self.output_shapes, self.padding_values = self.create_types_shapes_values()
        # We prefer list over shape here, as it's easier to type [] than ()
        # if isinstance(self.output_shapes, tuple):
        #     self.output_shapes = list(self.output_shapes)
        # for i, shapes in enumerate(self.output_shapes):
        #     if isinstance(shapes, tuple):
        #         self.output_shapes[i] = list(shapes)
        #     for j, shape in enumerate(shapes):
        #         if isinstance(shape, tuple):
        #             shapes[j] = list(shape)

    @abstractmethod
    def create_types_shapes_values(self) -> Tuple[Tuple, Tuple, Tuple]:
        """
        Create dataset related values,
        """
        raise NotImplementedError('%s.%s()' % (self.__class__.__name__, inspect.stack()[0][3]))

    @abstractmethod
    def file_to_inputs(self, filepath: str, gold=True):
        """
        Transform file to inputs. The inputs are defined as raw features (e.g. words) to be processed into more
        features (e.g. forms and characters)

        Parameters
        ----------
        filepath
        gold
        """
        raise NotImplementedError('%s.%s()' % (self.__class__.__name__, inspect.stack()[0][3]))

    def inputs_to_samples(self, inputs, gold=False):
        if gold:
            yield from inputs
        else:
            for x in inputs:
                yield x, self.padding_values[-1]

    def file_to_samples(self, filepath: str, gold=True):
        """
        Transform file to samples
        Parameters
        ----------
        filepath
        gold
        """
        filepath = get_resource(filepath)
        inputs = self.file_to_inputs(filepath, gold)
        yield from self.inputs_to_samples(inputs, gold)

    def file_to_dataset(self, filepath: str, gold=True, map_x=None, map_y=None, batch_size=32, shuffle=None,
                        repeat=None,
                        drop_remainder=False,
                        prefetch=1,
                        cache=True,
                        **kwargs) -> tf.data.Dataset:
        """
        Transform file to dataset

        Parameters
        ----------
        filepath
        gold : bool
            Whether it's processing gold data or not. Example: there is usually a column for gold answer
            when gold = True.
        map_x : bool
            Whether call map_x or not. Default to self.map_x
        map_y : bool
            Whether call map_y or not. Default to self.map_y
        batch_size
        shuffle
        repeat
        prefetch
        kwargs

        Returns
        -------

        """

        # debug
        # for sample in self.file_to_samples(filepath):
        #     pass

        def generator():
            inputs = self.file_to_inputs(filepath, gold)
            samples = self.inputs_to_samples(inputs, gold)
            yield from samples

        return self.samples_to_dataset(generator, map_x, map_y, batch_size, shuffle, repeat, drop_remainder, prefetch,
                                       cache)

    def inputs_to_dataset(self, inputs, gold=False, map_x=None, map_y=None, batch_size=32, shuffle=None, repeat=None,
                          drop_remainder=False,
                          prefetch=1, cache=False, **kwargs) -> tf.data.Dataset:
        # debug
        # for sample in self.inputs_to_samples(inputs):
        #     pass

        def generator():
            samples = self.inputs_to_samples(inputs, gold)
            yield from samples

        return self.samples_to_dataset(generator, map_x, map_y, batch_size, shuffle, repeat, drop_remainder, prefetch,
                                       cache)

    def samples_to_dataset(self, samples: Generator, map_x=None, map_y=None, batch_size=32, shuffle=None, repeat=None,
                           drop_remainder=False,
                           prefetch=1, cache=True) -> tf.data.Dataset:
        output_types, output_shapes, padding_values = self.output_types, self.output_shapes, self.padding_values
        if not all(v for v in [output_shapes, output_shapes,
                               padding_values]):
            # print('Did you forget to call build_config() on your transform?')
            self.build_config()
            output_types, output_shapes, padding_values = self.output_types, self.output_shapes, self.padding_values
        assert all(v for v in [output_shapes, output_shapes,
                               padding_values]), 'Your create_types_shapes_values returns None, which is not allowed'
        # if not callable(samples):
        #     samples = Transform.generator_to_callable(samples)
        if not hasattr(tf.compat.v1.get_default_graph(), '_py_funcs_used_in_graph'):
            tf.compat.v1.get_default_graph()._py_funcs_used_in_graph = []
        py_func_set_before = set(tf.compat.v1.get_default_graph()._py_funcs_used_in_graph)
        dataset = tf.data.Dataset.from_generator(samples, output_types=output_types, output_shapes=output_shapes)
        if cache:
            logger.debug('Dataset cache enabled')
            dataset = dataset.cache(cache if isinstance(cache, str) else '')
        if shuffle:
            if isinstance(shuffle, bool):
                shuffle = 1024
            dataset = dataset.shuffle(shuffle)
        if repeat:
            dataset = dataset.repeat(repeat)
        if batch_size:
            dataset = dataset.padded_batch(batch_size, output_shapes, padding_values, drop_remainder)
        if prefetch:
            dataset = dataset.prefetch(prefetch)
        if map_x is None:
            map_x = self.map_x
        if map_y is None:
            map_y = self.map_y
        if map_x or map_y:
            def mapper(X, Y):
                if map_x:
                    X = self.x_to_idx(X)
                if map_y:
                    Y = self.y_to_idx(Y)
                return X, Y

            dataset = dataset.map(mapper, num_parallel_calls=tf.data.experimental.AUTOTUNE)
        py_func_set_after = set(tf.compat.v1.get_default_graph()._py_funcs_used_in_graph) - py_func_set_before
        self.py_func_set_to_cleanup |= py_func_set_after
        return dataset

    @abstractmethod
    def x_to_idx(self, x) -> Union[tf.Tensor, Tuple]:
        raise NotImplementedError('%s.%s()' % (self.__class__.__name__, inspect.stack()[0][3]))

    @abstractmethod
    def y_to_idx(self, y) -> tf.Tensor:
        raise NotImplementedError('%s.%s()' % (self.__class__.__name__, inspect.stack()[0][3]))

    def lock_vocabs(self):
        for key, value in vars(self).items():
            if isinstance(value, VocabTF):
                value.lock()

    def summarize_vocabs(self, logger=None, header='Vocab summary:'):
        output = header + '\n'
        vocabs = {}
        for key, value in vars(self).items():
            if isinstance(value, VocabTF):
                vocabs[key] = value
        # tag vocab comes last usually
        for key, value in sorted(vocabs.items(), key=lambda kv: len(kv[1]), reverse=True):
            output += f'{key}' + value.summary(verbose=False) + '\n'
        output = output.strip()
        if logger:
            logger.info(output)
        else:
            print(output)

    @staticmethod
    def generator_to_callable(generator: Generator):
        return lambda: (x for x in generator)

    def str_to_idx(self, X, Y) -> Tuple[Union[tf.Tensor, Tuple], tf.Tensor]:
        return self.x_to_idx(X), self.y_to_idx(Y)

    def X_to_inputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]]) -> Iterable:
        return [repr(x) for x in X]

    def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None,
                     batch=None) -> Iterable:
        return [repr(y) for y in Y]

    def XY_to_inputs_outputs(self, X: Union[tf.Tensor, Tuple[tf.Tensor]],
                             Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False) -> Iterable:
        """
        Convert predicted tensors to outputs

        Parameters
        ----------
        X : Union[tf.Tensor, Tuple[tf.Tensor]]
            The inputs of model
        Y : Union[tf.Tensor, Tuple[tf.Tensor]]
            The outputs of model

        Returns
        -------

        """
        return [(x, y) for x, y in zip(self.X_to_inputs(X), self.Y_to_outputs(Y, gold))]

    def input_is_single_sample(self, input: Any) -> bool:
        return False

    def input_to_inputs(self, input: Any) -> Tuple[Any, bool]:
        """
        If input is one sample, convert it to a list which contains this unique sample

        Parameters
        ----------
        input :
            sample or samples

        Returns
        -------
        (inputs, converted) : Tuple[Any, bool]

        """
        flat = self.input_is_single_sample(input)
        if flat:
            input = [input]
        return input, flat

    def input_truth_output_to_str(self, input, truth, output):
        """
        Convert input truth output to string representation, usually for writing to file during evaluation

        Parameters
        ----------
        input
        truth
        output

        Returns
        -------

        """
        return '\t'.join([input, truth, output]) + '\n'

    def cleanup(self):
        new_py_funcs = set(tf.compat.v1.get_default_graph()._py_funcs_used_in_graph) - self.py_func_set_to_cleanup
        tf.compat.v1.get_default_graph()._py_funcs_used_in_graph = list(new_py_funcs)
        self.py_func_set_to_cleanup = set()
