import os, sys, time, random, shutil, json, re
import gradio as gr
import cv2
import numpy as np

from annotator.util import resize_image, HWC3

RES = os.path.join(os.path.dirname(__file__), "_res")

gr.set_static_paths(paths=["_res", "_res/assets/"])

custom_css = RES + "/_custom.css"
custom_js = RES + "/_custom.js"

custom_head = f"""
        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.9.0/css/all.min.css"/>
        <!--script src="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.9.0/js/all.min.js"></script-->
"""

# MARK: Gradio Theme
theme = gr.themes.Soft(
    primary_hue="emerald",
    radius_size="sm",
    neutral_hue=gr.themes.Color(c100="#a6adc8", c200="#9399b2", c300="#7f849c", c400="#6c7086", c50="#cdd6f4", c500="#585b70", c600="#45475a", c700="#313244", c800="#1e1e2e", c900="#181825", c950="#11111b"),
)

title = "ControlNet v1.1 Preprocessors Standalone"


# MARK: Funktionen

# Canny
model_canny = None


def canny(img, res, l, h, old_images=None):
    print("Old Images: ", old_images)
    result_images = []
    result = None
    img = resize_image(HWC3(img), res)
    global model_canny
    if model_canny is None:
        from annotator.canny import CannyDetector

        model_canny = CannyDetector()
    result = model_canny(img, l, h)
    result_images.append(result)

    if old_images is not None:
        result_images.extend(old_images)
    return result_images


# Hed
model_hed = None


def hed(img, res):
    img = resize_image(HWC3(img), res)
    global model_hed
    if model_hed is None:
        from annotator.hed import HEDdetector

        model_hed = HEDdetector()
    result = model_hed(img)
    return [result]


# Pidi
model_pidi = None


def pidi(img, res):
    img = resize_image(HWC3(img), res)
    global model_pidi
    if model_pidi is None:
        from annotator.pidinet import PidiNetDetector

        model_pidi = PidiNetDetector()
    result = model_pidi(img)
    return [result]


# MLSD
model_mlsd = None


def mlsd(img, res, thr_v, thr_d):
    img = resize_image(HWC3(img), res)
    global model_mlsd
    if model_mlsd is None:
        from annotator.mlsd import MLSDdetector

        model_mlsd = MLSDdetector()
    result = model_mlsd(img, thr_v, thr_d)
    return [result]


# Midas
model_midas = None


def midas(img, res):
    img = resize_image(HWC3(img), res)
    global model_midas
    if model_midas is None:
        from annotator.midas import MidasDetector

        model_midas = MidasDetector()
    result = model_midas(img)
    return [result]


# Zoe
model_zoe = None


def zoe(img, res):
    img = resize_image(HWC3(img), res)
    global model_zoe
    if model_zoe is None:
        from annotator.zoe import ZoeDetector

        model_zoe = ZoeDetector()
    result = model_zoe(img)
    return [result]


# Normal Bae
model_normalbae = None


def normalbae(img, res):
    img = resize_image(HWC3(img), res)
    global model_normalbae
    if model_normalbae is None:
        from annotator.normalbae import NormalBaeDetector

        model_normalbae = NormalBaeDetector()
    result = model_normalbae(img)
    return [result]


# DWPose
model_dwpose = None


def dwpose(img, res):
    img = resize_image(HWC3(img), res)
    global model_dwpose
    if model_dwpose is None:
        from annotator.dwpose import DWposeDetector

        model_dwpose = DWposeDetector()
    result = model_dwpose(img)
    return [result]


# OpenPose
model_openpose = None


def openpose(img, res, hand_and_face):
    img = resize_image(HWC3(img), res)
    global model_openpose
    if model_openpose is None:
        from annotator.openpose import OpenposeDetector

        model_openpose = OpenposeDetector()
    result = model_openpose(img, hand_and_face)
    return [result]


# model_uniformer = None
# def uniformer(img, res):
#    img = resize_image(HWC3(img), res)
#    global model_uniformer
#    if model_uniformer is None:
#        from annotator.uniformer import UniformerDetector
#        model_uniformer = UniformerDetector()
#    result = model_uniformer(img)
#    return [result]


# Lineart
model_lineart_anime = None
model_lineart = None


