import gradio as gr
import random
import csv
import datetime

class_names = ['cat', 'dog']

def update_dropdown(className):
    class_names.append(className)
    updated_choices = gr.Dropdown(choices=class_names)
    return updated_choices, updated_choices

def show_picked_class(className):
    return className

def image_classifier(inp):
    if inp is None:
        return {'cat': 0.3, 'dog': 0.7}

    num_class = len(class_names)

    # Generate random percentages between 0 and 1
    percentages = [random.random() for _ in range(num_class)]
    total = sum(percentages)

    # Normalize the percentages to ensure they sum up to 1
    normalized_percentages = [p / total for p in percentages]

    labeled_result = {name:score for name, score in zip(class_names, normalized_percentages)}

    return labeled_result

demo = gr.Blocks()

with demo as app:
    gr.Markdown("# Single Image")

    with gr.Row():
        with gr.Column():
            inp_img = gr.Image()
            with gr.Row():
                clear_btn = gr.Button(value="Clear")
                process_btn = gr.Button(value="Process", variant="primary")
        with gr.Column():
            out_txt = gr.Label(label="Probabilities", num_top_classes=3)

    text_input = gr.Textbox(label="Input the new class here")
    b1 = gr.Button("Add new class")

    text_options = gr.Dropdown(class_names, label="Class Label", multiselect=False)
    b2 = gr.Button("Show me the picked class")
    picked_class = gr.Textbox()

    b2.click(show_picked_class, inputs=text_options, outputs=picked_class)

    process_btn.click(image_classifier, inputs=inp_img, outputs=out_txt)
    clear_btn.click(lambda:(
        gr.update(value=None),
        gr.update(value=None)
        ),
        inputs=None,
        outputs=[inp_img, out_txt])

    gr.Markdown("# Multiple Images")

    def show_to_gallery(images):
        file_paths = [[file.name, class_names[0]] for file in images]
        # print(file_paths)
        return file_paths, file_paths

    def get_select_index(evt: gr.SelectData):
        # print("data",evt._data)
        # print("value",evt.value)
        return evt.index

    with gr.Column():
        imgs = gr.State()

        multiple_inputs = gr.UploadButton(label="Upload multiple images file here.", file_count="multiple", file_types=["image"])
        gallery = gr.Gallery()
        selected = gr.Textbox(label="Image Gallery Index")

        images_label = gr.Dropdown(class_names, label="Class Label", multiselect=False)
        b3 = gr.Button("Save and change the label using dropdown")

        b1.click(update_dropdown, inputs=text_input, outputs=[text_options, images_label])

        multiple_inputs.upload(show_to_gallery, inputs=multiple_inputs, outputs=[gallery, imgs])

        gallery.select(get_select_index, None, selected)

        def change_labels(imgs, index, images_label):
            index = int(index)
            label_idx = class_names.index(images_label)
            imgs[index][1] = class_names[label_idx]
            return imgs, imgs

        b3.click(change_labels, [imgs, selected, images_label], [imgs, gallery])

        gr.Markdown('### Save Metadata Into .csv')
        b4 = gr.Button("Upload to metadata")

        def upload_metadata(imgs):
            time_uploaded = datetime.datetime.now()
            time_str = time_uploaded.strftime("%m-%d-%Y_%H-%M-%S")

            with open(f'{time_str}.csv', mode='w', newline='') as csv_file:
                # Create a CSV writer
                csv_writer = csv.writer(csv_file)

                # Write the header row
                csv_writer.writerow(['image_path', 'ground_truth', 'time_uploaded', 'prediction_label', 'prediction_conf'])

                for image in imgs:
                    image.append(time_str)
                    model_output = image_classifier(image)
                    # Sort the label and confidence output in descending order
                    sorted_output = dict(sorted(model_output.items(), key=lambda item: item[1], reverse=True))

                    # Extract the label with the highest value
                    label_prediction = next(iter(sorted_output))
                    image.append(label_prediction)

                    label_confidence = model_output[label_prediction]
                    image.append(label_confidence)

                # Write the data rows
                csv_writer.writerows(imgs)

            print(f"Metadata CSV file has been created.")

        b4.click(upload_metadata, inputs=imgs)


demo.launch(debug=True)