from abc import abstractmethod
import typing

import tensorflow.keras as tfk

from .base_class import BaseClass


class ModelBuilderBase(BaseClass):
    """Building and compiling models to train with Trainer.

    Examples:
        >>> model_builder = ModelBuilder(config)
        >>> model = model_builder.get_compiled_model()
        >>> callbacks = model_builder.get_callbacks()
        >>> class_weight = model_builder.get_class_weight()
        >>> model.fit(train_gen, n_iter_train, callbacks=callbacks, class_weight=class_weight)

    """

    @abstractmethod
    def get_compiled_model(self) -> tfk.Model:
        """Generates the model for training, and returns the compiled model.

        Returns:
            A compiled ``tensorflow.keras`` model.
        """

        pass

    def get_callbacks(self) -> list:
        """Returns any callbacks for ``fit``.

        Returns:
            list of ``tf.keras.Callback`` objects. ``Orchestrator`` will handle the ``ModelCheckpoint`` and ``Tensorboard`` callbacks.
            Still, you can return each of these two callbacks, and orchestrator will modify your callbacks if needed.

        """
        return list()

    def get_class_weight(self) -> typing.Optional[dict]:
        """Set this if you want to pass ``class_weight`` to ``fit``.

        Returns:
           Optional dictionary mapping class indices (integers) to a weight (float) value.
           used for weighting the loss function (during training only).
           This can be useful to tell the model to "pay more attention" to samples from an under-represented class.

        """

        return None
    #
    # @abstractmethod
    # def post_process(self, y_pred):
    #     """Define your post-processing, used for evaluation.
    #
    #     If you have ``softmax`` as your final layer, but your labels are sparse-categorical labels, you will need to
    #      post-process the output of your model before comparing it to ``y_true``. In this case you should use
    #      ``return np.argmax(y_pred)``.
    #
    #      Note: make sure that this is compatible with your evaluation functions and ground truth labels generated
    #       using ``DataLoader``'s generators.
    #
    #      Args:
    #          y_pred: a tensor generated by ``model.predict`` method. Note that the first dimension is ``batch``.
    #
    #     Returns:
    #         post-processed batch of y_pred, ready to be compared to ground-truth.
    #
    #     """

    @classmethod
    def __subclasshook__(cls, c):
        """This defines the __subclasshook__ class method.
        This special method is called by the Python interpreter to answer the question,
         Is the class C a subclass of this class?
        """

        if cls is ModelBuilderBase:
            attrs = set(dir(c))
            if set(cls.__abstractmethods__) <= attrs:
                return True
        return NotImplemented
