##!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time    : 2024-01-29
# @Author  : Junjie He
import json
import os
import time
import uuid

import cv2
import gradio as gr
import numpy as np
from PIL import Image
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from src.generation import call_generation
from src.util import upload_np_2_oss, upload_json_string_2_oss, upload_preprocess, merge_images

universal_matting = pipeline(Tasks.universal_matting, model='damo/cv_unet_universal-matting')

img = "assets/image_gallery_en/"
files = os.listdir(img)
files = [file for file in files if file.lower().endswith(('.png', '.jpg', '.jpeg'))]
files = sorted(files)
basic_usage = []
showcases = []
for idx, name in enumerate(files):
    temp = os.path.join(os.path.dirname(__file__), img, name)
    if idx < 4:
        basic_usage.append(temp)
    else:
        showcases.append(temp)

# - - - - - examples  - - - - -  #
ep = "assets/examples"
# Layout, Style, Color, Subject, Prompt，Strict Layout Edge，Layout Content Scale, Automatic Image Matting
image_examples = [
    [0, f"{ep}/00_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0., False],
    [1, f"{ep}/01_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False],
    [2, f"{ep}/empty.png", f"{ep}/02_style.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False],
    [3, f"{ep}/empty.png", f"{ep}/03_style.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False],
    [4, f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/04_color.png", f"{ep}/empty.png", "A moose, Merry Christmas", True, 0.8, False],
    [5, f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/05_color.png", f"{ep}/empty.png", "A photo about cherry blossom",
     True, 0.8, False],
    [6, f"{ep}/06_layout.png", f"{ep}/06_style.jpeg", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False],
    [7, f"{ep}/07_layout.jpeg", f"{ep}/07_style.jpeg", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False],
    [8, f"{ep}/empty.png", f"{ep}/08_style.png", f"{ep}/08_color.jpeg", f"{ep}/empty.png", "", True, 0.8, False],
    [9, f"{ep}/empty.png", f"{ep}/09_style.png", f"{ep}/empty.png", f"{ep}/base_image1.jpeg", "", True, 0.8, True],
    [10, f"{ep}/10_layout.png", f"{ep}/10_style.png", f"{ep}/empty.png", f"{ep}/base_image2.png", "", True, 0.8, False],
    [11, f"{ep}/11_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/base_image4.png", "", False, 0.8, False],
    [12, f"{ep}/12_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/base_image3.png", "", False, 0.8, False],
]

example_images = [
    [0, f"{ep}/layout_image1.jpeg", None, None, None],
    [1, f"{ep}/layout_image1.jpeg", None, None, None],
    [2, None, f"{ep}/style_image1.jpeg", None, None],
    [3, None, f"{ep}/style_image1.jpeg", None, None],
    [4, None, None, f"{ep}/color_image1.jpeg", None],
    [5, None, None, f"{ep}/color_image3.jpeg", None],
    [6, f"{ep}/layout_image2.jpeg", f"{ep}/style_image2.jpeg", None, None],
    [7, f"{ep}/layout_image3.jpeg", f"{ep}/style_image3.jpeg", None, None],
    [8, None, f"{ep}/style_image4.jpeg", f"{ep}/color_image2.jpeg", None],
    [9, None, f"{ep}/style_image6.jpeg", None, f"{ep}/09_base.png"],
    [10, f"{ep}/layout_image5.jpeg", f"{ep}/style_image5.jpeg", None, f"{ep}/10_base.png"],
    [11, f"{ep}/layout_image7.jpeg", None, None, f"{ep}/11_base.png"],
    [12, f"{ep}/layout_image6.jpeg", None, None, f"{ep}/12_base.png"],
]

example_masks = [
    [0, None, None, None],
    [1, f"{ep}/layout_image1_mask.png", None, None],
    [2, None, f"{ep}/style_image1_mask.png", None],
    [3, None, f"{ep}/style_image1_mask2.png", None],
    [4, None, None, None],
    [5, None, None, f"{ep}/color_image3_mask.png"],
    [6, None, None, None],
    [7, None, None, None],
    [8, None, None, None],
    [9, None, f"{ep}/style_image6_mask.png", None],
    [10, f"{ep}/layout_image5_mask.png", None, None],
    [11, None, None, None],
    [12, f"{ep}/layout_image6_mask.png", None, None],
]


def process_example(example_idx, pil_layout_image, pil_style_image, pil_color_image, pil_base_image_rgba,
                    prompt, strict_edge, layout_scale, preprocess_base_image):
    _, layout_image, style_image, color_image, base_image_rgba = example_images[example_idx]
    _, layout_mask, style_mask, color_mask = example_masks[example_idx]

    pil_layout_image = None if layout_image is None else pil_layout_image
    pil_style_image = None if style_image is None else pil_style_image
    pil_color_image = None if color_image is None else pil_color_image
    pil_base_image_rgba = None if base_image_rgba is None else pil_base_image_rgba

    return pil_layout_image, layout_mask, pil_style_image, style_mask, pil_color_image, color_mask, \
           pil_base_image_rgba, prompt, strict_edge, layout_scale, preprocess_base_image


def process(pil_base_image_rgba=None, preprocess_base_image=False,
            pil_layout_image_dict=None, layout_scale=1.0, edge_consistency=0.5,
            strict_edge=False,
            pil_color_image_dict=None, color_scale=1.0,
            pil_style_image_dict=None, style_scale=1.0, prompt="best quality", negative_prompt="",
            pil_layout_mask=None, pil_style_mask=None, pil_color_mask=None):
    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())

    output_aspect_ratio = 1.
    matting_flag = False

    if pil_base_image_rgba is None:
        base_image_url = ""
        pil_fg_mask = None
    else:
        if preprocess_base_image:
            matting_flag = True
            orig_image = np.array(pil_base_image_rgba)
            orig_alpha = np.array(pil_base_image_rgba)[..., -1]
            matting_alpha = universal_matting(pil_base_image_rgba)[OutputKeys.OUTPUT_IMG][..., -1]
            orig_image[..., -1] = ((matting_alpha > 200) * (orig_alpha > 200) * 255.).astype(np.uint8)
            pil_base_image_rgba = Image.fromarray(orig_image)
        pil_fg_mask = pil_base_image_rgba.split()[-1]

        w, h = pil_base_image_rgba.size
        output_aspect_ratio = max(1.0 * w / h, 1.0 * h / w)
        if output_aspect_ratio > 2:
            raise gr.Error("Input of subject images with aspect ratio exceeding 2 is not supported")
        if min(w, h) > 1536:
            raise gr.Error("Input of subject images with the shorter side exceeding 1536 pixels is not supported")
        base_image_url = upload_np_2_oss(np.array(pil_base_image_rgba), request_id + "_base.png")

    if pil_layout_image_dict is None:
        layout_image_url = ""
    else:
        np_layout_image = np.array(pil_layout_image_dict["image"].convert("RGBA"))
        np_layout_image, np_layout_alpha = np_layout_image[..., :3], np_layout_image[..., 3]
        np_layout_mask = np.array(pil_layout_image_dict["mask"].convert("L"))
        if pil_layout_mask is None:
            np_layout_mask = ((np_layout_alpha > 127) * (np_layout_mask < 127) * 255.).astype(np.uint8)
        else:
            np_layout_mask = ((np_layout_alpha > 127) * (np_layout_mask < 127) *
                              (np.array(pil_layout_mask) > 127) * 255.).astype(np.uint8)
        layout_image_url = upload_np_2_oss(
            np.concatenate([np_layout_image, np_layout_mask[..., None]], axis=-1), request_id + "_layout.png"
        )
        if pil_base_image_rgba is None:
            h, w, c = np_layout_image.shape
            output_aspect_ratio = max(1.0 * w / h, 1.0 * h / w)
            if output_aspect_ratio > 2:
                raise gr.Error("Input of layout images with aspect ratio exceeding 2 is not supported")
            if min(w, h) > 1536:
                raise gr.Error("Input of layout images with the shorter side exceeding 1536 pixels is not supported")

    if pil_style_image_dict is None:
        style_image_url = ""
    else:
        np_style_image = np.array(pil_style_image_dict["image"].convert("RGBA"))
        np_style_image, np_style_alpha = np_style_image[..., :3], np_style_image[..., 3]
        np_style_mask = np.array(pil_style_image_dict["mask"].convert("L"))
        if pil_style_mask is None:
            np_style_mask = ((np_style_alpha > 127) * (np_style_mask < 127) * 255.).astype(np.uint8)
        else:
            np_style_mask = ((np_style_alpha > 127) * (np_style_mask < 127) *
                             (np.array(pil_style_mask) > 127) * 255.).astype(np.uint8)
        style_image_url = upload_np_2_oss(
            np.concatenate([np_style_image, np_style_mask[..., None]], axis=-1), request_id + "_style.png"
        )

    if pil_color_image_dict is None:
        color_image_url = ""
    else:
        np_color_image = np.array(pil_color_image_dict["image"].convert("RGBA"))
        np_color_image, np_color_alpha = np_color_image[..., :3], np_color_image[..., 3]
        np_color_mask = np.array(pil_color_image_dict["mask"].convert("L"))
        if pil_color_mask is None:
            np_color_mask = ((np_color_alpha > 127) * (np_color_mask < 127) * 255.).astype(np.uint8)
        else:
            np_color_mask = ((np_color_alpha > 127) * (np_color_mask < 127) *
                             (np.array(pil_color_mask) > 127) * 255.).astype(np.uint8)
        color_image_url = upload_np_2_oss(
            np.concatenate([np_color_image, np_color_mask[..., None]], axis=-1), request_id + "_color.png"
        )

    res = call_generation(base_image_url=base_image_url, layout_image_url=layout_image_url,
                          color_image_url=color_image_url, style_image_url=style_image_url,
                          strict_edge=int(strict_edge), layout_scale=int(layout_scale * 10),
                          edge_consistency=int(edge_consistency * 10), color_scale=int(color_scale * 10),
                          style_scale=int(style_scale * 10), prompt=prompt, negative_prompt=negative_prompt,
                          output_aspect_ratio=output_aspect_ratio)

    for idx, r in enumerate(res):
        upload_np_2_oss(np.array(r), request_id + f"_{idx}.jpg")

    if matting_flag:
        res.append(pil_base_image_rgba)
    return res, request_id, True, pil_fg_mask


if __name__ == "__main__":

    block = gr.Blocks(
        title="TransferAnything",
        css="assets/css/style.css",
        theme=gr.themes.Soft(
            radius_size=gr.themes.sizes.radius_none,
            text_size=gr.themes.sizes.text_md
        )).queue(concurrency_count=3)
    with block:
        with gr.Row():
            with gr.Column():
                gr.HTML(f"""
                        <div style="text-align: center;">
                            <h1> TransferAnything: Enabling Versatile Visual Information Transfer for Creative Image Synthesis </h1>
                            <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
                                <a href='https://aigcdesigngroup.github.io/transfer-anything/'><img src='https://img.shields.io/badge/Project_Page-TransferAnything-green' alt='Project Page'></a>
                            </div>
                            </br>
                            <h3> TransferAnything supports transferring various visual information from any area of any image to create new compositions, offering higher freedom and flexibility in image synthesis. Currently, it supports the transfer of layout, color, style, and pixel content, with more visual information transfer capabilities under continuous development. </h3>
                            </br>
                        </div>
                """)
                # <h5 style="margin: 0; color: red">如果你认为该项目有所帮助的话，不妨给我们Github点个Star以便获取最新的项目进展.</h5>
        with gr.Tabs(elem_classes=["Tab"]):
            with gr.TabItem("Image Gallery"):
                gr.Gallery(label="Basic Usage", value=basic_usage, height=400, columns=4, object_fit="scale-down")
                gr.Gallery(label="Advanced Combinations", value=showcases, height=1200, columns=4, object_fit="scale-down")
            with gr.TabItem("Image Creation"):
                with gr.Row():
                    with gr.Column(scale=1):
                        ...
                    with gr.Column(scale=3):
                        gr.Image(value="assets/banner/banner.png", width=1024, show_label=False,
                                 show_download_button=False)
                    with gr.Column(scale=1):
                        ...
                with gr.Accordion(label="🧭 Instructions:", open=True, elem_id="accordion"):
                    with gr.Row(equal_height=True):
                        #     with gr.Row(elem_id="ShowCase"):
                        #             gr.Image(value="assets/banner/ra.gif")
                        gr.Markdown("""
                        - ⭐️ <b>step1：</b> (Optional) Upload or select a set of images from the examples for the "Layout", "Style", and "Color" reference. Mix and match freely, no need to select all.
                        - ⭐️ <b>step2：</b> (Optional) Use brush to erase areas in the layout, style, and color reference images (if present) that you do not want transferred.
                        - ⭐️ <b>step3：</b> (Optional) Upload an image with the subject you wish to preserve at the pixel level to the "Subject" component, with a transparent channel required (or upload a regular RGB image and select automatic image matting, which will automatically segment the subject area for you).
                        - ⭐️ <b>step4：</b> Click "Run" to start the generation process.
                        - ⭐️ <b>step5：</b> (Optional) Additionally, prompt input is supported, as well as control over advanced parameters such as layout edge consistency and conditional weights. Feel free to try these features.
                        """)
                    # with gr.Row(equal_height=True):
                    #     gr.Markdown("""❤️ 温馨提示：上传图片最大边不超过1024，宽高比为1：1，将会生成更快哟""")
                with gr.Row():
                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Layout (Optional)"):
                                pil_layout_image_dict = gr.ImageMask(source='upload', type="pil", show_label=False,
                                                                     image_mode="RGBA")
                                pil_layout_image = gr.Image(interactive=False, type="pil", visible=False, label="Layout")
                                pil_layout_mask = gr.Image(interactive=False, type="pil", visible=False, label="Layout Mask",
                                                           image_mode="L")
                                with gr.Box():
                                    with gr.Accordion(label="Layout Parameters", open=False, elem_id="accordion"):
                                        with gr.TabItem("Layout Edge"):
                                            strict_edge = gr.Checkbox(label="Strict", value=True)
                                            strict_edge_mirror = gr.Checkbox(label="Strict Layout Edge", visible=False)
                                            edge_consistency = gr.Slider(label="Degree of Consistency (If Not Strict)", minimum=0.0,
                                                                         maximum=1.0,
                                                                         step=0.1, value=0.8, interactive=True)
                                        with gr.TabItem("Layout Content"):
                                            layout_scale = gr.Slider(label="Scale", minimum=0.0, maximum=1.0, step=0.1,
                                                                     value=0.8,
                                                                     interactive=True)
                                            layout_scale_mirror = gr.Slider(label="Layout Content Scale", visible=False)

                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Style (Optional)"):
                                pil_style_image_dict = gr.ImageMask(source='upload', type="pil", show_label=False,
                                                                    image_mode="RGBA")
                                pil_style_image = gr.Image(interactive=False, type="pil", visible=False, label="Style")
                                pil_style_mask = gr.Image(interactive=False, type="pil", visible=False, label="Style Mask",
                                                          image_mode="L")
                                with gr.Box():
                                    with gr.Accordion(label="Style Parameters", open=False, elem_id="accordion"):
                                        style_scale = gr.Slider(label="Scale", minimum=0.0, maximum=1.0, step=0.1,
                                                                value=0.8, interactive=True)
                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Color (Optional, recommended for use with prompt)"):
                                pil_color_image_dict = gr.ImageMask(source='upload', type="pil", show_label=False,
                                                                    image_mode="RGBA")
                                pil_color_image = gr.Image(interactive=False, type="pil", visible=False, label="Color")
                                pil_color_mask = gr.Image(interactive=False, type="pil", visible=False, label="Color Mask",
                                                          image_mode="L")
                                with gr.Box():
                                    with gr.Accordion(label="Color Parameters", open=False, elem_id="accordion"):
                                        color_scale = gr.Slider(label="Scale", minimum=0.0, maximum=1.0, step=0.1,
                                                                value=0.8, interactive=True)
                with gr.Row():
                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Subject (Optional)"):
                                pil_base_image_rgba = gr.Image(source='upload',
                                                               interactive=True, show_label=False,
                                                               type="pil", image_mode="RGBA", tool="editor")
                                pil_base_image_rgba_mirror = gr.Image(label="Subject", image_mode="RGBA", visible=False)
                                with gr.Box():
                                    preprocess_base_image = gr.Checkbox(label="Automatic Image Matting")
                                    pil_fg_mask = gr.Image(interactive=False, type="pil", image_mode="L", visible=False)
                        run_button = gr.Button("Run", elem_id="btn")
                        with gr.Accordion("", open=True, elem_id="accordion1"):
                            prompt = gr.Textbox(value="", label='Prompt', lines=1, interactive=True)
                            prompt_mirror = gr.Textbox(label='Prompt', visible=False)
                            negative_prompt = gr.Textbox(value="", label='Negative prompt', lines=1,
                                                         interactive=True)
                    with gr.Column(scale=2, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Outputs"):
                                result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery",
                                                            preview=True)
                                recommend = gr.Button("Recommend results to Image Gallery", elem_id="recBut")
                                request_id = gr.State(value="")
                                gallery_flag = gr.State(value=False)
                with gr.Row():
                    with gr.Box():
                        example_idx = gr.Slider(label="Index", visible=False, value=0)
                        example = gr.Examples(
                            label="Input Examples",
                            examples=image_examples,
                            inputs=[example_idx,
                                    pil_layout_image, pil_style_image, pil_color_image, pil_base_image_rgba_mirror,
                                    prompt_mirror, strict_edge_mirror, layout_scale_mirror, preprocess_base_image],
                            outputs=[pil_layout_image_dict, pil_layout_mask, pil_style_image_dict, pil_style_mask,
                                     pil_color_image_dict, pil_color_mask, pil_base_image_rgba,
                                     prompt, strict_edge, layout_scale, preprocess_base_image],
                            fn=process_example,
                            run_on_click=True,
                            examples_per_page=20
                        )


        def upload_to_img_gallery(pil_base_image_rgba, pil_layout_image_dict, pil_style_image_dict,
                                  pil_color_image_dict, pil_fg_mask, prompt, negative_prompt, res, re_id, flag,
                                  strict_edge, edge_consistency, layout_scale, style_scale, color_scale,
                                  preprocess_base_image):
            if flag:
                np_out_base_image, np_out_layout_image, np_out_style_image, np_out_color_image = upload_preprocess(
                    pil_base_image_rgba, pil_layout_image_dict, pil_style_image_dict, pil_color_image_dict, pil_fg_mask)
                np_out_images = [np_out_base_image, np_out_layout_image, np_out_style_image, np_out_color_image]
                for idx, r in enumerate(res):
                    if idx < 4:
                        r = cv2.imread(r['name'])
                        r = cv2.cvtColor(r, cv2.COLOR_BGR2RGB)
                        upload_np_2_oss(merge_images(*np_out_images, r, prompt, negative_prompt),
                                        name=re_id + f"_merge_{idx}.jpg", gallery=True)
                config = dict(
                    strict_edge=strict_edge,
                    edge_consistency=edge_consistency,
                    layout_scale=layout_scale,
                    style_scale=style_scale,
                    color_scale=color_scale,
                    preprocess_base_image=preprocess_base_image
                )
                upload_json_string_2_oss(json.dumps(config), name=re_id + f"_config.txt", gallery=True)

                flag = False
                gr.Info("Images have been uploaded and await review.")
            else:
                gr.Info("No images to recommend, or already suggested once.")
            return flag


        recommend.click(
            upload_to_img_gallery,
            [pil_base_image_rgba, pil_layout_image_dict, pil_style_image_dict, pil_color_image_dict, pil_fg_mask,
             prompt, negative_prompt, result_gallery, request_id, gallery_flag, strict_edge, edge_consistency,
             layout_scale, style_scale, color_scale, preprocess_base_image],
            [gallery_flag]
        )

        ips = [pil_base_image_rgba, preprocess_base_image,
               pil_layout_image_dict, layout_scale, edge_consistency, strict_edge,
               pil_color_image_dict, color_scale,
               pil_style_image_dict, style_scale, prompt, negative_prompt,
               pil_layout_mask, pil_style_mask, pil_color_mask]
        run_button.click(fn=process, inputs=ips, outputs=[result_gallery, request_id, gallery_flag, pil_fg_mask])

    block.launch(share=True)