def lineart(img, res, preprocessor_name="Lineart", invert=True, old_images=None):
    print("Old Images: ", old_images)
    result_images = []
    result = None
    img = resize_image(HWC3(img), res)
    ["Lineart", "Lineart Coarse", "Lineart Anime"]
    if preprocessor_name in ["Lineart", "Lineart Coarse"]:
        coarse = "Coarse" in preprocessor_name
        global model_lineart
        if model_lineart is None:
            from annotator.lineart import LineartDetector

            model_lineart = LineartDetector()
        if invert:
            result = (cv2.bitwise_not(model_lineart(img, coarse)), preprocessor_name)
        else:
            result = (model_lineart(img, coarse), preprocessor_name)
        # return [result]
    elif preprocessor_name == "Lineart Anime":
        global model_lineart_anime
        if model_lineart_anime is None:
            from annotator.lineart_anime import LineartAnimeDetector

            model_lineart_anime = LineartAnimeDetector()
        if invert:
            result = (cv2.bitwise_not(model_lineart_anime(img)), preprocessor_name)
        else:
            result = (model_lineart_anime(img), preprocessor_name)

    result_images.append(result)
    if old_images is not None:
        result_images.extend(old_images)
    return result_images


# OneFormer
model_oneformer_coco = None


def oneformer_coco(img, res):
    img = resize_image(HWC3(img), res)
    global model_oneformer_coco
    if model_oneformer_coco is None:
        from annotator.oneformer import OneformerCOCODetector

        model_oneformer_coco = OneformerCOCODetector()
    result = model_oneformer_coco(img)
    return [result]


# OneFormer ADE20k
model_oneformer_ade20k = None


def oneformer_ade20k(img, res):
    img = resize_image(HWC3(img), res)
    global model_oneformer_ade20k
    if model_oneformer_ade20k is None:
        from annotator.oneformer import OneformerADE20kDetector

        model_oneformer_ade20k = OneformerADE20kDetector()
    result = model_oneformer_ade20k(img)
    return [result]


# Content Shuffler
model_content_shuffler = None


def content_shuffler(img, res):
    img = resize_image(HWC3(img), res)
    global model_content_shuffler
    if model_content_shuffler is None:
        from annotator.shuffle import ContentShuffleDetector

        model_content_shuffler = ContentShuffleDetector()
    result = model_content_shuffler(img)
    return [result]


# Color Shuffler
model_color_shuffler = None


def color_shuffler(img, res):
    img = resize_image(HWC3(img), res)
    global model_color_shuffler
    if model_color_shuffler is None:
        from annotator.shuffle import ColorShuffleDetector

        model_color_shuffler = ColorShuffleDetector()
    result = model_color_shuffler(img)
    return [result]


# Inpaint
model_inpaint = None


def inpaint(image, invert):
    #    color = HWC3(image["image"])
    color = HWC3(image["background"])
    if invert:
        #        alpha = image["mask"][:, :, 0:1]
        alpha = image["layers"][0][:, :, 3:]
    else:
        #        alpha = 255 - image["mask"][:, :, 0:1]
        alpha = 255 - image["layers"][0][:, :, 3:]
    result = np.concatenate([color, alpha], axis=2)
    return [result]


# MARK: GRADIO UI

input_options = []

