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
import re

# import gradio
# import PIL
# import tensorflow as tf
# import numpy as np
# import matplotlib

# print("Gradio:", gradio.__version__)
# print("Pillow (PIL):", PIL.__version__)
# print("TensorFlow:", tf.__version__)
# print("NumPy:", np.__version__)
# print("Matplotlib:", matplotlib.__version__)

# 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')

fontsize_t = 15

# Plan visualization
def load_plan_vi(mapa_seleccionado, uploaded_file):

    if mapa_seleccionado == "Upload your own image" and uploaded_file is not None:
        plan_image = Image.open(uploaded_file.name)
    elif mapa_seleccionado == "Upload your own image" and uploaded_file is None:
        # raise gradio.Warning("Deafult plan. Image was not loaded 😒.", duration=5)
        
        image_plan_path1 = os.path.join(images_file, "80.JPG")
        plan_image = Image.open(image_plan_path1)
    else:
        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.imshow(plan_n)
    plt.xticks(np.arange(0, 256, 50), fontsize=fontsize_t)
    plt.yticks(np.arange(0, 256, 50), fontsize=fontsize_t)
    plt.xlabel("X Coordinate [Pixels]", fontsize=fontsize_t)
    plt.ylabel("Y Coordinate [Pixels]", fontsize=fontsize_t)
    plt.show()

    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()

    plan_im = Image.open(buf)

    tick_positions = np.linspace(0, 256, 11)
    tick_labels = ["0", "2", "4", "6", "8", "10", "12", "14", "16", "18", "20"]

    plt.xticks(tick_positions, tick_labels)

    buf = BytesIO()
    plt.imshow(plan_n)
    plt.xticks(tick_positions, tick_labels, fontsize=fontsize_t)
    plt.yticks(tick_positions, tick_labels, fontsize=fontsize_t)
    plt.xlabel("X Coordinate [meters]", fontsize=fontsize_t)
    plt.ylabel("Y Coordinate [meters]", fontsize=fontsize_t)
    plt.show()

    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()

    plan_im_meters = Image.open(buf)

    return plan_im, plan_im_meters

def validate_coords(num_aps, coords):
    matches = re.findall(r"\(\s*\d+\s*,\s*\d+\s*\)", coords)

    if len(matches) > int(num_aps):
        new_coords = ", ".join(matches[:int(num_aps)])
        return new_coords
    return coords

import re

def coordinates_process(texto, interference):
    a = False

    texto = re.sub(r'\s*,\s*', ', ', texto)  # Normalizar espacios entre comas
    texto = re.sub(r'\)\s*,\s*\(', '), (', texto)  # Asegurarse de que las tuplas estén separadas por "), ("
    texto = texto.strip()  # Eliminar cualquier espacio en exceso al principio o al final

    coordinates = texto.split("), ")

    resultado = []
    for coord in coordinates:
        try:
            coord = coord.replace("(", "").replace(")", "")  # Eliminar paréntesis
            x, y = map(int, coord.split(","))  # Convertir a enteros

            if 0 <= x <= 255 and 0 <= y <= 255:
                resultado.append((x, y))
            else:
                a = True
        except ValueError:
            a = True

    if a:
        resultado = [(0, 0), (0, 0), (0, 0)]

    while len(resultado) < 3 and interference == True:
        resultado.append((0, 0))
    while len(resultado) < 5 and interference == False:
        resultado.append((0, 0))

    return resultado  # Devolver como arreglo de tuplas

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

