#!/usr/bin/env python

from __future__ import annotations

import pathlib

import gradio as gr

from dualstylegan import Model

DESCRIPTION = """# Portrait Style Transfer with [DualStyleGAN](https://github.com/williamyang1991/DualStyleGAN)

<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'<img id="style-image" src="{url}" alt="style image">'


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)


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 **Preprocess** button.
- Choose the encoder version. Default is Z+ encoder which has better stylization performance. W+ encoder better reconstructs the input image to preserve more details.
- 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():
            encoder_type = gr.Radio(
                label="Encoder Type",
                choices=["Z+ encoder (better stylization)", "W+ encoder (better reconstruction)"],
                value="Z+ encoder (better stylization)",
            )
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    input_image = gr.Image(label="Input Image", type="filepath")
                with gr.Row():
                    preprocess_button = gr.Button("Preprocess")
            with gr.Column():
                with gr.Row():
                    aligned_face = gr.Image(label="Aligned Face", type="numpy", interactive=False)
            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=model.style_types[0])
                text = get_style_image_markdown_text("cartoon")
                style_image = gr.Markdown(value=text, latex_delimiters=[])
                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.
- Tips: For W+ encoder, better way of (Structure Only) is to uncheck (Structure Only) and set Color weight to 0.
- 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", value=False)
                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],
            )

    preprocess_button.click(
        fn=model.detect_and_align_face,
        inputs=[input_image],
        outputs=aligned_face,
    )
    aligned_face.change(
        fn=model.reconstruct_face,
        inputs=[aligned_face, encoder_type],
        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=20).launch()