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
from mnist_digits import classify_digit
from fashion_mnist_cloth import fashion_mnist_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)
    elif model_name == "mnist":
        return classify_digit(image)
    elif model_name == "fashion_mnist":
        return fashion_mnist_classification(image)
    else:
        return {"Error": "No model selected"}

# Function to update the selected model and button styles.
def select_model(model_name):
    model_variants = {
        "gender": "secondary", "emotion": "secondary", "dog breed": "secondary", "deepfake": "secondary",
        "gym workout": "secondary", "waste": "secondary", "age": "secondary", "mnist": "secondary", "fashion_mnist": "secondary"
    }
    model_variants[model_name] = "primary"
    return (model_name, *(gr.update(variant=model_variants[key]) for key in model_variants))

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.Markdown("# SigLIP2 224")
        with gr.Row():
            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")
            mnist_btn = gr.Button("Digit Classify (0-9)", variant="secondary")
            fashion_mnist_btn = gr.Button("Fashion MNIST", variant="secondary")
        
        selected_model = gr.State("age")
        gr.Markdown("### Current Model:")
        model_display = gr.Textbox(value="age", interactive=False)
        selected_model.change(lambda m: m, selected_model, model_display)

        buttons = [gender_btn, emotion_btn, dog_breed_btn, deepfake_btn, gym_workout_btn, waste_btn, age_btn, mnist_btn, fashion_mnist_btn]
        model_names = ["gender", "emotion", "dog breed", "deepfake", "gym workout", "waste", "age", "mnist", "fashion_mnist"]
        
        for btn, name in zip(buttons, model_names):
            btn.click(fn=lambda n=name: select_model(n), inputs=[], outputs=[selected_model] + buttons)
    
    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")
        analyze_btn.click(fn=classify, inputs=[image_input, selected_model], outputs=output_label)

demo.launch()