import cv2, os
import gradio as gr
import numpy as np
from demo.generation import call_generation
from demo.instructions import INSTRUCTIONS_VECTORIZE_SIMPLIFY

VERSION = 'v0.1'

GALLERY_LIST = [os.path.join('demo/gallery',path) for path in os.listdir('demo/gallery')]


def resize_image(image, size):
    # find the minimal size of the image, resize it to size
    # H, W, C = image.shape
    return cv2.resize(image, (size[0], size[1]), interpolation=cv2.INTER_LINEAR)

def HWC3(x):
    assert x.dtype == np.uint8
    if x.ndim == 2:
        x = x[:, :, None]
    assert x.ndim == 3
    H, W, C = x.shape
    assert C == 1 or C == 3 or C == 4
    if C == 3:
        return x
    if C == 1:
        return np.concatenate([x, x, x], axis=2)
    if C == 4:
        color = x[:, :, 0:3].astype(np.float32)
        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
        y = color * alpha + 255.0 * (1.0 - alpha)
        y = y.clip(0, 255).astype(np.uint8)
        return y

    


def process_vector(input_image, upsample_method, svg_simplify, svg_optimize, trace_mode, subsample_ratio, speckle_removal,sorting_method, sorting_order, use_gpu):
    print("Processing vector:",upsample_method, svg_simplify, svg_optimize, trace_mode)
    if input_image is not None:
        ## save input_image to a temp file
        
        ## process the image
        file_list = call_generation(input_image, 
                        preprocess=upsample_method, 
                        simplify=svg_simplify, 
                        optimize=svg_optimize, 
                        mode=trace_mode, 
                        subsample_ratio=subsample_ratio, 
                        speckle_removal=speckle_removal, 
                        sorting_method=sorting_method,
                        sorting_order=sorting_order,
                        use_gpu=use_gpu)
        
        return file_list





block = gr.Blocks(
        title = "VectorizeAnything",
        theme=gr.themes.Soft(
            radius_size=gr.themes.sizes.radius_none,
            text_size=gr.themes.sizes.text_md
        ),
        css="css/style.css",
        ).queue()

with block:
    state = gr.State(value={
        'gallery_selected_img_path': None,  # 当前选中的图片路径
        'gallery_selected_img_path_idx': 0,  # 当前选中的图片路径索引
    })
    with gr.Row():
        gr.HTML(f"""
                    </br>
                    <div>
                        <h1 style="font-size:3rem; "><center>Vectorize Anything: {VERSION} </center></h1>
                    </div>
                    </br>
            """)
    # tab_0 = gr.Tab(label="Gallery (画廊)")
    # with tab_0:
    #     with gr.Row():
    #         gr.Gallery(label='图像生成结果', value=GALLERY_LIST,show_label=False, elem_id="Gallery", columns=5, height=1000)


    tab_3 = gr.Tab(label="IMG to SVG")
    with tab_3:   
        with gr.Accordion('🕹Usage', open=True,):
            with gr.Tabs():
                gr.HTML(INSTRUCTIONS_VECTORIZE_SIMPLIFY)
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="numpy", image_mode="RGBA")
                run_vectorize = gr.Button(value="Vectorize",elem_id="btnVEC")

                with gr.Accordion("Vector options", open=True):
                    upsample_method = gr.Dropdown(choices=["None", "x4", "x2"], type="value", value="None", label="Upsample Method")
                    sorting_method = gr.Dropdown(choices=["brightness","area"], type="value", value="brightness", label="Sorting Method")
                    sorting_order = gr.Dropdown(choices=["ascend","descend"], type="value", value="descend", label="Sorting Order")
                    trace_mode = gr.Radio(choices=["overlap", "cutout"], type="value", value="overlap", label="Trace Mode")
                    use_gpu = gr.Checkbox(label='use GPU', value=False, visible=True)
                    svg_simplify = gr.Checkbox(label='Simplify SVG', value=False, visible=True)
                    svg_optimize = gr.Checkbox(label='Optimize SVG', value=False, visible=True)
                    speckle_removal = gr.Checkbox(label='Remove small speckle', value=False)
                    subsample_ratio = gr.Slider(label="Subsample Ratio", minimum=1, maximum=10000, value=12, step=1, visible=False)


            def exp_gen_click():
                return [gr.Slider(value=512), gr.Slider(value=512)]  # all examples are 512x512, refresh draw_img

            with gr.Column():
                result_vector_gallery = gr.Gallery(label='Output', show_label=False, elem_id="Gallery_vector")
            

            
                with gr.Tab("Image Examples"):
                    exp_gen_en = gr.Examples(
                    [
                        ["test_imgs/demo1.png"],
                        ["test_imgs/demo2.jpg"],
                        ["test_imgs/demo3.png"],
                        ["test_imgs/demo4.png"],
                        ["test_imgs/demo5.png"],
                        ["test_imgs/demo6.png"],
                        ["test_imgs/demo7.png"],
                        ["test_imgs/demo8.png"],
                        ["test_imgs/demo9.png"],
                        ["test_imgs/demo10.png"],
                        ["test_imgs/demo11.png"],
                        ["test_imgs/demo12.png"],
                    ],
                    [input_image],
                    examples_per_page=20,
                    label=''
                    )
                    exp_gen_en.dataset.click(exp_gen_click, None)

    
    vector_ips = [input_image, upsample_method, svg_simplify, svg_optimize, trace_mode, subsample_ratio, speckle_removal,sorting_method, sorting_order, use_gpu]
    run_vectorize.click(fn=process_vector, inputs=vector_ips, outputs=result_vector_gallery)

block.launch(server_name='0.0.0.0', share=False,debug=True, root_path=f"/{os.getenv('GRADIO_PROXY_PATH')}" if os.getenv('GRADIO_PROXY_PATH') else "")