import gradio as gr
from gender_classification import gender_classification
from emotion_classification import emotion_classification
from dog_breed import dog_breed_classification
from deepfake_vs_real import deepfake_classification
from gym_workout_classification import gym_workout_classification
from augmented_waste_classifier import waste_classification

# Functions to update the model state when a button is clicked.
def select_gender():
    return "gender"

def select_emotion():
    return "emotion"

def select_dog_breed():
    return "dog breed"

def select_deepfake():
    return "deepfake"

def select_gym_workout():
    return "gym workout"

def select_waste():
    return "waste"

# Main classification function that calls the appropriate model based on selection.
def classify(image, model_name):
    if model_name == "gender":
        return gender_classification(image)
    elif model_name == "emotion":
        return emotion_classification(image)
    elif model_name == "dog breed":
        return dog_breed_classification(image)
    elif model_name == "deepfake":
        return deepfake_classification(image)
    elif model_name == "gym workout":
        return gym_workout_classification(image)
    elif model_name == "waste":
        return waste_classification(image)
    else:
        return {"Error": "No model selected"}

with gr.Blocks() as demo:
    # Sidebar with title and model selection buttons.
    with gr.Sidebar():
        gr.Markdown("# SigLIP2 224")
        with gr.Row():
            gender_btn = gr.Button("Gender Classification")
            emotion_btn = gr.Button("Emotion Classification")
            dog_breed_btn = gr.Button("Dog Breed Classification")
            deepfake_btn = gr.Button("Deepfake vs Real")
            gym_workout_btn = gr.Button("Gym Workout Classification")
            waste_btn = gr.Button("Waste Classification")

        # State to hold the current model choice.
        selected_model = gr.State("gender")

        # Set model state when buttons are clicked.
        gender_btn.click(fn=select_gender, inputs=[], outputs=selected_model)
        emotion_btn.click(fn=select_emotion, inputs=[], outputs=selected_model)
        dog_breed_btn.click(fn=select_dog_breed, inputs=[], outputs=selected_model)
        deepfake_btn.click(fn=select_deepfake, inputs=[], outputs=selected_model)
        gym_workout_btn.click(fn=select_gym_workout, inputs=[], outputs=selected_model)
        waste_btn.click(fn=select_waste, inputs=[], outputs=selected_model)

        gr.Markdown("### Current Model:")
        model_display = gr.Textbox(value="gender", interactive=False)
        # Update display when state changes.
        selected_model.change(lambda m: m, selected_model, model_display)

    # Main interface: image input, analyze button, and prediction output.
    with gr.Column():
        image_input = gr.Image(type="numpy", label="Upload Image")
        analyze_btn = gr.Button("Analyze")
        output_label = gr.Label(label="Prediction Scores")

        # When the "Analyze" button is clicked, use the selected model to classify the image.
        analyze_btn.click(fn=classify, inputs=[image_input, selected_model], outputs=output_label)

demo.launch()