#!/usr/bin/env python

from __future__ import annotations

import pathlib

import gradio as gr

from dualstylegan import Model

DESCRIPTION = """# Portrait Style Transfer with <a href="https://github.com/williamyang1991/DualStyleGAN">DualStyleGAN</a>

<img id="overview" alt="overview" src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" />
"""


def get_style_image_url(style_name: str) -> str:
    base_url = "https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images"
    filenames = {
        "cartoon": "cartoon_overview.jpg",
        "caricature": "caricature_overview.jpg",
        "anime": "anime_overview.jpg",
        "arcane": "Reconstruction_arcane_overview.jpg",
        "comic": "Reconstruction_comic_overview.jpg",
        "pixar": "Reconstruction_pixar_overview.jpg",
        "slamdunk": "Reconstruction_slamdunk_overview.jpg",
    }
    return f"{base_url}/{filenames[style_name]}"


def get_style_image_markdown_text(style_name: str) -> str:
    url = get_style_image_url(style_name)
    return f'<center><img id="style-image" src="{url}" alt="style image"></center>'


def update_slider(choice: str) -> dict:
    max_vals = {
        "cartoon": 316,
        "caricature": 198,
        "anime": 173,
        "arcane": 99,
        "comic": 100,
        "pixar": 121,
        "slamdunk": 119,
    }
    return gr.Slider(maximum=max_vals[choice])


def update_style_image(style_name: str) -> dict:
    text = get_style_image_markdown_text(style_name)
    return gr.Markdown(value=text)


def set_example_image(example: list) -> dict:
    return gr.Image(value=example[0])


def set_example_styles(example: list) -> list[dict]:
    return [
        gr.Radio(value=example[0]),
        gr.Slider(value=example[1]),
    ]


def set_example_weights(example: list) -> list[dict]:
    return [
        gr.Slider(value=example[0]),
        gr.Slider(value=example[1]),
    ]


model = Model()

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Group():
        gr.Markdown(
            """## Step 1 (Preprocess Input Image)

- Drop an image containing a near-frontal face to the **Input Image**.
- If there are multiple faces in the image, hit the Edit button in the upper right corner and crop the input image beforehand.
- Hit the **Detect & Align Face** button.
- Hit the **Reconstruct Face** button.
- The final result will be based on this **Reconstructed Face**. So, if the reconstructed image is not satisfactory, you may want to change the input image.
"""
        )
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_image = gr.Image(label="Input Image", type="filepath")
                with gr.Row():
                    detect_button = gr.Button("Detect & Align Face")
            with gr.Column():
                with gr.Row():
                    aligned_face = gr.Image(label="Aligned Face", type="numpy", interactive=False)
                with gr.Row():
                    reconstruct_button = gr.Button("Reconstruct Face")
            with gr.Column():
                reconstructed_face = gr.Image(label="Reconstructed Face", type="numpy")
                instyle = gr.State()

        with gr.Row():
            paths = sorted(pathlib.Path("images").glob("*.jpg"))
            gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_image)

    with gr.Group():
        gr.Markdown(
            """## Step 2 (Select Style Image)

- Select **Style Type**.
- Select **Style Image Index** from the image table below.
"""
        )
        with gr.Row():
            with gr.Column():
                style_type = gr.Radio(label="Style Type", choices=model.style_types, value="cartoon")
                text = get_style_image_markdown_text("cartoon")
                style_image = gr.Markdown(value=text)
                style_index = gr.Slider(label="Style Image Index", minimum=0, maximum=316, step=1, value=26)

        with gr.Row():
            gr.Examples(
                examples=[
                    ["cartoon", 26],
                    ["caricature", 65],
                    ["arcane", 63],
                    ["pixar", 80],
                ],
                inputs=[style_type, style_index],
            )

    with gr.Group():
        gr.Markdown(
            """## Step 3 (Generate Style Transferred Image)

- Adjust **Structure Weight** and **Color Weight**.
- These are weights for the style image, so the larger the value, the closer the resulting image will be to the style image.
- Hit the **Generate** button.
"""
        )
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    structure_weight = gr.Slider(label="Structure Weight", minimum=0, maximum=1, step=0.1, value=0.6)
                with gr.Row():
                    color_weight = gr.Slider(label="Color Weight", minimum=0, maximum=1, step=0.1, value=1)
                with gr.Row():
                    structure_only = gr.Checkbox(label="Structure Only")
                with gr.Row():
                    generate_button = gr.Button("Generate")

            with gr.Column():
                result = gr.Image(label="Result")

        with gr.Row():
            gr.Examples(
                examples=[
                    [0.6, 1.0],
                    [0.3, 1.0],
                    [0.0, 1.0],
                    [1.0, 0.0],
                ],
                inputs=[structure_weight, color_weight],
            )

    detect_button.click(fn=model.detect_and_align_face, inputs=input_image, outputs=aligned_face)
    reconstruct_button.click(fn=model.reconstruct_face, inputs=aligned_face, outputs=[reconstructed_face, instyle])
    style_type.change(fn=update_slider, inputs=style_type, outputs=style_index)
    style_type.change(fn=update_style_image, inputs=style_type, outputs=style_image)
    generate_button.click(
        fn=model.generate,
        inputs=[
            style_type,
            style_index,
            structure_weight,
            color_weight,
            structure_only,
            instyle,
        ],
        outputs=result,
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()