import gradio as gr

import os
import json
from filelock import Timeout, FileLock
from random import sample, choice

def random_images(gen_images):
    return choice(gen_images)

def vote(method):
    lock_path = "votes.lock"
    file_path = "votes.json"
    lock = FileLock(lock_path, timeout=10)
    
    with lock:
        # Check if the votes file exists, create if not
        if not os.path.exists(file_path):
            with open(file_path, 'w') as file:
                json.dump({}, file)
        
        # Read the current votes
        with open(file_path, 'r') as file:
            votes = json.load(file)
        
        # Update the vote count
        try:
            option = (method)
            if option in votes:
                votes[option] += 1
            else:
                votes[option] = 1
        except:
            return "Error"
        # Write the updated votes back to the file
        with open(file_path, 'w') as file:
            json.dump(votes, file)

def vote_fn(method):
    vote(method)
    image_1, image_2, first_method, second_method = refresh()
    return image_1, image_2, first_method, second_method

css = """
.gradio-container {
  width: 50%;        /* Set the width of the container to 50% of the parent element */
  margin-left: auto; /* Set the left margin to auto to enable horizontal centering */
  margin-right: auto;/* Set the right margin to auto to enable horizontal centering */
}
.gen_images {
    padding: 30px;
  }
"""

METHODS = ["baseline", "lora"]


baseline_imgs = [f"gen_images/baseline/{path}" for path in os.listdir("gen_images/baseline")]
lora_imgs = [f"gen_images/lora/{path}" for path in os.listdir("gen_images/lora")]

og_images = [os.path.join("og_images", path) for path in os.listdir("og_images")]

print(baseline_imgs)

def refresh():
    random_baseline_img = random_images(baseline_imgs)
    random_lora_img = random_images(lora_imgs)

    first_method = choice(METHODS)

    if first_method == "baseline":
        image1 = random_baseline_img
        image2 = random_lora_img
        second_method = "lora"
    else:
        image1 = random_lora_img
        image2 = random_baseline_img
        second_method = "baseline"

    return image1, image2, first_method, second_method

with gr.Blocks(css=css) as demo: 
    first_method = gr.State()
    first_method.value = "baseline"
    second_method = gr.State()
    second_method.value = "lora"

    # show the title of the demo
    gr.components.HTML("<h1>Image Scoring</h1>")
    
    # show the description of the demo
    gr.components.HTML("<p>Select the best image among the proposed.</p>")
    
    # show the images to be scored on the same row
    with gr.Blocks() as selection:
            with gr.Row(scale=3):
                
                random_baseline_img = random_images(baseline_imgs)
                random_lora_img = random_images(lora_imgs)

                with gr.Row(scale=6, elem_id="gen_images"):
                    if first_method.value == "baseline":
                        
                        with gr.Column(scale=6):
                            image_1 = gr.Image(random_baseline_img, container=False, show_label=False)
                            button1 = gr.Button(value="First Image 👍")
                        with gr.Column(scale=6):
                            image_2 = gr.Image(random_lora_img, container=False, show_label=False)
                            button2 = gr.Button(value="Second Image 👍") 
                    
                    else:
                        with gr.Column(scale=6):
                            image_1 = gr.Image(random_baseline_img, container=False, show_label=False)
                            button1 = gr.Button(value="First Image 👍")
                        with gr.Column(scale=6):
                            image_2 = gr.Image(random_lora_img, container=False, show_label=False)
                            button2 = gr.Button(value="Second Image 👍") 


            button1.click(vote_fn, inputs=[first_method], outputs=[image_1, image_2, first_method, second_method])
            button2.click(vote_fn, inputs=[second_method], outputs=[image_1, image_2, first_method, second_method])


demo.launch()