import numpy as np
import pandas as pd
import gradio as gr
import cv2
from tensorflow import keras as k

# ============================== variables ==============================
image_size = 256
num_classes = 3

in_channel_tool = 3
in_channel_spec = 9
in_channel_chip = 3
img_rows, img_cols = image_size, image_size

# ============================== paths ==============================
model_class_path = f"Models/siren_base_multi_tsc.h5"
model_reg_path = f"Models/siren_regression_mse_multi_tsc.h5"

# dataset
csv_path = "Dataset/labels_sample.csv"
csv_path_reg = "Dataset/labels_reg_sample.csv"
tool_path = "Dataset/tool"
spec_path = "Dataset/spec"
chip_path = "Dataset/chip"

# ============================== load/prepare dataset ==============================
df = pd.read_csv(csv_path)
df_reg = pd.read_csv(csv_path_reg)

df["tool"] = df.id.map(lambda id: f"{tool_path}/{id}.jpg")
df["spec_x"] = df.id.map(lambda id: f"{spec_path}/x/{id}.jpg")
df["spec_y"] = df.id.map(lambda id: f"{spec_path}/y/{id}.jpg")
df["spec_z"] = df.id.map(lambda id: f"{spec_path}/z/{id}.jpg")
df["chip"] = df.id.map(lambda id: f"{chip_path}/{id}.jpg")

# ============================== examples ===========================================
exs = []
for i in range(len(df)):
    row = df.iloc[i,:]
    tool_id = row.id
    image_label = row.image_label 
    tool = row.tool 
    spec_x = row.spec_x 
    spec_y = row.spec_y 
    spec_z = row.spec_z 
    chip = row.chip 

    task = "Regression"
    if i % 2 == 0:
        task = "Classification"
    example = [tool_id, image_label, task, tool, spec_x, spec_y, spec_z, chip]
    exs.append(example)

