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 workout_classification
from augmented_waste_classifier import waste_classification
from age_classification import age_classification

# 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 workout_classification(image)
    elif model_name == "waste":
        return waste_classification(image)
    elif model_name == "age":
        return age_classification(image)
    else:
        return {"Error": "No model selected"}

# Function to update the selected model and button styles.
def select_model(model_name):
    # Set each button's variant to "primary" if selected, otherwise "secondary"
    gender_variant = "primary" if model_name == "gender" else "secondary"
    emotion_variant = "primary" if model_name == "emotion" else "secondary"
    dog_breed_variant = "primary" if model_name == "dog breed" else "secondary"
    deepfake_variant = "primary" if model_name == "deepfake" else "secondary"
    gym_workout_variant = "primary" if model_name == "gym workout" else "secondary"
    waste_variant = "primary" if model_name == "waste" else "secondary"
    age_variant = "primary" if model_name == "age" else "secondary"
    # Return new state and update objects for each button in the specified order.
    return (
       model_name,
       gr.update(variant=gender_variant),
       gr.update(variant=emotion_variant),
       gr.update(variant=dog_breed_variant),
       gr.update(variant=deepfake_variant),
       gr.update(variant=gym_workout_variant),
       gr.update(variant=waste_variant),
       gr.update(variant=age_variant)
    )

with gr.Blocks() as demo:
    # Sidebar with title and model selection buttons.
    with gr.Sidebar():
        gr.Markdown("# SigLIP2 224")
        with gr.Row():
            # Initialize buttons with variants. Default is "age" set to primary.
            age_btn = gr.Button("Age Classification", variant="primary")
            gender_btn = gr.Button("Gender Classification", variant="secondary")
            emotion_btn = gr.Button("Emotion Classification", variant="secondary")
            dog_breed_btn = gr.Button("Dog Breed Classification", variant="secondary")
            deepfake_btn = gr.Button("Deepfake vs Real", variant="secondary")
            gym_workout_btn = gr.Button("Gym Workout Classification", variant="secondary")
            waste_btn = gr.Button("Waste Classification", variant="secondary")

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

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

        # Set up click events for each button, updating state and button variants.
        gender_btn.click(fn=lambda: select_model("gender"),
                         inputs=[],
                         outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
        emotion_btn.click(fn=lambda: select_model("emotion"),
                          inputs=[],
                          outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
        dog_breed_btn.click(fn=lambda: select_model("dog breed"),
                            inputs=[],
                            outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
        deepfake_btn.click(fn=lambda: select_model("deepfake"),
                           inputs=[],
                           outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
        gym_workout_btn.click(fn=lambda: select_model("gym workout"),
                              inputs=[],
                              outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
        waste_btn.click(fn=lambda: select_model("waste"),
                        inputs=[],
                        outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])
        age_btn.click(fn=lambda: select_model("age"),
                      inputs=[],
                      outputs=[selected_model, gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn])

    # 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("Classify / Predict")
        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()