from typing import Tuple

import tensorflow as tf
from tensorflow.keras.layers import Conv2D, UpSampling2D
from tensorflow.keras.layers import Input, concatenate
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import AveragePooling2D, LeakyReLU
from tensorflow.keras.layers import GlobalAveragePooling2D
from tensorflow.keras.metrics import MeanIoU
from tensorflow.keras.models import Model
from tensorflow.python.keras.optimizer_v2 import optimizer_v2

from omdenalore.computer_vision.loss_functions_semantic import LossFunctions


class SemanticSegemtationModel:
    """
    Returns a Semantic segmentation model

    :param pretrained_weights: pretrained weight matrix
    :type pretrained_weights: numpy ndarray
    :param input_size: img_height, img_width, input_channel
    :type input_size: tuple
    :param optimizer: optimization strategy
    :type optimizer: optimizer object
    :param num_classes: number of target classes
    :type num_classes: int
    :returns compiled Keras model
    :rtype: Keras model object

    :Example:

    from omdenalore.computer_vision.semantic_segmentation import
    SemanticSegemtationModel
    from tensorflow.keras.optimizers import Adam
    >>> num_classes = 10
    >>> optimizer = Adam(learning_rate=0.01)
    >>> input_size = (224, 224, 3)
    >>> semantic_segmentation_model = Model(num_classes, optimizer, input_size)
    >>> model = semantic_segmentation_model()
    """

    def __init__(
        self, num_classes: int, optimizer: optimizer_v2, input_size: Tuple[int]
    ) -> None:
        self.num_classes = num_classes
        self.input_size = input_size
        self.optimizer = optimizer

    def __call__(self):
        inputs = Input(self.input_size)

        lrelu = LeakyReLU(0.2)

        # Encoder
        conv1 = Conv2D(
            16,
            1,
            dilation_rate=(1, 1),
            padding="same",
            kernel_initializer="he_normal",
        )(inputs)
        conv1 = BatchNormalization()(conv1)
        conv1 = lrelu(conv1)
        conv1 = Conv2D(
            16,
            3,
            dilation_rate=(2, 2),
            padding="same",
            kernel_initializer="he_normal",
        )(conv1)
        conv1 = BatchNormalization()(conv1)
        conv1 = lrelu(conv1)

        conv2 = Conv2D(
            32,
            1,
            dilation_rate=(1, 1),
            padding="same",
            kernel_initializer="he_normal",
        )(conv1)
        conv2 = BatchNormalization()(conv2)
        conv2 = lrelu(conv2)
        conv2 = Conv2D(
            32,
            3,
            dilation_rate=(2, 2),
            padding="same",
            kernel_initializer="he_normal",
        )(conv2)
        conv2 = BatchNormalization()(conv2)
        conv2 = lrelu(conv2)

        conv3 = Conv2D(
            64,
            1,
            dilation_rate=(1, 1),
            padding="same",
            kernel_initializer="he_normal",
        )(conv2)
        conv3 = BatchNormalization()(conv3)
        conv3 = lrelu(conv3)
        conv3 = Conv2D(
            64,
            3,
            dilation_rate=(2, 2),
            padding="same",
            kernel_initializer="he_normal",
        )(conv3)
        conv3 = BatchNormalization()(conv3)
        conv3 = lrelu(conv3)

        conv4 = Conv2D(
            128,
            1,
            dilation_rate=(1, 1),
            padding="same",
            kernel_initializer="he_normal",
        )(conv3)
        conv4 = BatchNormalization()(conv4)
        conv4 = lrelu(conv4)
        conv4 = Conv2D(
            128,
            3,
            dilation_rate=(2, 2),
            padding="same",
            kernel_initializer="he_normal",
        )(conv4)
        conv4 = BatchNormalization()(conv4)
        conv4 = lrelu(conv4)

        # Decoder - Pyramid Module Pooling
        red = GlobalAveragePooling2D(name="red_pool")(conv4)
        red = tf.keras.layers.Reshape((1, 1, 128))(red)
        red = Conv2D(64, 1, name="red_1_by_1")(red)
        red = UpSampling2D(
            size=(400, 1920),
            interpolation="bilinear",
            name="red_upsampling",
        )(red)

        yellow = AveragePooling2D(
            pool_size=(2, 2),
            name="yellow_pool",
        )(conv4)
        yellow = Conv2D(64, 1, name="yellow_1_by_1")(yellow)
        yellow = UpSampling2D(
            size=2,
            interpolation="bilinear",
            name="yellow_upsampling",
        )(yellow)

        blue = AveragePooling2D(
            pool_size=(4, 4),
            name="blue_pool",
            padding="valid",
        )(conv4)
        blue = Conv2D(64, 1, name="blue_1_by_1")(blue)
        blue = UpSampling2D(
            size=4,
            interpolation="bilinear",
            name="blue_upsampling",
        )(blue)

        green = AveragePooling2D(
            pool_size=(4, 8),
            name="green_pool",
            padding="valid",
        )(conv4)
        green = Conv2D(64, 1, name="green_1_by_1")(green)
        green = UpSampling2D(
            size=(4, 8), interpolation="bilinear", name="green_upsampling"
        )(green)

        merge9 = concatenate([conv4, red, yellow, blue, green])

        output = Conv2D(
            256,
            3,
            padding="same",
            kernel_initializer="he_normal",
        )(merge9)
        output = lrelu(output)
        output = Conv2D(
            3,
            1,
            activation="softmax",
            padding="same",
        )(output)

        model = Model(inputs, output)
        model.compile(
            optimizer=self.optimizer,
            loss=LossFunctions.focal_tversky_loss,
            metrics=[MeanIoU(num_classes=self.num_classes)],
        )

        return model