# MAIN FUNCTION ****************************************************************
def main_function(plan_name, uploaded_file, interference = True, aps_int = 0, aps_coor = '(0,0)',
                  apch1  = 0, apch6 = 0, apch11 = 0, coord1 = '(0,0)', coord6 = '(0,0)', coord11 = '(0,0)'):

    plan_name = str(plan_name)
    interference = bool(interference)
    aps_int = int(aps_int)
    aps_coor = str(aps_coor)
    apch1 = int(apch1)
    apch6 = int(apch6)
    apch11 = int(apch11)
    coord1 = str(coord1)
    coord6 = str(coord6)
    coord11 = str(coord11)

    aps_coor = validate_coords(aps_int, aps_coor)
    coord1 = validate_coords(apch1, coord1)
    coord6 = validate_coords(apch6, coord6)
    coord11 = validate_coords(apch11, coord11)

    # **************************************************************************
    imagencober = {}

    prediction_rgb = np.zeros((256, 256))

    for k in range(5):
        plt.imshow(prediction_rgb, cmap='gray')
        plt.title(f'No coverage', fontsize=fontsize_t + 2)
        plt.axis("off")
        # plt.show()

        buf = BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        plt.close()

        imagencober[k] = Image.open(buf)
    # **************************************************************************

    # Load plan
    if plan_name == "Upload your own image" and uploaded_file is not None:
        plan_image = np.array(Image.open(uploaded_file.name))/255

        dimensiones = plan_image.shape
        if len(dimensiones) > 2:
            # raise gradio.Error("Error in dimensions of the uploaded image. Must [256,256,3] 💣🎆.", duration=5)
            raise ValueError("Error in image size. Must [256,256].")
        
        plan_grayscale = plan_image[:, :, 0]
        plan_in = 1 - plan_grayscale

    elif plan_name == "Upload your own image" and uploaded_file is None:
        # raise gradio.Warning("Deafult plan processed. Image was not loaded 😒.", duration=5)
        numero = "80"
        plan_in = np.array(Image.open(f"{path_main}Scennarios init/Scennarios B/{numero}.png")) / 255

    else:
        numero = plan_name.split('.')[0]
        plan_in = np.array(Image.open(f"{path_main}Scennarios init/Scennarios B/{numero}.png")) / 255

    # Some variables init
    deep_count = 0
    deep_coverage = []
    dimension = 256

    if interference:
        # if apch1 == 0 and apch6 == 0 and apch11 == 0:
        #     raise gradio.Warning("There are not APs for estimation 😒.", duration=5)
            
        channels_c = [1, 6 , 11]
        channels = 3
        num_APs = np.zeros(channels, dtype=int)
        num_APs[0] = apch1
        num_APs[1] = apch6
        num_APs[2] = apch11
        aps_chs = np.zeros((dimension, dimension, channels))
        coords = [coord1, coord6, coord11]

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

    if not interference:
        channels = aps_int
        aps_chs = np.zeros((dimension, dimension, 5))  # Crear la matriz
        coordinates = coordinates_process(aps_coor, interference)

        for att, (x, y) in enumerate(coordinates):
            if x != 0 and y != 0:
                aps_chs[int(y), int(x), att] = 1

    # Coverage process
    deep_coverage = []
    ap_images = []
    layer_indices = []

    for k in range(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)

        if np.all(prediction == 0):
            plt.imshow(prediction_rgb, cmap='gray')
            plt.title(f'No coverage', fontsize=fontsize_t)
            plt.axis("off")
            plt.show()
        else:
            plt.imshow(prediction_rgb, cmap='jet')
            if interference:
                plt.title(f'Coverage CH {channels_c[k]}', fontsize=fontsize_t)
                cbar = plt.colorbar(ticks=np.linspace(0, 1, num=6),)
                cbar.set_label('SINR [dB]', fontsize=fontsize_t)
                cbar.set_ticklabels(['-3.01', '20.29', '43.60', '66.90', '90.20', '113.51'])
            if not interference:
                plt.title(f'Coverage AP {k}', fontsize=fontsize_t)
                cbar = plt.colorbar(ticks=np.linspace(0, 1, num=6),)
                cbar.set_label('Power [dBm]', fontsize=fontsize_t)
                cbar.set_ticklabels(['-94.94', '-70.75', '-46.56', '-22.38', '1.81', '26.00'])
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

        # 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)

    # 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))
        if interference:
            cbar.set_label('SINR [dB]', fontsize=fontsize_t)
            cbar.set_ticklabels(['-3.01', '20.29', '43.60', '66.90', '90.20', '113.51'])
        if not interference:
            cbar.set_label('Power [dBm]', fontsize=fontsize_t)
            cbar.set_ticklabels(['-94.94', '-70.75', '-46.56', '-22.38', '1.81', '26.00'])
        cbar.ax.tick_params(labelsize=fontsize_t)
        plt.axis("off")
        plt.show()

        # 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 interference == True:
        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)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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)

    # **************************************************************************

    if interference == False:
        if aps_int == 5:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red', 'green', 'yellow', 'violet'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1, 2, 3, 4])
            cbar.set_ticklabels(['1', '2', '3', '4', '5'])
            cbar.set_label('Cell ID', fontsize=fontsize_t)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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 aps_int == 4:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red', 'green', 'yellow'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1, 2, 3])
            cbar.set_ticklabels(['1', '2', '3', '4'])
            cbar.set_label('Cell ID', fontsize=fontsize_t)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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 aps_int == 3:
            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', '2', '3'])
            cbar.set_label('Cell ID', fontsize=fontsize_t)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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 aps_int == 2:
            cmap = plt.cm.colors.ListedColormap(['blue', 'red'])
            plt.imshow(celdas, cmap=cmap)
            cbar = plt.colorbar()
            cbar.set_ticks([0, 1])
            cbar.set_ticklabels(['1', '2'])
            cbar.set_label('Cell ID', fontsize=fontsize_t)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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)
            cbar.ax.tick_params(labelsize=fontsize_t)
            plt.axis("off")
            plt.show()

            # 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], imagencober[3], imagencober[4], imagen3, imagen4]

