# -*- coding: utf-8 -*-
"""Funcional WLAN_design_gradio.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1MIfY3UkK4eSXOiPx3gtMSPoZMtnyJxux
"""

# from google.colab import drive
# drive.mount('/content/drive')

# Commented out IPython magic to ensure Python compatibility.
# %%capture
# !pip install gradio

import gradio as gr
from PIL import Image
import os
from tensorflow import keras
from keras.models import load_model
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from io import BytesIO

# Images path
path_main = 'Data/'
images_file = path_main + 'Scennarios init/Scennarios W'

# Load DL models
modelo_1ap = load_model(path_main + 'Models/modelo_1ap_app.keras')
modelo_2ap = load_model(path_main + 'Models/modelo_2ap_app.keras')

plt.rc('font', family='Times New Roman')
fontsize_t = 15

def coordinates_process(texto):
    coordinates = texto.split("), ")

    resultado = []
    for coord in coordinates:
        try:
            coord = coord.replace("(", "").replace(")", "")
            x, y = map(int, coord.split(","))
            # Validate range
            if 0 <= x <= 255 and 0 <= y <= 255:
                resultado.append((x, y))
            else:
                return False
        except ValueError:
            return False

    while len(resultado) < 3:
        resultado.append((0, 0))

    return resultado

# plan images path
def plan_images_list():
    return [file_ for file_ in os.listdir(images_file) if file_.endswith((".JPG", ".jpg", ".jpeg", ".png"))]

# Valdate inputs
def validate_input(value):
    if value == "" or value is None:
        return 0
    elif value >= 0 or value <= 2:
        return value

# MAIN FUNCTION ****************************************************************
def main_function(plan_name, apch1, apch6, apch11, coord1, coord6, coord11):
    image_plan_path = os.path.join(images_file, plan_name)
    imagen1 = Image.open(image_plan_path)

    # No negative number as input
    if not (0 <= apch1 <= 2):
        return False
    if not (0 <= apch6 <= 2):
        return False
    if not (0 <= apch11 <= 2):
        return False

    # Some variables init
    deep_count = 0
    deep_coverage = []
    channels = [1, 6, 11]
    num_APs = np.zeros(len(channels), dtype=int)
    num_APs[0] = apch1
    num_APs[1] = apch6
    num_APs[2] = apch11
    dimension = 256
    aps_chs = np.zeros((dimension, dimension, len(channels)))

    # Load plan
    numero = plan_name[:1]
    plan_in = np.array(Image.open(f"{path_main}Scennarios init/Scennarios B/{numero}.png")) / 255

    coords = [coord1, coord6, coord11]
    for att, channel in enumerate(channels):
      if num_APs[att] > 0:
          coordinates = coordinates_process(coords[att])
          for x, y in coordinates:
            if x != 0 and y != 0:
              aps_chs[int(y), int(x), att] = 1

    # Coverage process
    deep_coverage = []
    ap_images = []
    layer_indices = []
    imagencober = {}
    for k in range(len(channels)):
        capa = aps_chs[:, :, k]
        filas, columnas = np.where(capa == 1)

        if len(filas) == 2:
            # For 2 AP
            deep_count += 1
            layer_1 = np.zeros_like(capa)
            layer_2 = np.zeros_like(capa)
            layer_1[filas[0], columnas[0]] = 1
            layer_2[filas[1], columnas[1]] = 1

            datos_entrada = np.stack([plan_in, layer_1, layer_2], axis=-1)
            prediction = modelo_2ap.predict(datos_entrada[np.newaxis, ...])[0]

        elif len(filas) == 1:
            # For 1 AP
            deep_count += 1
            layer_1 = np.zeros_like(capa)
            layer_1[filas[0], columnas[0]] = 1

            datos_entrada = np.stack([plan_in, layer_1], axis=-1)
            prediction = modelo_1ap.predict(datos_entrada[np.newaxis, ...])[0]

        else:
            # Whitout AP
            prediction = np.zeros((dimension,dimension,1))

        # print(prediction.shape)
        deep_coverage.append(prediction)
        prediction_rgb = np.squeeze((Normalize()(prediction)))
        ap_images.append(prediction_rgb)  # Guardar la imagen de cobertura del AP

        if np.all(prediction == 0):
            plt.imshow(prediction_rgb)
            plt.title('No coverage', fontsize=fontsize_t + 2, family='Times New Roman')
            plt.axis("off")
        else:
            plt.imshow(prediction_rgb, cmap='jet')
            cbar = plt.colorbar(ticks=np.linspace(0, 1, num=6),)
            cbar.set_label('SINR [dB]', fontsize=fontsize_t, fontname='Times New Roman')
            cbar.set_ticklabels(['-3.01', '20.29', '43.60', '66.90', '90.20', '113.51'])
            cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
            plt.axis("off")

        # Save the plot to a buffer
        buf = BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        plt.close()

        # Convert buffer to an image
        imagencober[k] = Image.open(buf)

        # Cell map estimation
        layer_indices.append(np.argmax(prediction, axis=0))

    # Final coverage
    if deep_coverage:
        deep_coverage = np.array(deep_coverage)
        nor_matrix = np.max(deep_coverage, axis=0)
        celdas = np.argmax(deep_coverage, axis=0)

        resultado_rgb = np.squeeze((Normalize()(nor_matrix)))

        plt.imshow(resultado_rgb, cmap='jet')
        cbar = plt.colorbar(ticks=np.linspace(0, 1, num=6))
        cbar.set_label('SINR [dB]', fontsize=fontsize_t, fontname='Times New Roman')
        cbar.set_ticklabels(['-3.01', '20.29', '43.60', '66.90', '90.20', '113.51'])
        cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
        plt.axis("off")

        # Save the plot to a buffer
        buf = BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        plt.close()

        # Convert buffer to an image
        imagen3 = Image.open(buf)

        if num_APs[0] > 0 and num_APs[1] > 0 and num_APs[2] > 0:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red', 'green'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1, 2])
            cbar.set_ticklabels(['1', '6', '11'])
            cbar.set_label('Cell ID', fontsize=fontsize_t, fontname='Times New Roman')
            cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
            plt.axis("off")

            # Save the plot to a buffer
            buf = BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            plt.close()

            # Convert buffer to an image
            imagen4 = Image.open(buf)

        elif num_APs[0] > 0 and num_APs[1] > 0:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1])
            cbar.set_ticklabels(['1', '6'])
            cbar.set_label('Cell ID', fontsize=fontsize_t, fontname='Times New Roman')
            cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
            plt.axis("off")

            # Save the plot to a buffer
            buf = BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            plt.close()

            # Convert buffer to an image
            imagen4 = Image.open(buf)

        elif num_APs[0] > 0 and num_APs[2] > 0:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1])
            cbar.set_ticklabels(['1', '11'])
            cbar.set_label('Cell ID', fontsize=fontsize_t, fontname='Times New Roman')
            cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
            plt.axis("off")

            # Save the plot to a buffer
            buf = BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            plt.close()

            # Convert buffer to an image
            imagen4 = Image.open(buf)

        elif num_APs[1] > 0 and num_APs[2] > 0:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1])
            cbar.set_ticklabels(['6', '11'])
            cbar.set_label('Cell ID', fontsize=fontsize_t, fontname='Times New Roman')
            cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
            plt.axis("off")

            # Save the plot to a buffer
            buf = BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            plt.close()

            # Convert buffer to an image
            imagen4 = Image.open(buf)

        else:
            cmap = plt.cm.colors.ListedColormap(['blue'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0])
            cbar.set_ticklabels(['1'])
            cbar.set_label('Cell ID', fontsize=fontsize_t, fontname='Times New Roman')
            cbar.ax.tick_params(labelsize=fontsize_t, labelfontfamily = 'Times New Roman')
            plt.axis("off")

            # Save the plot to a buffer
            buf = BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            plt.close()

            # Convert buffer to an image
            imagen4 = Image.open(buf)

    return [imagencober[0], imagencober[1], imagencober[2], imagen3, imagen4]

