import os
import onnxruntime
import numpy as np
import cv2
import gradio as gr
from huggingface_hub import hf_hub_download
import spaces


# Set your Hugging Face access token here
# HF_TOKEN = "your_hugging_face_token"  # Replace with your actual Hugging Face token
HF_TOKEN = os.environ.get("HF_TOKEN")


# Preprocessing function
def preprocess_image(image):
    input_size = (224, 224)  # Modify based on model's expected input
    image = cv2.resize(image, input_size)  # Resize the image
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert to RGB
    image = image.astype(np.float32) / 255.0  # Normalize
    image = np.transpose(image, (2, 0, 1))  # Change to CHW format
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Class Dictionary
cls_dict = {
    16: "2-products-in-one-offer-+-coupon",
    3: "2-products-in-one-offer",
    2: "an-offer-with-a-coupon",
    17: "availability-of-additional-products",
    21: "offers-with-a-preliminary-promotional-price",
    15: "offers-with-an-additional-deposit-price",
    4: "offers-with-an-additional-shipping",
    12: "offers-with-dealtype-special_price",
    7: "offers-with-different-sizes",
    14: "offers-with-money_rebate",
    6: "offers-with-percentage_rebate",
    18: "offers-with-price-characteristic-(statt)",
    22: "offers-with-price-characterization-(uvp)",
    13: "offers-with-product-number-(sku)",
    1: "offers-with-reward",
    24: "offers-with-the-condition:-available-from-such-and-such-a-number / offers-with-the-old-price-crossed-out",
    19: "offers-with-the-old-price-crossed-out",
    8: "regular",
    11: "scene-with-multiple-offers-+-uvp-price-for-each-offers",
    0: "several-products-in-one-offer-with-different-prices",
    5: "simple-offers",
    20: "stock-offers",
    10: "stocks",
    9: "travel-booklets",
    23: "with-a-product-without-a-price",
    25: "with-the-price-of-the-supplemental-bid"
}



session = None

# Inference function

@spaces.GPU()
def predict(file):

    image = cv2.imread(file.name)  # Use the file's path from Gradio

    
    global session

    if session is None:

        # Download ONNX model from Hugging Face Hub
        onnx_model_path = hf_hub_download(repo_id="limitedonly41/offers_26cls", 
                                           filename="offers_26_best_v2.onnx", 
                                           use_auth_token=HF_TOKEN)
        
        # Load ONNX model
        session = onnxruntime.InferenceSession(onnx_model_path)
        
    # Preprocess the image
    input_image = preprocess_image(image)

    # Get input name for the ONNX model
    input_name = session.get_inputs()[0].name

    # Perform inference
    outputs = session.run(None, {input_name: input_image})

    # Post-process the outputs
    list_values = list(outputs[0][0])
    max_index = list_values.index(max(list_values))
    
    # Return the predicted class label
    return cls_dict[max_index]

# Gradio interface
interface = gr.Interface(fn=predict, 
                         inputs=gr.File(label="Upload an Image"),  # Change to File input
                         outputs=gr.Label(num_top_classes=1, label="Predicted Class"),
                         title="ONNX Model Image Classification",
                         description="Upload an image and get the predicted class using the ONNX model.",
                         live=True)

# Launch the interface
interface.launch()