import time

import gradio as gr
import numpy as np
from pathlib import Path
import time
from anomalib.deploy import OpenVINOInferencer
from openvino.runtime import Core
# Initialize the Core
core = Core() 

# Get the available devices
devices = core.available_devices

inferencer = None

example_list = [["bottle/examples/000.png", "anomaly_map", "bottle", "CPU"],
               ["pill/examples/010.png", "heat_map", "pill", "CPU"],
               ["zipper/examples/001.png", "pred_mask", "zipper", "CPU"],
               ["grid/examples/005.png", "segmentations", "grid", "CPU"],
               ["cubes/examples/005.jpg", "heat_map", "cubes", "CPU"]]

def OV_compilemodel(category_choice, device):
    global inferencer
    #Get the available models
    openvino_model_path = Path.cwd() / category_choice / "run" / "weights" / "openvino" / "model.bin"
    metadata_path = Path.cwd() / category_choice / "run" / "weights" / "openvino" / "metadata.json"
    
    inferencer = OpenVINOInferencer(
        path=openvino_model_path,  # Path to the OpenVINO IR model.
        metadata_path=metadata_path,  # Path to the metadata file.
        device=device,  # We would like to run it on an Intel CPU.
        config= {"INFERENCE_PRECISION_HINT": "f16" } if device != "CPU" else {}
    )
    
    return inferencer

def OV_inference(input_img, operation, category_choice, device):
         
    start_time = time.time()
    predictions = inferencer.predict(image=input_img)
    stop_time = time.time()
    inference_time = stop_time - start_time
    confidence = predictions.pred_score
    
    if operation == "original":
        output_img1 = predictions.image
    elif operation == "anomaly_map":
        output_img1 = predictions.anomaly_map
    elif operation == "heat_map":
        output_img1 = predictions.heat_map
    elif operation == "pred_mask":
        output_img1 = predictions.pred_mask
    elif operation == "segmentations":
        output_img1 = predictions.segmentations
    else: 
        output_img1 = predictions.image
    return output_img1, round(inference_time*1000), round(confidence*100,2)

with gr.Blocks() as demo:
    gr.Markdown(
    """
    <img align="left" width="150" src= "https://github.com/openvinotoolkit/anomalib/assets/10940214/7e61a627-d1b0-4ad4-b602-da9b348c0cbe">   
    <img align="right" width="150" src= "https://github.com/openvinotoolkit/anomalib/assets/10940214/5d6dd038-b40c-441f-ad38-1cf526137de2">   
    
    <h1 align="center"> 🚀 Anomaly detection 🚀 </h1>  
    
    Experience the power of the state-of-the-art anomaly detection with Anomalib-OpenVINO Anomaly detection toolbox. This interactive APP leverages the robust capabilities of Anomalib and OpenVINO.
    
    All model are FP32 precision, if you select GPU it will automatically change precision to FP16. Using Anomalib you can also quantize your model in INT8 using NNCF.
    
    
     
    ![](https://github.com/openvinotoolkit/anomalib/assets/10940214/ce78346f-4d27-4f99-bea7-75b87e2ac02a)
    
    
    
    """
    )
    
    gr.Markdown("## 1. Select the category over you want to detect anormalities.")
    category_choice = gr.Radio(["bottle", "grid", "pill", "zipper", "cubes"], label="Choose the category")
    
    gr.Markdown(
        """
        ## 2. Select the Intel device
        Device Name   | CPU  | GPU.0  | GPU.1 
        ------------- | ------------  |------------- | -------------
        Intel Device  | CPU  |  Integrated GPU  | Discrete GPU  
        
        
        """
    )
    device_choice = gr.Dropdown(devices, label="Choose the device")
    
    gr.Markdown("## 3. Compile the model")
    compile_btn = gr.Button("Compile Model")

    gr.Markdown("## 4. Choose the output you want to visualize.")
    output_choice = gr.Radio(["original", "anomaly_map", "heat_map", "pred_mask", "segmentations"], label="Choose the output")
    
    gr.Markdown("## 5. Drop the image in the input image box and run the inference")
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="numpy", label= "Input image")
            
            
        with gr.Column():
            output_img = gr.outputs.Image(type="numpy", label="Anomalib Output")

    inference_btn = gr.Button("Run Inference")
    
    with gr.Row():
        # Create your output components
        #output_prediction = gr.Textbox(label="Prediction")
        output_confidence = gr.Textbox(label="Confidence [%]")
        output_time = gr.Textbox(label="Inference Time [ms]")
    
    gr.Markdown("Note: Change the image and run the inference again. If you want to change the object you need to recompile the model, that means you need to start from step 1.")    
    gr.Markdown("## Image Examples")
    
    gr.Examples(
        examples=example_list,
        inputs=[image, output_choice, category_choice, device_choice],
        outputs=[output_img, output_time, output_confidence],
        fn=OV_inference,
    )
    
    compile_btn.click(OV_compilemodel, inputs=[category_choice, device_choice])
    inference_btn.click(OV_inference, inputs=[image, output_choice], outputs=[output_img, output_time, output_confidence])

demo.launch(share=True, enable_queue=True)