# ============================== preprocess ===========================================
def process_img(img, img_rows, img_cols, channels):
    """
    Reads the spectogram files from disk and normalizes the pixel values
      @params:
        img - Data of the image
        img_rows - The image height.
        img_cols - The image width.
        as_grey - Read the image as Greyscale or RGB.
        channels - Number of channels.
      @returns:
        The created and compiled model (Model)
    """
    img = cv2.imread(img)
    img = cv2.resize(img, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img = np.asarray(img, dtype=np.float32)

    # normalize
    #print(np.max(img))
    img = img / 255.0 #np.max(img)

    # reshape to match Keras expectaions
    img = img.reshape(img_rows, img_cols, channels)

    return img

def process_specs(img_x, img_y, img_z, img_rows, img_cols, channels):
    img_x = cv2.imread(img_x)
    img_y = cv2.imread(img_y)
    img_z = cv2.imread(img_z)

    img = []
    img_x = cv2.resize(img_x, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img_y = cv2.resize(img_y, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img_z = cv2.resize(img_z, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img.append([img_x, img_y, img_z])
    img = np.asarray(img, dtype=np.float32)

    # normalize
    #print(np.max(img))
    img = img / 255.0 #np.max(img)

    # reshape to match Keras expectaions
    img = img.reshape(img_rows, img_cols, channels)

    return img


# ============================== load model ==============================================
model_class = k.models.load_model(model_class_path, compile=False)
model_reg = k.models.load_model(model_reg_path, compile=False)
    
# ============================== app =====================================================

def change_output_labels(choice):
  if choice == "Classification":
    return [
        gr.Label(value=None, label="Actual Label", visible=True), gr.Label(value=None, label="Predicted Label", visible=True), 
        gr.Label(label="Actual Gaps", visible=False), gr.Label(label="Predicted Gaps", visible=False),
        gr.Label(label="Actual Flank Wear", visible=False), gr.Label(label="Predicted Flank Wear", visible=False),
        gr.Label(label="Actual Overhang", visible=False), gr.Label(label="Predicted Overhang", visible=False)
      ]
  else:
    return [
        gr.Label("Sharp", label="Actual Label", visible=False), gr.Label("Sharp", label="Predicted Label", visible=False), 
        gr.Label(value=None, label="Actual Gaps", visible=True), gr.Label(value=None, label="Predicted Gaps", visible=True),
        gr.Label(value=None, label="Actual Flank Wear", visible=True), gr.Label(value=None, label="Predicted Flank Wear", visible=True),
        gr.Label(value=None, label="Actual Overhang", visible=True), gr.Label(value=None, label="Predicted Overhang", visible=True)
      ]


def predict(tool_id, task, label, tool, spec_x, spec_y, spec_z, chip):
    if task is None:
      raise gr.Error("Choose a task first!")

    labels = ['sharp', 'used', 'dulled']
    tool = process_img(tool, img_rows, img_cols, in_channel_tool)
    spec = process_specs(spec_x, spec_y, spec_z, img_rows, img_cols, in_channel_spec)
    chip = process_img(chip, img_rows, img_cols, in_channel_chip)


    inputs = [np.array([tool,]), np.array([spec,]), np.array([chip,])]
    if task == "Classification":
      y_score = model_class.predict(inputs)
      y_pred = {label:float(score) for label, score in zip(labels, y_score[0])}
      return [
          gr.Label(value=label, label="Actual Label", visible=True), gr.Label(value=y_pred, label="Predicted Label", visible=True), 
          gr.Label(label="Actual Gaps", visible=False), gr.Label(label="Predicted Gaps", visible=False),
          gr.Label(label="Actual Flank Wear", visible=False), gr.Label(label="Predicted Flank Wear", visible=False),
          gr.Label(label="Actual Overhang", visible=False), gr.Label(label="Predicted Overhang", visible=False)
        ]
    else:
      y_score = model_reg.predict(inputs)
      gaps_pred = str(y_score[0][0])
      flank_wear_pred = str(y_score[0][1])
      overhang_pred = str(y_score[0][2])

      actual = df_reg[df_reg["id"] == tool_id].values
      gaps_actual = str(actual[0][1])
      flank_wear_actual = str(actual[0][2])
      overhang_actual = str(actual[0][3])
      return [
        gr.Label("Sharp", label="Actual Label", visible=False), gr.Label("Sharp", label="Predicted Label", visible=False), 
        gr.Label(value=gaps_actual, label="Actual Gaps", visible=True), gr.Label(value=gaps_pred, label="Predicted Gaps", visible=True),
        gr.Label(value=flank_wear_actual, label="Actual Flank Wear", visible=True), gr.Label(value=flank_wear_pred, label="Predicted Flank Wear", visible=True),
        gr.Label(value=overhang_actual, label="Actual Overhang", visible=True), gr.Label(value=overhang_pred, label="Predicted Overhang", visible=True)
      ]

# ============================== blocks ========================================================
title = r"""
<h1 align="center">Siren</h1>
"""
description = r"""
<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/hubtru/Siren' 
target='_blank'><b>SIREN: A Scalable Isotropic Recursive Column Multimodal Neural Architecture with Device State Recognition Use-Case</b></a>.<br>
"""

with gr.Blocks() as demo:
  gr.Markdown(value=title)
  gr.Markdown(description)
  with gr.Row():
    with gr.Column():
      with gr.Row():
        tool_id = gr.Textbox("T1R2B1", label="Tool")
        label_input = gr.Textbox("Sharp", label="Label")
      task = gr.Radio(["Classification", "Regression"], label="Task")
      with gr.Row():
        tool = gr.Image(label="Tool", type="filepath")
      with gr.Row():
        chip = gr.Image(label="Chip", type="filepath")
      with gr.Row():
        spec_x = gr.Image(label="Spec_x", type="filepath")
        spec_y = gr.Image(label="Spec_y", type="filepath")
        spec_z = gr.Image(label="Spec_z", type="filepath")
      submit_btn = gr.Button("Submit", variant="primary")
      
    with gr.Column():
      output_labels = [
        gr.Label("Sharp", label="Actual Label"),
        gr.Label("Sharp", label="Predicted Label"),
        gr.Label(label="Actual Gaps", visible=False), 
        gr.Label(label="Predicted Gaps", visible=False), 
        gr.Label(label="Actual Flank Wear", visible=False), 
        gr.Label(label="Predicted Flank Wear", visible=False), 
        gr.Label(label="Actual Overhang", visible=False), 
        gr.Label(label="Predicted Overhang", visible=False), 
      ]
  
  examples = gr.Examples(examples=exs, inputs=[tool_id, label_input, task, tool, spec_x, spec_y, spec_z, chip])
  task.change(fn=change_output_labels, inputs=task, outputs=output_labels)
  submit_btn.click(fn=predict, inputs=[tool_id, task, label_input, tool, spec_x, spec_y, spec_z, chip], outputs=output_labels)
demo.launch()