import gradio as gr
from datasets import load_dataset

dataset = iter(load_dataset('philipp-zettl/mtg_cards-2025-04-04', split='train'))


def get_next_elem():
    try:
        elem = next(dataset)
    except StopIteration:
        return

    return elem


def greet(img1, img2, label):
    return str(label)


with gr.Blocks() as demo:
    gr.Markdown('''# MTG Similarity Labeler

    ''')
    elem = get_next_elem()
    img1_src = elem.get('image_1')
    img2_src = elem.get('image_2')
    print(elem)
    with gr.Row():
        with gr.Column():
            img1 = gr.Image(img1_src, interactive=False)
        with gr.Column():
            img2 = gr.Image(img2_src, interactive=False)
    with gr.Row():
        with gr.Column():
            btn = gr.Button('Similar')
        with gr.Column():
            btn2 = gr.Button('Not Similar')

    uuid = gr.State("")
    btn.click(lambda x, y: greet(x, y, 1), inputs=[img1, img2, uuid], outputs=[img1, img2, uuid])
    btn2.click(lambda x, y: greet(x, y, 0), inputs=[img1, img2, uuid], outputs=[img1, img2, uuid])

demo.launch()