# Copyright 2023-2025 Marigold Team, ETH Zürich. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# --------------------------------------------------------------------------
# More information about Marigold:
#   https://marigoldmonodepth.github.io
#   https://marigoldcomputervision.github.io
# Efficient inference pipelines are now part of diffusers:
#   https://huggingface.co/docs/diffusers/using-diffusers/marigold_usage
#   https://huggingface.co/docs/diffusers/api/pipelines/marigold
# Examples of trained models and live demos:
#   https://huggingface.co/prs-eth
# Related projects:
#   https://rollingdepth.github.io/
#   https://marigolddepthcompletion.github.io/
# Citation (BibTeX):
#   https://github.com/prs-eth/Marigold#-citation
# If you find Marigold useful, we kindly ask you to cite our papers.
# --------------------------------------------------------------------------
import os
import tempfile

import gradio as gr
from PIL import Image

from extrude import extrude_depth_3d
from gradio_patches.examples import Examples

default_seed = 2024
default_batch_size = 4
default_bas_plane_near = 0.0
default_bas_plane_far = 1.0
default_bas_embossing = 20
default_bas_size_longest_px = 512
default_bas_size_longest_cm = 10
default_bas_filter_size = 3
default_bas_frame_thickness = 5
default_bas_frame_near = 1
default_bas_frame_far = 1


def process_bas(
    path_input_depth,
    path_input_rgb=None,
    plane_near=default_bas_plane_near,
    plane_far=default_bas_plane_far,
    embossing=default_bas_embossing,
    size_longest_px=default_bas_size_longest_px,
    size_longest_cm=default_bas_size_longest_cm,
    filter_size=default_bas_filter_size,
    frame_thickness=default_bas_frame_thickness,
    frame_near=default_bas_frame_near,
    frame_far=default_bas_frame_far,
):
    if path_input_depth is None:
        raise gr.Error(
            "Missing image in the first pane: upload a file or use one from the gallery below."
        )

    input_depth = Image.open(path_input_depth)
    if input_depth.mode not in ("I", "I;16"):
        raise gr.Error(
            f"Input depth must be a 16-bit PNG image of a depth map, found {input_depth.mode}"
        )
    depth_longest_px = max(input_depth.size)

    input_rgb = None
    if path_input_rgb is not None:
        input_rgb = Image.open(path_input_rgb).convert("RGB")
        if (
            input_depth.size[0] * input_rgb.size[1]
            != input_depth.size[0] * input_rgb.size[1]
        ):
            raise gr.Error(
                f"Inputs have incompatible dimensions: {input_depth.size} and {input_rgb.size}"
            )

    if plane_near >= plane_far:
        raise gr.Error("NEAR plane must have a value smaller than the FAR plane")

    name_base, name_ext = os.path.splitext(os.path.basename(path_input_depth))
    print(f"Processing bas-relief {name_base}{name_ext}")

    path_output_dir = tempfile.mkdtemp()

    def _process_3d(
        size_longest_px,
        filter_size,
        vertex_colors,
        scene_lights,
        output_model_scale=None,
        prepare_for_3d_printing=False,
        zip_outputs=False,
    ):
        image_new_w = size_longest_px * input_depth.width // depth_longest_px
        image_new_h = size_longest_px * input_depth.height // depth_longest_px
        image_new_sz = (image_new_w, image_new_h)

        path_depth_new = os.path.join(
            path_output_dir, f"{name_base}_depth_{size_longest_px}.png"
        )
        (
            input_depth.convert(mode="F")
            .resize(image_new_sz, Image.BILINEAR)
            .convert("I")
            .save(path_depth_new)
        )
        path_rgb_new = None
        if input_rgb is not None:
            path_rgb_new = os.path.join(
                path_output_dir, f"{name_base}_rgb_{size_longest_px}{name_ext}"
            )
            input_rgb.resize(image_new_sz, Image.LANCZOS).save(path_rgb_new)

        path_glb, path_stl, path_obj = extrude_depth_3d(
            path_depth_new,
            path_rgb_new,
            output_model_scale=(
                size_longest_cm * 10
                if output_model_scale is None
                else output_model_scale
            ),
            filter_size=filter_size,
            coef_near=plane_near,
            coef_far=plane_far,
            emboss=embossing / 100,
            f_thic=frame_thickness / 100,
            f_near=frame_near / 100,
            f_back=frame_far / 100,
            vertex_colors=vertex_colors,
            scene_lights=scene_lights,
            prepare_for_3d_printing=prepare_for_3d_printing,
            zip_outputs=zip_outputs,
        )

        return path_glb, path_stl, path_obj

    path_viewer_glb, _, _ = _process_3d(
        256, filter_size, vertex_colors=False, scene_lights=True, output_model_scale=1
    )
    path_files_glb, path_files_stl, path_files_obj = _process_3d(
        size_longest_px,
        filter_size,
        vertex_colors=True,
        scene_lights=False,
        prepare_for_3d_printing=True,
        zip_outputs=True,
    )

    return path_viewer_glb, [path_files_glb, path_files_stl, path_files_obj]