# plan visualization
def load_plan_vi(mapa_seleccionado):

    image_plan_path1 = os.path.join(images_file, mapa_seleccionado)
    plan_image = Image.open(image_plan_path1)

    plan_n = np.array(plan_image.convert('RGB'))
    plt.figure(figsize=(3, 3))
    plt.imshow(plan_n)
    plt.xticks(np.arange(0, 256, 50))
    plt.yticks(np.arange(0, 256, 50))

    # Save the plot to a buffer
    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()

    # Convert buffer to an image
    plan_im = Image.open(buf)

    return plan_im

with gr.Blocks() as demo:

    gr.Markdown("""
    ## Fast Radio Propagation Prediction in WLANs Using Deep Learning
    This app use deep learning models to radio map estimation (RME). RME entails estimating the received RF power based on spatial information maps.

    Instructions for use:

    - A predefined list of indoor floor plans is available for use.
    - A maximum of (255,255) pixels is allowed for image size.
    - Negative numbers are not allowed.
    - The established format for the coordinates of each access point (AP) must be maintained.
    - A maximum of 2 APs are allowed per channel.
    """)

    with gr.Row():
        # Input left panel
        with gr.Column(scale=1):  # Scale to resize
            map_dropdown = gr.Dropdown(choices=plan_images_list(), label="Select indoor plan")
            ch1_input = gr.Number(label="APs CH 1")
            ch6_input = gr.Number(label="APs CH 6")
            ch11_input = gr.Number(label="APs CH 11")
            coords_ch1_input = gr.Textbox(label="Coordinate CH 1", placeholder="Format: (x1, y1), (x2, y2)")
            coords_ch6_input = gr.Textbox(label="Coordinate CH 6", placeholder="Format: (x1, y1), (x2, y2)")
            coords_ch11_input = gr.Textbox(label="Coordinate CH 11", placeholder="Format: (x1, y1), (x2, y2)")
            button1 = gr.Button("Load plan")
            button2 = gr.Button("Predict coverage")

        # Rigth panel
        a_images = 320 # Size putput images

        with gr.Column(scale=3):
            first_image_output = gr.Image(label="plan image", height=a_images, width=a_images)
            # with gr.Row():
            with gr.Row():
                image_ch1 = gr.Image(label="CH 1 coverage", height=a_images, width=a_images)
                image_ch6 = gr.Image(label="CH 6 coverage", height=a_images, width=a_images)
                image_ch11 = gr.Image(label="CH 11 coverage", height=a_images, width=a_images)

            with gr.Row():
                image_cover_final = gr.Image(label="Final coverage", height=a_images, width=a_images)
                image_cells = gr.Image(label="Cells coverage", height=a_images, width=a_images)


    # Buttons
    button1.click(load_plan_vi, inputs=[map_dropdown], outputs=first_image_output)
    button2.click(
    lambda map_dropdown, ch1, ch6, ch11, coords1, coords6, coords11: main_function(
        map_dropdown,
        validate_input(ch1),
        validate_input(ch6),
        validate_input(ch11),
        coords1,
        coords6,
        coords11,
    ),
    inputs=[map_dropdown, ch1_input, ch6_input, ch11_input, coords_ch1_input, coords_ch6_input, coords_ch11_input],
    outputs=[image_ch1, image_ch6, image_ch11, image_cover_final, image_cells]
)


demo.launch()