import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import (
    Conv2D,
    MaxPool2D,
    Dropout,
    Conv2DTranspose,
    concatenate,
)
import matplotlib.pyplot as plt


class EncoderBlock(tf.keras.layers.Layer):
    def __init__(self, filters, rate=None, pooling=True, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        self.filters = filters
        self.rate = rate
        self.pooling = pooling
        self.conv1 = Conv2D(
            self.filters,
            kernel_size=3,
            strides=1,
            padding="same",
            activation="relu",
            kernel_initializer="he_normal",
        )
        self.conv2 = Conv2D(
            self.filters,
            kernel_size=3,
            strides=1,
            padding="same",
            activation="relu",
            kernel_initializer="he_normal",
        )
        if self.pooling:
            self.pool = MaxPool2D(pool_size=(2, 2))
        if self.rate is not None:
            self.drop = Dropout(rate)

    def call(self, inputs):
        x = self.conv1(inputs)
        if self.rate is not None:
            x = self.drop(x)
        x = self.conv2(x)
        if self.pooling:
            y = self.pool(x)
            return y, x
        else:
            return x

    def get_config(self):
        base_config = super().get_config()
        return {
            **base_config,
            "filters": self.filters,
            "rate": self.rate,
            "pooling": self.pooling,
        }


class DecoderBlock(tf.keras.layers.Layer):
    def __init__(self, filters, rate=None, axis=-1, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.filters = filters
        self.rate = rate
        self.axis = axis
        self.convT = Conv2DTranspose(
            self.filters, kernel_size=3, strides=2, padding="same"
        )
        self.conv1 = Conv2D(
            self.filters,
            kernel_size=3,
            activation="relu",
            kernel_initializer="he_normal",
            padding="same",
        )
        if rate is not None:
            self.drop = Dropout(self.rate)
        self.conv2 = Conv2D(
            self.filters,
            kernel_size=3,
            activation="relu",
            kernel_initializer="he_normal",
            padding="same",
        )

    def call(self, inputs):
        X, short_X = inputs
        ct = self.convT(X)
        c_ = concatenate([ct, short_X], axis=self.axis)
        x = self.conv1(c_)
        if self.rate is not None:
            x = self.drop(x)
        y = self.conv2(x)
        return y

    def get_config(self):
        base_config = super().get_config()
        return {
            **base_config,
            "filters": self.filters,
            "rate": self.rate,
            "axis": self.axis,
        }


# Load the model with custom layers
unet = tf.keras.models.load_model(
    "final.h5",
    custom_objects={
        "EncoderBlock": EncoderBlock,
        "DecoderBlock": DecoderBlock,
    },
)


def show_image(image, cmap=None, title=None):
    plt.imshow(image, cmap=cmap)
    if title is not None:
        plt.title(title)
    plt.axis("off")


def predict(image):
    real_img = tf.image.resize(image, [128, 128])
    real_img = real_img / 255.0
    real_img = np.expand_dims(real_img, axis=0)
    pred_mask = unet.predict(real_img).reshape(128, 128)
    real_img = real_img[0]

    fig, ax = plt.subplots(1, 2, figsize=(10, 5))

    ax[0].imshow(real_img)
    ax[0].set_title("Original Image")
    ax[0].axis("off")

    ax[1].imshow(pred_mask, cmap="gray")
    ax[1].set_title("Predicted Mask")
    ax[1].axis("off")

    plt.tight_layout()
    plt.show()
    return pred_mask


# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="numpy"),
    examples=["./images/water_body_11.jpg", "./images/water_body_1011.jpg"],
    title="Water Body Segmentation",
)

iface.launch()