def update_interface(enable_interference):
    if enable_interference:
        return {
            map_dropdown : gr.update(visible=True),
            upload_image : gr.update(visible=True),
            ch1_input: gr.update(visible=True),
            ch6_input: gr.update(visible=True),
            ch11_input: gr.update(visible=True),
            coords_ch1_input: gr.update(visible=True),
            coords_ch6_input: gr.update(visible=True),
            coords_ch11_input: gr.update(visible=True),
            button1: gr.update(visible=True),
            button2: gr.update(visible=True),
            image_ap1 : gr.update(visible=False),
            image_ap2 : gr.update(visible=False),
            image_ch1 : gr.update(visible=True),
            image_ch6 : gr.update(visible=True),
            image_ch11 : gr.update(visible=True),
            simple_dropdown: gr.update(visible=False),
            simple_coords: gr.update(visible=False)
        }
    else:
        return {
            map_dropdown: gr.update(visible=True),
            upload_image : gr.update(visible=True),
            ch1_input: gr.update(visible=False),
            ch6_input: gr.update(visible=False),
            ch11_input: gr.update(visible=False),
            coords_ch1_input: gr.update(visible=False),
            coords_ch6_input: gr.update(visible=False),
            coords_ch11_input: gr.update(visible=False),
            simple_dropdown: gr.update(visible=True, interactive=True),
            simple_coords: gr.update(visible=True, interactive=True),
            image_ap1 : gr.update(visible=True),
            image_ap2 : gr.update(visible=True),
            image_ch1 : gr.update(visible=True),
            image_ch6 : gr.update(visible=True),
            image_ch11 : gr.update(visible=True)
        }

with gr.Blocks() as demo:
    gr.Markdown(
    """
    ## Fast Indoor Radio Propagation Prediction using Deep Learning
    This app uses deep learning models for radio map estimation (RME) with and without interference, simulating 2.4 GHz and 5 GHz bands. RME involves estimating the received RF power based on spatial information maps.
    
    Instructions for use:

    - A predefined list of indoor floor plans is available for use.
    - You can upload your own indoor floor plan.
    - Negative numbers are not allowed.
    - The established format for the coordinates of each access point (AP) must be followed.
    - A maximum of 2 APs per channel is allowed for the interference case.
    - A maximum of 5 APs is allowed for the non-interference case.
    - The uploaded plan must meet the dimensions [256,256], with free spaces as white pixels and walls as black pixels.
    """
    )

    enable_interference = gr.Checkbox(label="Enable Interference Analysis", value=True)

    with gr.Row():
        with gr.Column(scale=1):
            map_dropdown = gr.Dropdown(choices=plan_images_list() + ["Upload your own image"], label="Select indoor plan", value="80.JPG")
            upload_image = gr.File(label="Or upload your own image", file_types=[".JPG", ".jpg", ".jpeg", ".png", ".PNG"])
            ch1_input = gr.Dropdown(choices=[i for i in range(0, 3)], label="Select APs CH 1", value=0)
            ch6_input = gr.Dropdown(choices=[i for i in range(0, 3)], label="Select APs CH 6", value=0)
            ch11_input = gr.Dropdown(choices=[i for i in range(0, 3)], label="Select APs CH 11", value=0)
            coords_ch1_input = gr.Textbox(label="Coordinate CH 1", placeholder="Format (Pixels): (x1, y1), (x2, y2)")
            coords_ch6_input = gr.Textbox(label="Coordinate CH 6", placeholder="Format (Pixels): (x1, y1), (x2, y2)")
            coords_ch11_input = gr.Textbox(label="Coordinate CH 11", placeholder="Format (Pixels): (x1, y1), (x2, y2)")

            simple_dropdown = gr.Dropdown(choices=[str(i) for i in range(1, 6)], label="Select APs number", visible=False)
            simple_coords = gr.Textbox(label="Enter APs coordinates", placeholder="Format (Pixels): (x1, y1), (x1, y1)...", visible=False,)

            button1 = gr.Button("Load plan")
            button2 = gr.Button("Predict coverage")

        with gr.Column(scale=3):
            with gr.Row():
                first_image_output = gr.Image(label="Plan image pixels")
                second_image_output = gr.Image(label="Plan image meters")
            with gr.Row():
                image_ch1 = gr.Image(label="Coverage 1")
                image_ch6 = gr.Image(label="Coverage 2")
                image_ch11 = gr.Image(label="Coverage 3")

                image_ap1 = gr.Image(label="Coverage 4", visible=False)
                image_ap2 = gr.Image(label="Coverage 5", visible=False)
            with gr.Row():
                image_cover_final = gr.Image(label="Final coverage")
                image_cells = gr.Image(label="Cells coverage")

    enable_interference.change(update_interface, inputs=[enable_interference],
                               outputs=[map_dropdown, upload_image,
                                        ch1_input, ch6_input, ch11_input,
                                        coords_ch1_input, coords_ch6_input, coords_ch11_input,
                                        button1, button2,
                                        simple_dropdown, simple_coords,
                                        image_ch1, image_ch6, image_ch11,
                                        image_ap1, image_ap2])

    button1.click(load_plan_vi,
                  inputs=[map_dropdown, upload_image], outputs=[first_image_output, second_image_output])

    # Único bloque para el clic de button2
    button2.click(main_function,
                  inputs=[map_dropdown,
                          upload_image,
                          enable_interference,
                          simple_dropdown,
                          simple_coords,
                          ch1_input, ch6_input, ch11_input,
                          coords_ch1_input, coords_ch6_input, coords_ch11_input],
                  outputs=[image_ch1, image_ch6, image_ch11, image_ap1, image_ap2, image_cover_final, image_cells])

demo.launch()