import tensorflow as tf


@tf.function
def BinaryAccuracy_Infiltrates(y_true, y_pred, i=0):
    return tf.keras.metrics.binary_accuracy(y_true[:, i], y_pred[:, i])


@tf.function
def BinaryAccuracy_Pneumonia(y_true, y_pred, i=1):
    return tf.keras.metrics.binary_accuracy(y_true[:, i], y_pred[:, i])


@tf.function
def BinaryAccuracy_Covid19(y_true, y_pred, i=2):
    return tf.keras.metrics.binary_accuracy(y_true[:, i], y_pred[:, i])


@tf.function
def BinaryAccuracy_Normal(y_true, y_pred, i=3):
    return tf.keras.metrics.binary_accuracy(y_true[:, i], y_pred[:, i])


class CustomModelCheckpoint(tf.keras.callbacks.ModelCheckpoint):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.wait_epoch_warmup = kwargs.get("wait_epoch_warmup")

    def on_epoch_end(self, epoch, logs=None):
        if self.wait_epoch_warmup:
            if (epoch + 1) >= self.wait_epoch_warmup:
                super().on_epoch_end(epoch, logs)
            else:
                self.epochs_since_last_save += 1
                print(f"Skipping save model (wait_epoch_warmup={self.wait_epoch_warmup - (epoch + 1)})")
        else:
            super().on_epoch_end(epoch, logs)


class CustomEarlyStopping(tf.keras.callbacks.EarlyStopping):
    def __init__(self, *args, **kwargs):
        self.minimum_epochs = kwargs.get("minimum_epochs", 0)
        kwargs.pop('minimum_epochs', None)  # Problems with EarlyStopping kwargs
        super().__init__(*args, **kwargs)

    def on_epoch_end(self, epoch, logs=None):
        if epoch >= self.minimum_epochs:
            super().on_epoch_end(epoch, logs)


def get_losses():
    losses = [tf.keras.losses.BinaryCrossentropy()]
    return losses


def get_metrics(single_output_idx, add_normal=False):
    metrics = []
    if single_output_idx is None:  # Multi-label
        print("###### Multi-label classification ######")
        metrics += [
            BinaryAccuracy_Infiltrates,
            BinaryAccuracy_Pneumonia,
            BinaryAccuracy_Covid19
        ]

        # Add normal class
        if add_normal:
            metrics.append(BinaryAccuracy_Normal)
    else:
        print(f"###### Multi-class classification (cls: '{single_output_idx}') ######")
        metrics = [
            tf.keras.metrics.BinaryAccuracy(),
            tf.keras.metrics.AUC(),
            tf.keras.metrics.Precision(),
            tf.keras.metrics.Recall()
        ]
    return metrics


def get_model(backbone, classes=None, target_size=None, freeze_base_model=True, ignore_model=None):
    istrainable = not freeze_base_model

    # Select backbone
    if backbone == "resnet50":
        from tensorflow.keras.applications.resnet import ResNet50 as TFModel
        from tensorflow.keras.applications.resnet import preprocess_input
    elif backbone == "resnet50v2":
        from tensorflow.keras.applications.resnet_v2 import ResNet50V2 as TFModel
        from tensorflow.keras.applications.resnet_v2 import preprocess_input
    elif backbone == "resnet101v2":
        from tensorflow.keras.applications.resnet_v2 import ResNet101V2 as TFModel
        from tensorflow.keras.applications.resnet_v2 import preprocess_input
    elif backbone == "vgg16":
        from tensorflow.keras.applications.vgg16 import VGG16 as TFModel
        from tensorflow.keras.applications.vgg16 import preprocess_input
    elif backbone == "efficientnetb0":
        from tensorflow.keras.applications.efficientnet import EfficientNetB0 as TFModel
        from tensorflow.keras.applications.efficientnet import preprocess_input
    elif backbone == "efficientnetb7":
        from tensorflow.keras.applications.efficientnet import EfficientNetB7 as TFModel
        from tensorflow.keras.applications.efficientnet import preprocess_input
    else:
        raise ValueError(f"Unknown backbone: {backbone}")

    if ignore_model:
        model = None
    else:
        # Instantiate base model with pre-trained weights
        base_model = TFModel(input_shape=(*target_size, 3), include_top=False, weights="imagenet")

        # Freeze base model
        # base_model.trainable = istrainable
        for layers in base_model.layers:
            layers.trainable = istrainable

        # Create a new model on top
        inputs = base_model.input
        x = base_model(inputs)

        # Option A
        x = tf.keras.layers.GlobalAveragePooling2D(name='avg_pool')(x)

        # Option B
        # x = tf.keras.layers.Flatten(name='flatten')(x)
        # x = tf.keras.layers.Dense(512, activation='relu', name='fc1')(x)
        # x = tf.keras.layers.Dense(512, activation='relu', name='fc2')(x)

        # Outputs
        outputs = tf.keras.layers.Dense(classes, activation="sigmoid", name='predictions')(x)
        model = tf.keras.Model(inputs, outputs)

    return model, preprocess_input