with gr.Blocks(theme=theme, css=custom_css, js=custom_js, head=custom_head, title=title) as demo:
    with gr.Row(elem_classes="row-header"):
        gr.Markdown(
            f"""
                    <h1>{title}</h1>
                    <p>Erstelle "Control Images" für Stable Diffusion und andere Tools die ControlNet verwenden.<br/>
                    Diese Demo läuft nur auf CPU, eine Inference wird daher sehr lange dauern. Du kannst diesen Space clonen und lokal auf deinem Computer verwenden.</p>
                    <p><i class="winking-hand-emoji"></i> Sebastian, gib dem Space gerne ein <i class="heart-beat-emoji"></i></p>
                    
                    """,
            elem_classes="md-header",
        )
    with gr.Tab("Preprocessors"):
        with gr.Row(elem_classes="row-main"):
            with gr.Column(scale=1, elem_id="input_column", elem_classes="input-column"):
                input_image = gr.Image(label="Dein Bild", type="numpy")

                ## TAB LINEART
                with gr.Tab("Lineart") as tab_lineart:
                    with gr.Row():
                        with gr.Column():
                            invert_toggle_info = ["Schwarzer Hintergrund, Weiße Linien", "Weißer Hintergrund, Schwarze Linien"]
                            preprocessor_name = gr.Radio(label="Preprocessor", show_label=False, choices=["Lineart", "Lineart Coarse", "Lineart Anime"], type="value", value="Lineart", elem_classes="radio-btn-group")

                            invert = gr.Checkbox(label="Farbe invertieren?", info=invert_toggle_info[0], value=True, elem_classes="toggle-btn")
                            # resolution = gr.Slider(label="Auflösung", minimum=256, maximum=1024, value=512, step=64)
                            invert.change(lambda x: {"info": invert_toggle_info[0] if x else invert_toggle_info[1], "__type__": "update"}, inputs=invert, outputs=invert)

                ## TAB Canny
                with gr.Tab("Canny Edge") as tab_canny:
                    # with gr.Row():
                    #     gr.Markdown("## Canny Edge")
                    with gr.Row():
                        with gr.Column():

                            low_threshold = gr.Slider(label="niedriger Schwellenwert", minimum=1, maximum=255, value=100, step=1)
                            high_threshold = gr.Slider(label="hoher Schwellenwert", minimum=1, maximum=255, value=200, step=1)

                with gr.Row():
                    resolution = gr.Slider(label="Auflösung (Pixel Breiet)", minimum=256, maximum=1024, value=512, step=64)
                with gr.Row():
                    run_btn_lineart = gr.Button("Los", variant="primary", visible=True)
                    run_btn_canny = gr.Button("Los", variant="primary", visible=False)

                all_run_btns = [run_btn_lineart, run_btn_canny]

                def set_inputs(tab):
                    tab = "tab_lineart" if tab is None else tab
                    # global input_options
                    lineart_btn_visible = True if tab == "tab_lineart" else False
                    canny_btn_visible = True if tab == "tab_canny" else False

                    return {"visible": lineart_btn_visible, "__type__": "update"}, {"visible": canny_btn_visible, "__type__": "update"}

                tab_lineart.select(fn=lambda: set_inputs("tab_lineart"), inputs=None, outputs=[run_btn_lineart, run_btn_canny])
                tab_canny.select(fn=lambda: set_inputs("tab_canny"), inputs=None, outputs=[run_btn_lineart, run_btn_canny])

            with gr.Column(scale=2):
                gallery = gr.Gallery(label="Generated images", show_label=False, interactive=False, format="png", elem_id="output_gallery", elem_classes="output-gallery", columns=[3], rows=[2], object_fit="contain", height="auto", type="filepath")

    with gr.Tab("Tutorial (demnächst!)"):
        with gr.Row(elem_classes="row-main"):
            with gr.Column():
                gr.Markdown(
                    f"""
                    # Das Tutorial kommt bald.
                    """
                )

    # MARK: Button Runs
    run_btn_lineart.click(fn=lineart, inputs=[input_image, resolution, preprocessor_name, invert, gallery], outputs=[gallery])
    run_btn_canny.click(fn=canny, inputs=[input_image, resolution, low_threshold, high_threshold, gallery], outputs=[gallery])


