import logging
import math
import warnings
from typing import Optional

import numpy as np
import tensorflow as tf
from keras import layers
from tensorflow import keras

from ivus_segmentation.architectures import inception_resnetv2_unet, unet
from ivus_segmentation.augment import (RandomFlip, RandomRotation, RandomZoom,
                                       perform_augments)
from ivus_segmentation.evaluation import custom_classification_report
from ivus_segmentation.metrics import (categorical_crossentropy,
                                       dice_for_class, jaccard_coef_loss)
from ivus_segmentation.preprocessing import IdentityPreprocessor, Preprocessor

warnings.filterwarnings("ignore")
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)


# from ivus_segmentation.augment import gen_dtype, get_image_generator
CUSTOM_OBJECTS = {
    "jaccard_coef_loss": jaccard_coef_loss,
    "dice_background": dice_for_class(0),
    "dice_media": dice_for_class(1),
    "dice_lumen": dice_for_class(2),
}


class IvusModel:
    preprocessor: Preprocessor

    def __init__(
        self,
        input_shape,
        arch=None,
        loss=None,
        augments=None,
        num_classes=3,
        preprocessor: Optional[Preprocessor] = None,
    ):
        if arch is None:
            arch = unet
        if loss is None:
            loss = categorical_crossentropy
        self.arch = arch
        self.augments = list(augments) if augments is not None else []
        self.loss = loss
        self.num_classes = num_classes
        self.preprocessor = (
            preprocessor
            if preprocessor is not None
            else IdentityPreprocessor(input_shape)
        )
        shape = self.preprocessor.input_shape()
        self.model = self.build_model(shape)

    def build_model(self, input_shape):
        inputs = keras.Input(input_shape)

        x = self.arch(inputs)
        outputs = layers.Conv2D(
            self.num_classes, 3, activation="softmax", padding="same"
        )(x)
        return keras.Model(inputs, outputs)

    def __getattr__(self, name):
        """
        This method enables to access an attribute/method of self.model.
        Thus, any method of keras.Model() can be used transparently from a SubModel object
        """
        return getattr(self.model, name)

    def fit(self, x, y, validation_data=None, batch_size=2, **kwargs):
        x = np.array([self.preprocessor(img) for img in x])
        if validation_data is not None:
            assert len(validation_data) == 2, "Expected tuple of x and y"
            x_val, y_val = validation_data
            x_val = np.array([self.preprocessor(img) for img in x_val])
            validation_data = (x_val, y_val)
        ds = perform_augments(x, y, batch_size, self.augments)
        return self.model.fit(ds, validation_data=validation_data, **kwargs)

    def predict(self, x, **kwargs):
        x = np.array([self.preprocessor(img) for img in x])
        return self.model.predict(x, **kwargs)

    def load_model(self, filename, *args, **kwargs):
        custom_objects = kwargs.pop("custom_objects", {})
        custom_objects.update(CUSTOM_OBJECTS)
        self.model = keras.models.load_model(
            filename, custom_objects=custom_objects,
            *args, **kwargs
        )
        return self

    @classmethod
    def from_saved_model(cls, input_shape, filename, *args, **kwargs):
        instance = cls(input_shape)
        instance.load_model(filename, *args, **kwargs)
        return instance

    def compile(self, **kwargs):
        metrics = [dice_for_class(i) for i in range(3)]

        return self.model.compile(loss=self.loss, metrics=metrics, **kwargs)

    def evaluate(self, X, y, **kwargs):
        X = np.array([self.preprocessor(img) for img in X])
        score = {
            k: v
            for k, v in zip(
                ["loss", "dice_media", "dice_lumen"],
                self.model.evaluate(X, y, **kwargs),
            )
        }
        masks_test_from_pred = np.argmax(y, axis=-1)
        prediction_masks_test = np.argmax(self.model.predict(X), axis=-1)
        df = custom_classification_report(
            masks_test_from_pred, prediction_masks_test, verbose=True
        )
        return {**df, **score}
