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
from indian_western_food_classify import food_classification
from bird_species import bird_classification
from alphabet_sign_language_detection import sign_language_classification
from rice_leaf_disease import classify_leaf_disease
from traffic_density import traffic_density_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)
    elif model_name == "food":
        return food_classification(image)
    elif model_name == "bird":
        return bird_classification(image)
    elif model_name == "leaf disease":
        return classify_leaf_disease(image)
    elif model_name == "sign language":
        return sign_language_classification(image)
    elif model_name == "traffic density":
        return traffic_density_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", "food": "secondary", "bird": "secondary", "leaf disease": "secondary",
        "sign language": "secondary", "traffic density": "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("# Choose Domain")
        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 Classification", variant="secondary")
            food_btn = gr.Button("Indian/Western Food", variant="secondary")
            bird_btn = gr.Button("Bird Species Classification", variant="secondary")
            leaf_disease_btn = gr.Button("Rice Leaf Disease", variant="secondary")
            sign_language_btn = gr.Button("Alphabet Sign Language", variant="secondary")
            traffic_density_btn = gr.Button("Traffic Density", 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, food_btn, bird_btn, leaf_disease_btn, sign_language_btn, traffic_density_btn]
        model_names = ["gender", "emotion", "dog breed", "deepfake", "gym workout", "waste", "age", "mnist", "fashion_mnist", "food", "bird", "leaf disease", "sign language", "traffic density"]
        
        for btn, name in zip(buttons, model_names):
            btn.click(fn=lambda n=name: select_model(n), inputs=[], outputs=[selected_model] + buttons)
    
    with gr.Row():
        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")  # Output on the right

    analyze_btn.click(fn=classify, inputs=[image_input], outputs=output_label)

demo.launch()