"""
with gr.Blocks(theme=theme, css="custom.css", js="javascript.js") as demo:
    gr.Markdown(DESCRIPTION, elem_classes="top-description")
    with gr.Tab("Canny Edge", elem_id="tab_wrapper", elem_classes="tab_wrapper"):
        with gr.Row():
            gr.Markdown("## Canny Edge")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                low_threshold = gr.Slider(label="low_threshold", minimum=1, maximum=255, value=100, step=1)
                high_threshold = gr.Slider(label="high_threshold", minimum=1, maximum=255, value=200, step=1)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_btn_canny = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=canny, inputs=[input_image, resolution, low_threshold, high_threshold], outputs=[gallery])

    with gr.Tab("HED Edge"):
        with gr.Row():
            gr.Markdown("## HED Edge&nbsp;&quot;SoftEdge&quot;")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=hed, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("Pidi Edge"):
        with gr.Row():
            gr.Markdown("## Pidi Edge&nbsp;&quot;SoftEdge&quot;")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=pidi, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("MLSD Edge"):
        with gr.Row():
            gr.Markdown("## MLSD Edge")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                value_threshold = gr.Slider(label="value_threshold", minimum=0.01, maximum=2.0, value=0.1, step=0.01)
                distance_threshold = gr.Slider(label="distance_threshold", minimum=0.01, maximum=20.0, value=0.1, step=0.01)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=384, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=mlsd, inputs=[input_image, resolution, value_threshold, distance_threshold], outputs=[gallery])

    with gr.Tab("MIDAS Depth"):
        with gr.Row():
            gr.Markdown("## MIDAS Depth")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=384, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=midas, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("ZOE Depth"):
        with gr.Row():
            gr.Markdown("## Zoe Depth")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=zoe, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("Normal Bae"):
        with gr.Row():
            gr.Markdown("## Normal Bae")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=normalbae, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("DWPose"):
        with gr.Row():
            gr.Markdown("## DWPose")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=dwpose, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("Openpose"):
        with gr.Row():
            gr.Markdown("## Openpose")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                hand_and_face = gr.Checkbox(label="Hand and Face", value=False)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=openpose, inputs=[input_image, resolution, hand_and_face], outputs=[gallery])

    ## TAB LINEART
    with gr.Tab("Lineart"):
        with gr.Row():
            gr.Markdown("## Lineart \n<p>Check Invert to use with Mochi Diffusion.  Inverted image can also be created here for use with ControlNet Scribble.")
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                preprocessor_name = gr.Radio(label="Preprocessor", show_label=False, choices=["Lineart", "Lineart Coarse", "Lineart Anime"], type="value", value="Lineart")
                invert = gr.Checkbox(label="Invert", value=True)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_btn_lineart = gr.Button("Run")
            with gr.Column():
                gallery = gr.Gallery(label="Generated images", show_label=False, interactive=False, format="png", elem_id="output_gallery", elem_classes="output-gallery", columns=[3], rows=[2], object_fit="contain", height="auto", type="filepath")

        run_btn_lineart.click(fn=lineart, inputs=[input_image, resolution, preprocessor_name, invert, gallery], outputs=[gallery])

    with gr.Tab("InPaint"):
        with gr.Row():
            gr.Markdown("## InPaint")
        with gr.Row():
            with gr.Column():
                input_image = gr.ImageMask(sources="upload", type="numpy", height="auto")
                invert = gr.Checkbox(label="Invert Mask", value=False)
                run_button = gr.Button("Run")

            with gr.Column():

                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=inpaint, inputs=[input_image, invert], outputs=[gallery])

    #    with gr.Row():
    #        gr.Markdown("## Uniformer Segmentation")
    #    with gr.Row():
    #        with gr.Column():
    #            input_image = gr.Image(source='upload', type="numpy")
    #            resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
    #            run_button = gr.Button(label="Run")
    #        with gr.Column():
    #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
    #    run_button.click(fn=uniformer, inputs=[input_image, resolution], outputs=[gallery])

    #    with gr.Row():
    #        gr.Markdown("## Oneformer COCO Segmentation")
    #    with gr.Row():
    #        with gr.Column():
    #            input_image = gr.Image(source='upload', type="numpy")
    #            input_image = gr.Image(label="Input Image", type="numpy", height=512)
    #            resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
    #            run_button = gr.Button("Run")
    #            run_button = gr.Button(label="Run")
    #        with gr.Column():
    #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
    #            gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
    #    run_button.click(fn=oneformer_coco, inputs=[input_image, resolution], outputs=[gallery])

    #    with gr.Row():
    #        gr.Markdown("## Oneformer ADE20K Segmentation")
    #    with gr.Row():
    #        with gr.Column():
    #            input_image = gr.Image(source='upload', type="numpy")
    #            input_image = gr.Image(label="Input Image", type="numpy", height=512)
    #            resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=640, step=64)
    #            run_button = gr.Button("Run")
    #            run_button = gr.Button(label="Run")
    #        with gr.Column():
    #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
    #            gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
    #    run_button.click(fn=oneformer_ade20k, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("Content Shuffle"):
        with gr.Row():
            gr.Markdown("## Content Shuffle")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")
        run_button.click(fn=content_shuffler, inputs=[input_image, resolution], outputs=[gallery])

    with gr.Tab("Color Shuffle"):
        with gr.Row():
            gr.Markdown("## Color Shuffle")
        with gr.Row():
            with gr.Column():
                #            input_image = gr.Image(source='upload', type="numpy")
                input_image = gr.Image(label="Input Image", type="numpy", height=512)
                resolution = gr.Slider(label="resolution", minimum=256, maximum=1024, value=512, step=64)
                run_button = gr.Button("Run")
            #            run_button = gr.Button(label="Run")
            with gr.Column():
                #            gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")
                gallery = gr.Gallery(label="Generated images", show_label=False, height="auto")

    run_button.click(fn=color_shuffler, inputs=[input_image, resolution], outputs=[gallery])
"""

demo.launch()