def add_tabular_input(model, classes):
    # Input1
    input1 = model.input
    input2 = tf.keras.layers.Input(shape=(2,), name="input_2b")

    # Pre-outputs 1x3 + 1x3
    output1 = model.output
    output2 = tf.keras.layers.Dense(classes, activation="sigmoid", name='output_tab')(input2)

    # Outputs
    x = tf.keras.layers.Concatenate(axis=1)([output1, output2])
    output = tf.keras.layers.Dense(classes, activation="sigmoid", name='final_predictions')(x)
    model = tf.keras.Model([input1, input2], output)

    return model


def unfreeze_base_model(model, n=None, unfreeze=True):
    base_model = model.layers[1].layers

    # Select number of layers to unfreeze
    idx = 0
    if n is not None:
        if isinstance(n, int):
            idx = n
            print(f"Unfreezing {len(base_model) - idx} layers")
        elif isinstance(n, float) and 0.0 < n <= 1.0:
            idx = int(len(base_model) * n)
            print(f"Unfreezing {idx} layers")
        else:
            raise ValueError("Invalid number of layers")

    # We unfreeze all layers but BatchNorm (to not destroy the non-trainable weights)
    for layer in base_model[-idx:]:
        if not isinstance(layer, tf.keras.layers.BatchNormalization):
            layer.trainable = True

[(6, 'tensorflow.keras.metrics.binary_accuracy', 'tf.keras.metrics.binary_accuracy', True, 'import tensorflow as tf\n'), (11, 'tensorflow.keras.metrics.binary_accuracy', 'tf.keras.metrics.binary_accuracy', True, 'import tensorflow as tf\n'), (16, 'tensorflow.keras.metrics.binary_accuracy', 'tf.keras.metrics.binary_accuracy', True, 'import tensorflow as tf\n'), (21, 'tensorflow.keras.metrics.binary_accuracy', 'tf.keras.metrics.binary_accuracy', True, 'import tensorflow as tf\n'), (139, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', True, 'import tensorflow as tf\n'), (148, 'tensorflow.keras.Model', 'tf.keras.Model', True, 'import tensorflow as tf\n'), (53, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', True, 'import tensorflow as tf\n'), (110, 'tensorflow.keras.applications.efficientnet.EfficientNetB7', 'TFModel', True, 'from tensorflow.keras.applications.efficientnet import EfficientNetB7 as TFModel\n'), (131, 'tensorflow.keras.Model', 'tf.keras.Model', True, 'import tensorflow as tf\n'), (143, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', True, 'import tensorflow as tf\n'), (146, 'tensorflow.keras.layers.Concatenate', 'tf.keras.layers.Concatenate', True, 'import tensorflow as tf\n'), (147, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', True, 'import tensorflow as tf\n'), (73, 'tensorflow.keras.metrics.BinaryAccuracy', 'tf.keras.metrics.BinaryAccuracy', True, 'import tensorflow as tf\n'), (74, 'tensorflow.keras.metrics.AUC', 'tf.keras.metrics.AUC', True, 'import tensorflow as tf\n'), (75, 'tensorflow.keras.metrics.Precision', 'tf.keras.metrics.Precision', True, 'import tensorflow as tf\n'), (76, 'tensorflow.keras.metrics.Recall', 'tf.keras.metrics.Recall', True, 'import tensorflow as tf\n'), (122, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'tf.keras.layers.GlobalAveragePooling2D', True, 'import tensorflow as tf\n'), (130, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', True, 'import tensorflow as tf\n')]