from enum import Enum

import tensorflow as tf
from tensorflow.keras import models, layers

BUFFER_SIZE = 10000


class ModelArchitecture(Enum):
    SA_MIRI = 1
    RESNET101 = 2
    MOBILENET = 3

    def __str__(self):
        return self.name.lower()

    def __repr__(self):
        return str(self)

    @staticmethod
    def argparse(s):
        try:
            return ModelArchitecture[s.upper()]
        except KeyError:
            return s


def get_model(show_summary=True, architecture=ModelArchitecture.SA_MIRI):
    model = None

    if architecture == ModelArchitecture.SA_MIRI:
        model = models.Sequential(name="samiri2019")
        model.add(layers.Conv2D(32, (5, 5), activation='relu', input_shape=(28, 28, 1)))
        model.add(layers.MaxPooling2D((2, 2)))
        model.add(layers.Conv2D(64, (5, 5), activation='relu'))
        model.add(layers.MaxPooling2D((2, 2)))
        model.add(layers.Flatten())
        model.add(layers.Dense(10, activation='softmax'))
    elif architecture == ModelArchitecture.MOBILENET:
        model = tf.keras.applications.MobileNetV2(classes=10, weights=None)
    elif architecture == ModelArchitecture.RESNET101:
        model = tf.keras.applications.ResNet101V2(classes=10, weights=None)

    if model and show_summary:
        model.summary()

    return model


def make_datasets_unbatched(datasets, set_name='train', architecture=ModelArchitecture.SA_MIRI):
    # Scaling MNIST data from (0, 255] to (0., 1.]
    def scale(image, label):
        image = tf.cast(image, tf.float32)

        if architecture is not ModelArchitecture.SA_MIRI:
            image = tf.image.resize(image, [224, 224], tf.image.ResizeMethod.NEAREST_NEIGHBOR)

        image /= 255

        label = tf.one_hot(label, depth=10)
        return image, label

    if 'train' in set_name:
        return datasets['train'].map(scale, num_parallel_calls=tf.data.experimental.AUTOTUNE).\
            cache().repeat().shuffle(BUFFER_SIZE)
    else:
        return datasets['test'].map(scale, num_parallel_calls=tf.data.experimental.AUTOTUNE)

tensorflow.keras.layers.Dense tensorflow.keras.applications.ResNet101V2 tensorflow.keras.layers.Conv2D tensorflow.cast tensorflow.image.resize tensorflow.one_hot tensorflow.keras.layers.MaxPooling2D tensorflow.keras.models.Sequential tensorflow.keras.layers.Flatten tensorflow.keras.applications.MobileNetV2
[(32, 'tensorflow.keras.models.Sequential', 'models.Sequential', ([], {'name': '"""samiri2019"""'}), False, 'from tensorflow.keras import models, layers\n'), (53, 'tensorflow.cast', 'tf.cast', (['image', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (60, 'tensorflow.one_hot', 'tf.one_hot', (['label'], {'depth': '(10)'}), True, 'import tensorflow as tf\n'), (33, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(32)', '(5, 5)'], {'activation': '"""relu"""', 'input_shape': '(28, 28, 1)'}), False, 'from tensorflow.keras import models, layers\n'), (34, 'tensorflow.keras.layers.MaxPooling2D', 'layers.MaxPooling2D', (['(2, 2)'], {}), False, 'from tensorflow.keras import models, layers\n'), (35, 'tensorflow.keras.layers.Conv2D', 'layers.Conv2D', (['(64)', '(5, 5)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras import models, layers\n'), (36, 'tensorflow.keras.layers.MaxPooling2D', 'layers.MaxPooling2D', (['(2, 2)'], {}), False, 'from tensorflow.keras import models, layers\n'), (37, 'tensorflow.keras.layers.Flatten', 'layers.Flatten', ([], {}), False, 'from tensorflow.keras import models, layers\n'), (38, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(10)'], {'activation': '"""softmax"""'}), False, 'from tensorflow.keras import models, layers\n'), (40, 'tensorflow.keras.applications.MobileNetV2', 'tf.keras.applications.MobileNetV2', ([], {'classes': '(10)', 'weights': 'None'}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.image.resize', 'tf.image.resize', (['image', '[224, 224]', 'tf.image.ResizeMethod.NEAREST_NEIGHBOR'], {}), True, 'import tensorflow as tf\n'), (42, 'tensorflow.keras.applications.ResNet101V2', 'tf.keras.applications.ResNet101V2', ([], {'classes': '(10)', 'weights': 'None'}), True, 'import tensorflow as tf\n')]