import os

import gradio as gr

from ccip import _VALID_MODEL_NAMES, _DEFAULT_MODEL_NAMES, ccip_difference, ccip_default_threshold


def _compare(imagex, imagey, model_name):
    threshold = ccip_default_threshold(model_name)
    diff = ccip_difference(imagex, imagey)

    return diff, 'Same' if diff <= threshold else 'Not Same'


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        gr_input_x = gr.Image(type='pil', label='Image X')
                    with gr.Column():
                        gr_input_y = gr.Image(type='pil', label='Image Y')
                with gr.Row():
                    gr_model_name = gr.Dropdown(_VALID_MODEL_NAMES, value=_DEFAULT_MODEL_NAMES, label='Model')

                gr_button = gr.Button(value='Compare', variant='primary')

            with gr.Column():
                with gr.Row():
                    gr_diff = gr.Number(value=0.0, label='Difference')
                with gr.Row():
                    gr_prediction = gr.Text(value='', label='Prediction')

            gr_button.click(
                _compare,
                inputs=[gr_input_x, gr_input_y, gr_model_name],
                outputs=[gr_diff, gr_prediction],
            )

    demo.queue(os.cpu_count()).launch()