with gr.Blocks(
    title="Depth To 3D Print",
    css="""
        #download {
            height: 118px;
        }
        .viewport {
            aspect-ratio: 4/3;
        }
        h1 {
            text-align: center;
            display: block;
        }
        h2 {
            text-align: center;
            display: block;
        }
        h3 {
            text-align: center;
            display: block;
        }
        a {
            display: inline-block;
        }
        .md_feedback li {
            margin-bottom: 0px !important;
        }
        ol {
            margin: 0 auto;
            width: fit-content;
            text-align: left;
        }
        ol li {
            margin-bottom: 0px;
        }
    """,
    head="""
        <script async src="https://www.googletagmanager.com/gtag/js?id=G-1FWSVCGZTG"></script>
        <script>
            window.dataLayer = window.dataLayer || [];
            function gtag() {dataLayer.push(arguments);}
            gtag('js', new Date());
            gtag('config', 'G-1FWSVCGZTG');
        </script>
    """,
) as demo:
    gr.Markdown(
        """
        # Depth To 3D Print
        <p align="center">
        <a title="Get Depth" href="https://huggingface.co/spaces/prs-eth/marigold" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
            <img src="https://img.shields.io/badge/🤗%20Create%20Your%20-Depth%20from%20Image-blue" alt="Get Depth">
        </a>
        <a title="Website" href="https://marigoldcomputervision.github.io/" target="_blank" rel="noopener noreferrer" 
                style="display: inline-block;">
            <img src="https://img.shields.io/badge/%F0%9F%A4%8D%20Project%20-Website-af2928" alt="Website Badge">
        </a>
        <a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" 
                style="display: inline-block;">
            <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
        </a><br>
            Start exploring the interactive bas-relief examples at the bottom of the page! 
            To create your own watertight 3D-printable bas-relief depth map:
        </p>
        <ol>
            <li>Click the "Create Your Depth from Image" badge above</li>
            <li>Upload the input there and download the 16-bit PNG (right click save)</li>
            <li>Return to this demo</li>
            <li>Load your 16-bit depth PNG file in the top-left pane</li>
            <li>Click "Create 3D"</li>
        </ol>
        """
    )
    with gr.Row():
        with gr.Column():
            bas_depth = gr.Image(
                label="Depth",
                type="filepath",
                format="png",
                image_mode=None,
                sources=["upload", "clipboard"],
                show_fullscreen_button=False,
            )
            bas_rgb = gr.Image(
                label="Image (optional)",
                type="filepath",
                sources=["upload", "clipboard"],
                show_fullscreen_button=False,
            )
            with gr.Row():
                bas_submit_btn = gr.Button(value="Create 3D", variant="primary")
                bas_reset_btn = gr.Button(value="Reset")
            with gr.Accordion("3D printing demo: Main options", open=True):
                bas_plane_near = gr.Slider(
                    label="Near plane",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.001,
                    value=default_bas_plane_near,
                )
                bas_plane_far = gr.Slider(
                    label="Far plane",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.001,
                    value=default_bas_plane_far,
                )
                bas_embossing = gr.Slider(
                    label="Embossing level",
                    minimum=0,
                    maximum=100,
                    step=1,
                    value=default_bas_embossing,
                )
            with gr.Accordion("3D printing demo: Advanced options", open=False):
                bas_size_longest_px = gr.Slider(
                    label="Longest side (px)",
                    minimum=256,
                    maximum=1024,
                    step=256,
                    value=default_bas_size_longest_px,
                )
                bas_size_longest_cm = gr.Slider(
                    label="Longest side (cm)",
                    minimum=1,
                    maximum=100,
                    step=1,
                    value=default_bas_size_longest_cm,
                )
                bas_filter_size = gr.Slider(
                    label="Smooth radius",
                    minimum=1,
                    maximum=5,
                    step=2,
                    value=default_bas_filter_size,
                )
                bas_frame_thickness = gr.Slider(
                    label="Frame thickness",
                    minimum=0,
                    maximum=100,
                    step=1,
                    value=default_bas_frame_thickness,
                )
                bas_frame_near = gr.Slider(
                    label="Near offset",
                    minimum=-100,
                    maximum=100,
                    step=1,
                    value=default_bas_frame_near,
                )
                bas_frame_far = gr.Slider(
                    label="Far offset",
                    minimum=1,
                    maximum=10,
                    step=1,
                    value=default_bas_frame_far,
                )
        with gr.Column():
            bas_output_viewer = gr.Model3D(
                camera_position=(75.0, 90.0, 1.25),
                elem_classes="viewport",
                label="3D preview",
                interactive=False,
            )
            bas_output_files = gr.Files(
                label="3D models",
                elem_id="download",
                interactive=False,
            )
    Examples(
        fn=process_bas,
        examples=[
            [
                "files/einstein_depth_16bit.png",  # input_depth
                "files/einstein_rgb.jpg",  # input_rgb
                0.0,  # plane_near
                0.5,  # plane_far
                50,  # embossing
                512,  # size_longest_px
                10,  # size_longest_cm
                3,  # filter_size
                5,  # frame_thickness
                -25,  # frame_near
                1,  # frame_far
            ],
        ],
        inputs=[
            bas_depth,
            bas_rgb,
            bas_plane_near,
            bas_plane_far,
            bas_embossing,
            bas_size_longest_px,
            bas_size_longest_cm,
            bas_filter_size,
            bas_frame_thickness,
            bas_frame_near,
            bas_frame_far,
        ],
        outputs=[bas_output_viewer, bas_output_files],
        cache_examples=True,
        directory_name="outputs",
    )

    bas_submit_btn.click(
        fn=process_bas,
        inputs=[
            bas_depth,
            bas_rgb,
            bas_plane_near,
            bas_plane_far,
            bas_embossing,
            bas_size_longest_px,
            bas_size_longest_cm,
            bas_filter_size,
            bas_frame_thickness,
            bas_frame_near,
            bas_frame_far,
        ],
        outputs=[bas_output_viewer, bas_output_files],
    )

    bas_reset_btn.click(
        fn=lambda: (
            gr.Button(interactive=True),
            None,
            None,
            None,
            None,
            default_bas_plane_near,
            default_bas_plane_far,
            default_bas_embossing,
            default_bas_size_longest_px,
            default_bas_size_longest_cm,
            default_bas_filter_size,
            default_bas_frame_thickness,
            default_bas_frame_near,
            default_bas_frame_far,
        ),
        inputs=[],
        outputs=[
            bas_submit_btn,
            bas_depth,
            bas_rgb,
            bas_output_viewer,
            bas_output_files,
            bas_plane_near,
            bas_plane_far,
            bas_embossing,
            bas_size_longest_px,
            bas_size_longest_cm,
            bas_filter_size,
            bas_frame_thickness,
            bas_frame_near,
            bas_frame_far,
        ],
    )


if __name__ == "__main__":
    demo.queue(
        api_open=False,
    ).launch(
        share=True,
        server_port=7860,
    )