from PIL import Image
import gradio as gr

def preview_images(files):
    return gr.Gallery(files)

def merge_images(image_paths, direction, b_resize):
    images = [Image.open(p) for p in image_paths]
    if b_resize:
        w, h = images[0].size
        images = [i.resize((w, h)) for i in images]

    widths, heights = zip(*(i.size for i in images))
    if direction == 'horizontal':
        total_width = sum(widths)
        max_height = max(heights)
        new_im = Image.new('RGB', (total_width, max_height))
        x_offset = 0
        for im in images:
            new_im.paste(im, (x_offset, 0))
            x_offset += im.width
    else:  # vertical
        total_height = sum(heights)
        max_width = max(widths)
        new_im = Image.new('RGB', (max_width, total_height))
        y_offset = 0
        for im in images:
            new_im.paste(im, (0, y_offset))
            y_offset += im.height

    return [new_im]

def split_image(image, parts, direction='horizontal'):
    width, height = image.size
    if direction == 'horizontal':
        part_width = width // parts
        return [image.crop((i * part_width, 0, (i + 1) * part_width, height)) for i in range(parts)]
    else:  # vertical
        part_height = height // parts
        return [image.crop((0, i * part_height, width, (i + 1) * part_height)) for i in range(parts)]


css = "#file {height: 10vh;}"
result_image = gr.Gallery(height="45vh")
with gr.Blocks(css=css) as demo: 
    with gr.Tabs():
        with gr.Tab(label="Split") as tab_split:
            with gr.Row():
                split_input = gr.Image(type="pil", height="40vh")
                with gr.Column():
                    split_direction = gr.Radio(label="Direction", choices=["vertical", "horizontal"], value="vertical", interactive=True) 
                    split_num = gr.Number(label="Number of parts", value=2, interactive=True)
                    split_btn = gr.Button(scale=1)
            split_btn.click(fn=split_image, inputs=[split_input, split_num, split_direction], outputs=result_image)
        with gr.Tab(label="Merge") as tab_split:
            with gr.Row():
                with gr.Column():
                    merge_files = gr.File(file_count="multiple", elem_id="file")
                    merge_preview = gr.Gallery(label="prevew", height="35vh")
                with gr.Column():
                    merge_direction = gr.Radio(label="Direction", choices=["vertical", "horizontal"], value="vertical", interactive=True) 
                    merge_b_resize = gr.Checkbox(label="Resize to first image", value=True, interactive=True)
                    merge_btn = gr.Button(scale=1)
            merge_files.change(fn=preview_images, inputs=merge_files, outputs=merge_preview)
            merge_btn.click(fn=merge_images, inputs=[merge_files, merge_direction, merge_b_resize], outputs=result_image)
    result_image.render()
    
demo.launch()