import os
import numpy as np
import random
import time
import gradio as gr
from runner import Runner
import matplotlib.pyplot as plt

def show_mask(mask, ax, color='blue'):
    if color == 'blue':
        # reference, blue
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    else:
        # target, green
        color = np.array([78 / 255, 238 / 255, 148 / 255, 0.6])

    # if random_color:
    #     color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    # else:
    #     color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)

def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))


def show_img_point_box_mask(img, input_point=None, input_label=None, box=None, masks=None, save_path=None, mode='mask', color='blue'):

    if mode == 'point':
        # point
        plt.figure(figsize=(10, 10))
        plt.imshow(img)
        show_points(input_point, input_label, plt.gca())
        plt.axis('on')
        plt.savefig(save_path, bbox_inches='tight')
    elif mode == 'box':
        # box
        plt.figure(figsize=(10, 10))
        plt.imshow(img)
        show_box(box, plt.gca())
        plt.axis('on')
        plt.savefig(save_path, bbox_inches='tight')
    else:
        # mask
        plt.figure(figsize=(10, 10))
        plt.imshow(img)
        show_mask(masks, plt.gca(), color=color)
        plt.axis('off')
        plt.savefig(save_path, bbox_inches='tight')
    plt.close()


def create_oss_demo(
        runner: Runner,
        pipe: None = None
) -> gr.Blocks:

    examples = [
        ['./gradio_demo/images/horse1.png', './gradio_demo/images/horse2.png', './gradio_demo/images/horse3.png'],
        ['./gradio_demo/images/hmbb1.png', './gradio_demo/images/hmbb2.png', './gradio_demo/images/hmbb3.png'],
        ['./gradio_demo/images/earth1.png', './gradio_demo/images/earth2.png', './gradio_demo/images/earth3.png'],
        ['./gradio_demo/images/elephant1.png', './gradio_demo/images/elephant2.png', './gradio_demo/images/elephant3.png'],
        ['./gradio_demo/images/dinosaur1.png', './gradio_demo/images/dinosaur2.png', './gradio_demo/images/dinosaur3.png'],
    ]

    with gr.Blocks() as oss_demo:
        with gr.Column():

            # inputs
            with gr.Row():
                img_input_prompt = gr.ImageMask(label='Prompt (提示图)')
                img_input_target1 = gr.Image(label='Target 1 (测试图1)')
                img_input_target2 = gr.Image(label='Target 2 (测试图2)')

            version = gr.inputs.Radio(['version 1 (🔺 multiple instances  🔻 whole, 🔻 part)',
                                       'version 2 (🔻 multiple instances  🔺 whole, 🔻 part)',
                                       'version 3 (🔻 multiple instances  🔻 whole, 🔺 part)'],
                                      type="value", default='version 1 (🔺 whole, 🔻 part)',
                                      label='Multiple Instances (version 1), Single Instance (version 2), Part of a object (version 3)')

            with gr.Row():
                submit1 = gr.Button("提交 (Submit)")
                clear = gr.Button("清除 (Clear)")
            info = gr.Text(label="Processing result: ", interactive=False)

            # decision
            K = gr.Slider(0, 10, 10, step=1, label="Controllable mask output", interactive=True)
            submit2 = gr.Button("提交 (Submit)")

            # outputs
            with gr.Row():
                img_output_pmt = gr.Image(label='Prompt (提示图)')
                img_output_tar1 = gr.Image(label='Output 1 (输出图1)')
                img_output_tar2 = gr.Image(label='Output 2 (输出图2)')

        # images
        gr.Examples(
            examples=examples,
            fn=runner.inference_oss_ops,
            inputs=[img_input_prompt, img_input_target1, img_input_target2],
            outputs=info
        )

        submit1.click(
            fn=runner.inference_oss_ops,
            inputs=[img_input_prompt, img_input_target1, img_input_target2, version],
            outputs=info
        )
        submit2.click(
            fn=runner.controllable_mask_output,
            inputs=K,
            outputs=[img_output_pmt, img_output_tar1, img_output_tar2]
        )

        clear.click(
            fn=runner.clear_fn,
            inputs=None,
            outputs=[img_input_prompt, img_input_target1, img_input_target2, info, img_output_pmt, img_output_tar1, img_output_tar2],
            queue=False
        )

    return oss_demo


def create_vos_demo(
        runner: Runner,
        pipe: None = None
) -> gr.Interface:

    raise NotImplementedError

def create_demo(
        runner: Runner,
        pipe: None = None
) -> gr.TabbedInterface:

    title = "Matcher🎯: Segment Anything with One Shot Using All-Purpose Feature Matching<br> \
    <div align='center'> \
    <h2><a href='https://arxiv.org/abs/2305.13310' target='_blank' rel='noopener'>[paper]</a> \
    <a href='https://github.com/aim-uofa/Matcher' target='_blank' rel='noopener'>[code]</a></h2> \
    <h2>Matcher can segment anything with one shot by integrating an all-purpose feature extraction model and a class-agnostic segmentation model.</h2> \
    <br> \
    </div> \
    "

    oss_demo = create_oss_demo(runner=runner, pipe=pipe)
    # vos_demo = create_vos_demo(runner=runner, pipe=pipe)
    demo = gr.TabbedInterface(
        [oss_demo,],
        ['OSS+OPS',], title=title)
    return demo


if __name__ == '__main__':
    pipe = None
    HF_TOKEN = os.getenv('HF_TOKEN')
    runner = Runner(HF_TOKEN)
    # runner = None
    demo = create_demo(runner, pipe)
    demo.launch(enable_queue=False)