import cv2 as cv
import numpy as np
import gradio as gr
from mobilenet import MobileNet
from huggingface_hub import hf_hub_download

# Download ONNX model from Hugging Face
model_path = hf_hub_download(repo_id="opencv/image_classification_mobilenet", filename="image_classification_mobilenetv1_2022apr.onnx")
top_k = 10
backend_id = cv.dnn.DNN_BACKEND_OPENCV
target_id = cv.dnn.DNN_TARGET_CPU

# Load MobileNet model
model = MobileNet(modelPath=model_path, topK=top_k, backendId=backend_id, targetId=target_id)

def add_hsv_noise(image, hue_noise=0, saturation_noise=0, value_noise=0):
    """Add HSV noise to an image"""
    if image is None:
        return None
    
    # Convert BGR to HSV (OpenCV uses BGR by default)
    hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV).astype(np.float32)
    
    # Add noise to each channel
    hsv[:, :, 0] = np.clip(hsv[:, :, 0] + hue_noise, 0, 179)  # Hue: 0-179
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] + saturation_noise, 0, 255)  # Saturation: 0-255
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] + value_noise, 0, 255)  # Value: 0-255
    
    # Convert back to BGR
    bgr = cv.cvtColor(hsv.astype(np.uint8), cv.COLOR_HSV2BGR)
    
    return bgr

def classify_image_with_noise(input_image, top_n, hue_noise, saturation_noise, value_noise):
    """Classify image with HSV noise applied and return exact confidence scores"""
    if input_image is None:
        return None, "Please upload an image first."
    
    # Apply HSV noise
    noisy_image = add_hsv_noise(input_image, hue_noise, saturation_noise, value_noise)
    
    # Resize and crop as in original code
    image = cv.resize(noisy_image, (256, 256))
    image = image[16:240, 16:240, :]
    
    # Preprocess manually to get raw scores
    input_blob = model._preprocess(image)
    
    # Forward pass
    model.model.setInput(input_blob, model.input_names)
    output_blob = model.model.forward(model.output_names)
    
    # Get raw probabilities (apply softmax if needed)
    raw_scores = output_blob[0]  # First batch
    probabilities = np.exp(raw_scores) / np.sum(np.exp(raw_scores))  # Softmax
    
    # Get top N indices and their scores
    top_indices = np.argsort(probabilities)[::-1][:top_n]
    
    # Format results with exact confidence scores
    result_lines = []
    for i, idx in enumerate(top_indices):
        label = model._labels[idx]
        confidence = probabilities[idx]
        result_lines.append(f"{i+1}. {label}: {confidence:.6f} ({confidence*100:.4f}%)")
    
    result_str = "\n".join(result_lines)
    
    # Convert BGR to RGB for display in Gradio
    display_image = cv.cvtColor(noisy_image, cv.COLOR_BGR2RGB)
    
    return display_image, result_str

def clear_output_on_change(img):
    return gr.update(value=""), None

def clear_all():
    return None, None, ""

with gr.Blocks(css='''.example * {
    font-style: italic;
    font-size: 18px !important;
    color: #0ea5e9 !important;
    }''') as demo:

    gr.Markdown("### Image Classification with MobileNet + HSV Noise Analysis")
    gr.Markdown("Upload an image and adjust HSV noise sliders to see how it affects MobileNet predictions in real-time.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Upload Image")
            
            gr.Markdown("### Classification Settings")
            top_n = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Top N Classes")
            
            gr.Markdown("### HSV Noise Controls")
            hue_noise = gr.Slider(minimum=-50, maximum=50, value=0, step=1, label="Hue Noise (-50 to 50)")
            saturation_noise = gr.Slider(minimum=-100, maximum=100, value=0, step=5, label="Saturation Noise (-100 to 100)")
            value_noise = gr.Slider(minimum=-100, maximum=100, value=0, step=5, label="Value/Brightness Noise (-100 to 100)")
        
        with gr.Column():
            noisy_image_output = gr.Image(label="Image with Noise Applied")
            output_box = gr.Textbox(label="Top Predictions with Confidence Scores", lines=10, max_lines=15)

    image_input.change(fn=clear_output_on_change, inputs=image_input, outputs=[output_box, noisy_image_output])

    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
        clear_btn = gr.Button("Clear")

    inputs = [image_input, top_n, hue_noise, saturation_noise, value_noise]
    outputs = [noisy_image_output, output_box]
    
    for slider in [top_n, hue_noise, saturation_noise, value_noise]:
        slider.change(fn=classify_image_with_noise, inputs=inputs, outputs=outputs)
    
    submit_btn.click(fn=classify_image_with_noise, inputs=inputs, outputs=outputs)
    clear_btn.click(fn=clear_all, outputs=[image_input, noisy_image_output, output_box])

    gr.Markdown("Click on any example to try it.", elem_classes=["example"])

    gr.Examples(
        examples=[
            ["examples/squirrel_cls.jpg"],
            ["examples/baboon.jpg"]
        ],
        inputs=image_input
    )

if __name__ == "__main__":
    demo.launch()