import gradio as gr
from utils import load_specific_model, inference
from datetime import datetime
import pytz

# current_model = None  # Initialize the current model as None
MODEL_NAMES = ["EfficientNet-B3", "EfficientNet-B4", "vgg19", "resnet50", "dinov2_vits14"]

# Define a set of example images
example_images = [
    ("Beispielbild Glas", "src/examples/Glas.jpg"),
    ("Beispielbild Organic", "src/examples/Organic.jpg"),
    ("Beispielbild Papier", "src/examples/Papier.jpg"),
    ("Beispielbild Restmüll", "src/examples/Restmuell.jpg"),
    ("Beispielbild Wertstoff", "src/examples/Wertstoff.jpg")
]

def predict(inp_image, inp_dropdown):
    if inp_dropdown is None:
        raise gr.Error("No model selected!")
    if inp_image is None:
        raise gr.Error("No image uploaded!")
    if inp_dropdown not in MODEL_NAMES:
        raise gr.Error("Invalid model selected!")
    
    # Get the current time in UTC
    utc_now = datetime.now(pytz.utc)

    # Convert UTC time to German time zone
    german_timezone = pytz.timezone('Europe/Berlin')
    german_time = utc_now.astimezone(german_timezone)

    # Format and print the German time and date
    formatted_time = german_time.strftime('%Y-%m-%d %H:%M:%S %Z')
    print('Current time and date:', formatted_time)
    
    print(f"\nInput: {inp_dropdown}\n")
    current_model = load_specific_model(inp_dropdown)
    confidences = inference(current_model, inp_image)
    print(f"\nConfidences: {confidences}\n")
    return confidences

with gr.Blocks() as demo:
    with open('src/app_template.md', 'r') as f:
        markdown_string = f.read()
    header = gr.Markdown(markdown_string)
    
    with gr.Row(variant="panel", equal_height=True):

        user_image = gr.Image(
            type="pil",
            label="Upload Your Own Image",
            info="You can also upload your own image for prediction.",
            scale=2,
            height=350,
        )
        
        with gr.Column():
            output = gr.Label(
                num_top_classes=3,
                label="Output",
                info="Top three predicted classes and their confidences.",
                scale=2,
            )

            model_dropdown = gr.Dropdown(
                MODEL_NAMES,
                value="EfficientNet-B3",
                label="Model",
                info="Select a model to use.",
                scale=1,
            )
            
            predict_button = gr.Button(value="Predict", label="Predict", info="Click to make a prediction.", scale=1)
            predict_button.click(fn=predict, inputs=[user_image, model_dropdown], outputs=output, queue=True)

    gr.Markdown("## Example Images")
    gr.Markdown("You can just drag and drop these images into the image uploader above!")
    
    with gr.Row():
        for name, image_path in example_images:
            example_image = gr.Image(
                value=image_path,
                label=name,
                type="pil",
                height=220,
                interactive=False,
            )

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