#!/usr/bin/env python3
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# sparse gradio demo functions
# --------------------------------------------------------
import math
import gradio
import os
import numpy as np
import functools
import trimesh
import copy
from scipy.spatial.transform import Rotation
import tempfile
import shutil

from mast3r.cloud_opt.sparse_ga import sparse_global_alignment
from mast3r.cloud_opt.tsdf_optimizer import TSDFPostProcess

import mast3r.utils.path_to_dust3r  # noqa
from dust3r.image_pairs import make_pairs
from dust3r.utils.image import load_images
from dust3r.utils.device import to_numpy
from dust3r.viz import add_scene_cam, CAM_COLORS, OPENGL, pts3d_to_trimesh, cat_meshes
from dust3r.demo import get_args_parser as dust3r_get_args_parser

import matplotlib.pyplot as pl


class SparseGAState:
    def __init__(self, sparse_ga, should_delete=False, cache_dir=None, outfile_name=None):
        self.sparse_ga = sparse_ga
        self.cache_dir = cache_dir
        self.outfile_name = outfile_name
        self.should_delete = should_delete

    def __del__(self):
        if not self.should_delete:
            return
        if self.cache_dir is not None and os.path.isdir(self.cache_dir):
            shutil.rmtree(self.cache_dir)
        self.cache_dir = None
        if self.outfile_name is not None and os.path.isfile(self.outfile_name):
            os.remove(self.outfile_name)
        self.outfile_name = None


def get_args_parser():
    parser = dust3r_get_args_parser()
    parser.add_argument("--share", action="store_true")
    parser.add_argument(
        "--gradio_delete_cache",
        default=None,
        type=int,
        help="age/frequency at which gradio removes the file. If >0, matching cache is purged",
    )

    actions = parser._actions
    for action in actions:
        if action.dest == "model_name":
            action.choices = ["MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"]
    # change defaults
    parser.prog = "mast3r demo"
    return parser


def _convert_scene_output_to_glb(
    outfile,
    imgs,
    pts3d,
    mask,
    focals,
    cams2world,
    cam_size=0.05,
    cam_color=None,
    as_pointcloud=False,
    transparent_cams=False,
    silent=False,
):
    assert len(pts3d) == len(mask) <= len(imgs) <= len(cams2world) == len(focals)
    pts3d = to_numpy(pts3d)
    imgs = to_numpy(imgs)
    focals = to_numpy(focals)
    cams2world = to_numpy(cams2world)

    scene = trimesh.Scene()

    # full pointcloud
    if as_pointcloud:
        pts = np.concatenate([p[m.ravel()] for p, m in zip(pts3d, mask)]).reshape(-1, 3)
        col = np.concatenate([p[m] for p, m in zip(imgs, mask)]).reshape(-1, 3)
        valid_msk = np.isfinite(pts.sum(axis=1))
        pct = trimesh.PointCloud(pts[valid_msk], colors=col[valid_msk])
        scene.add_geometry(pct)
    else:
        meshes = []
        for i in range(len(imgs)):
            pts3d_i = pts3d[i].reshape(imgs[i].shape)
            msk_i = mask[i] & np.isfinite(pts3d_i.sum(axis=-1))
            meshes.append(pts3d_to_trimesh(imgs[i], pts3d_i, msk_i))
        mesh = trimesh.Trimesh(**cat_meshes(meshes))
        scene.add_geometry(mesh)

    # add each camera
    for i, pose_c2w in enumerate(cams2world):
        if isinstance(cam_color, list):
            camera_edge_color = cam_color[i]
        else:
            camera_edge_color = cam_color or CAM_COLORS[i % len(CAM_COLORS)]
        add_scene_cam(
            scene,
            pose_c2w,
            camera_edge_color,
            None if transparent_cams else imgs[i],
            focals[i],
            imsize=imgs[i].shape[1::-1],
            screen_width=cam_size,
        )

    rot = np.eye(4)
    rot[:3, :3] = Rotation.from_euler("y", np.deg2rad(180)).as_matrix()
    scene.apply_transform(np.linalg.inv(cams2world[0] @ OPENGL @ rot))
    if not silent:
        print("(exporting 3D scene to", outfile, ")")
    scene.export(file_obj=outfile)
    return outfile


def get_3D_model_from_scene(
    silent,
    scene_state,
    min_conf_thr=2,
    as_pointcloud=False,
    mask_sky=False,
    clean_depth=False,
    transparent_cams=False,
    cam_size=0.05,
    TSDF_thresh=0,
):
    """
    extract 3D_model (glb file) from a reconstructed scene
    """
    if scene_state is None:
        return None
    outfile = scene_state.outfile_name
    if outfile is None:
        return None

    # get optimized values from scene
    scene = scene_state.sparse_ga
    rgbimg = scene.imgs
    focals = scene.get_focals().cpu()
    cams2world = scene.get_im_poses().cpu()

    # 3D pointcloud from depthmap, poses and intrinsics
    if TSDF_thresh > 0:
        tsdf = TSDFPostProcess(scene, TSDF_thresh=TSDF_thresh)
        pts3d, _, confs = to_numpy(tsdf.get_dense_pts3d(clean_depth=clean_depth))
    else:
        pts3d, _, confs = to_numpy(scene.get_dense_pts3d(clean_depth=clean_depth))
    msk = to_numpy([c > min_conf_thr for c in confs])
    return _convert_scene_output_to_glb(
        outfile,
        rgbimg,
        pts3d,
        msk,
        focals,
        cams2world,
        as_pointcloud=as_pointcloud,
        transparent_cams=transparent_cams,
        cam_size=cam_size,
        silent=silent,
    )


def get_reconstructed_scene(
    outdir,
    gradio_delete_cache,
    model,
    device,
    silent,
    image_size,
    current_scene_state,
    filelist,
    optim_level,
    lr1,
    niter1,
    lr2,
    niter2,
    min_conf_thr,
    matching_conf_thr,
    as_pointcloud,
    mask_sky,
    clean_depth,
    transparent_cams,
    cam_size,
    scenegraph_type,
    winsize,
    win_cyclic,
    refid,
    TSDF_thresh,
    shared_intrinsics,
    **kw,
):
    """
    from a list of images, run mast3r inference, sparse global aligner.
    then run get_3D_model_from_scene
    """
    imgs = load_images(filelist, size=image_size, verbose=not silent)
    if len(imgs) == 1:
        imgs = [imgs[0], copy.deepcopy(imgs[0])]
        imgs[1]["idx"] = 1
        filelist = [filelist[0], filelist[0] + "_2"]

    scene_graph_params = [scenegraph_type]
    if scenegraph_type in ["swin", "logwin"]:
        scene_graph_params.append(str(winsize))
    elif scenegraph_type == "oneref":
        scene_graph_params.append(str(refid))
    if scenegraph_type in ["swin", "logwin"] and not win_cyclic:
        scene_graph_params.append("noncyclic")
    scene_graph = "-".join(scene_graph_params)
    pairs = make_pairs(imgs, scene_graph=scene_graph, prefilter=None, symmetrize=True)
    if optim_level == "coarse":
        niter2 = 0
    # Sparse GA (forward mast3r -> matching -> 3D optim -> 2D refinement -> triangulation)
    if (
        current_scene_state is not None
        and not current_scene_state.should_delete
        and current_scene_state.cache_dir is not None
    ):
        cache_dir = current_scene_state.cache_dir
    elif gradio_delete_cache:
        cache_dir = tempfile.mkdtemp(suffix="_cache", dir=outdir)
    else:
        cache_dir = os.path.join(outdir, "cache")
    os.makedirs(cache_dir, exist_ok=True)
    scene = sparse_global_alignment(
        filelist,
        pairs,
        cache_dir,
        model,
        lr1=lr1,
        niter1=niter1,
        lr2=lr2,
        niter2=niter2,
        device=device,
        opt_depth="depth" in optim_level,
        shared_intrinsics=shared_intrinsics,
        matching_conf_thr=matching_conf_thr,
        **kw,
    )
    if (
        current_scene_state is not None
        and not current_scene_state.should_delete
        and current_scene_state.outfile_name is not None
    ):
        outfile_name = current_scene_state.outfile_name
    else:
        outfile_name = tempfile.mktemp(suffix="_scene.glb", dir=outdir)

    scene_state = SparseGAState(scene, gradio_delete_cache, cache_dir, outfile_name)
    outfile = get_3D_model_from_scene(
        silent, scene_state, min_conf_thr, as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, TSDF_thresh
    )
    return scene_state, outfile


def set_scenegraph_options(inputfiles, win_cyclic, refid, scenegraph_type):
    num_files = len(inputfiles) if inputfiles is not None else 1
    show_win_controls = scenegraph_type in ["swin", "logwin"]
    show_winsize = scenegraph_type in ["swin", "logwin"]
    show_cyclic = scenegraph_type in ["swin", "logwin"]
    max_winsize, min_winsize = 1, 1
    if scenegraph_type == "swin":
        if win_cyclic:
            max_winsize = max(1, math.ceil((num_files - 1) / 2))
        else:
            max_winsize = num_files - 1
    elif scenegraph_type == "logwin":
        if win_cyclic:
            half_size = math.ceil((num_files - 1) / 2)
            max_winsize = max(1, math.ceil(math.log(half_size, 2)))
        else:
            max_winsize = max(1, math.ceil(math.log(num_files, 2)))
    winsize = gradio.Slider(
        label="Scene Graph: Window Size",
        value=max_winsize,
        minimum=min_winsize,
        maximum=max_winsize,
        step=1,
        visible=show_winsize,
    )
    win_cyclic = gradio.Checkbox(value=win_cyclic, label="Cyclic sequence", visible=show_cyclic)
    win_col = gradio.Column(visible=show_win_controls)
    refid = gradio.Slider(
        label="Scene Graph: Id", value=0, minimum=0, maximum=num_files - 1, step=1, visible=scenegraph_type == "oneref"
    )
    return win_col, winsize, win_cyclic, refid


def main_demo(
    tmpdirname,
    model,
    device,
    image_size,
    server_name,
    server_port,
    silent=False,
    share=False,
    gradio_delete_cache=False,
):
    if not silent:
        print("Outputing stuff in", tmpdirname)

    recon_fun = functools.partial(
        get_reconstructed_scene, tmpdirname, gradio_delete_cache, model, device, silent, image_size
    )
    model_from_scene_fun = functools.partial(get_3D_model_from_scene, silent)

    def get_context(delete_cache):
        css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
        title = "MASt3R Demo"
        if delete_cache:
            return gradio.Blocks(css=css, title=title, delete_cache=(delete_cache, delete_cache))
        else:
            return gradio.Blocks(css=css, title="MASt3R Demo")  # for compatibility with older versions

    with get_context(gradio_delete_cache) as demo:
        # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
        scene = gradio.State(None)
        gradio.HTML('<h2 style="text-align: center;">MASt3R Demo</h2>')
        with gradio.Column():
            inputfiles = gradio.File(file_count="multiple")
            with gradio.Row():
                with gradio.Column():
                    with gradio.Row():
                        lr1 = gradio.Slider(label="Coarse LR", value=0.07, minimum=0.01, maximum=0.2, step=0.01)
                        niter1 = gradio.Number(
                            value=500,
                            precision=0,
                            minimum=0,
                            maximum=10_000,
                            label="num_iterations",
                            info="For coarse alignment!",
                        )
                        lr2 = gradio.Slider(label="Fine LR", value=0.014, minimum=0.005, maximum=0.05, step=0.001)
                        niter2 = gradio.Number(
                            value=200,
                            precision=0,
                            minimum=0,
                            maximum=100_000,
                            label="num_iterations",
                            info="For refinement!",
                        )
                        optim_level = gradio.Dropdown(
                            ["coarse", "refine", "refine+depth"],
                            value="refine+depth",
                            label="OptLevel",
                            info="Optimization level",
                        )
                    with gradio.Row():
                        matching_conf_thr = gradio.Slider(
                            label="Matching Confidence Thr",
                            value=5.0,
                            minimum=0.0,
                            maximum=30.0,
                            step=0.1,
                            info="Before Fallback to Regr3D!",
                        )
                        shared_intrinsics = gradio.Checkbox(
                            value=False,
                            label="Shared intrinsics",
                            info="Only optimize one set of intrinsics for all views",
                        )
                        scenegraph_type = gradio.Dropdown(
                            [
                                ("complete: all possible image pairs", "complete"),
                                ("swin: sliding window", "swin"),
                                ("logwin: sliding window with long range", "logwin"),
                                ("oneref: match one image with all", "oneref"),
                            ],
                            value="complete",
                            label="Scenegraph",
                            info="Define how to make pairs",
                            interactive=True,
                        )
                        with gradio.Column(visible=False) as win_col:
                            winsize = gradio.Slider(
                                label="Scene Graph: Window Size", value=1, minimum=1, maximum=1, step=1
                            )
                            win_cyclic = gradio.Checkbox(value=False, label="Cyclic sequence")
                        refid = gradio.Slider(
                            label="Scene Graph: Id", value=0, minimum=0, maximum=0, step=1, visible=False
                        )
            run_btn = gradio.Button("Run")

            with gradio.Row():
                # adjust the confidence threshold
                min_conf_thr = gradio.Slider(label="min_conf_thr", value=1.5, minimum=0.0, maximum=10, step=0.1)
                # adjust the camera size in the output pointcloud
                cam_size = gradio.Slider(label="cam_size", value=0.2, minimum=0.001, maximum=1.0, step=0.001)
                TSDF_thresh = gradio.Slider(label="TSDF Threshold", value=0.0, minimum=0.0, maximum=1.0, step=0.01)
            with gradio.Row():
                as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud")
                # two post process implemented
                mask_sky = gradio.Checkbox(value=False, label="Mask sky")
                clean_depth = gradio.Checkbox(value=True, label="Clean-up depthmaps")
                transparent_cams = gradio.Checkbox(value=False, label="Transparent cameras")

            outmodel = gradio.Model3D()

            # events
            scenegraph_type.change(
                set_scenegraph_options,
                inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
                outputs=[win_col, winsize, win_cyclic, refid],
            )
            inputfiles.change(
                set_scenegraph_options,
                inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
                outputs=[win_col, winsize, win_cyclic, refid],
            )
            win_cyclic.change(
                set_scenegraph_options,
                inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
                outputs=[win_col, winsize, win_cyclic, refid],
            )
            run_btn.click(
                fn=recon_fun,
                inputs=[
                    scene,
                    inputfiles,
                    optim_level,
                    lr1,
                    niter1,
                    lr2,
                    niter2,
                    min_conf_thr,
                    matching_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    scenegraph_type,
                    winsize,
                    win_cyclic,
                    refid,
                    TSDF_thresh,
                    shared_intrinsics,
                ],
                outputs=[scene, outmodel],
            )
            min_conf_thr.release(
                fn=model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
            cam_size.change(
                fn=model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
            TSDF_thresh.change(
                fn=model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
            as_pointcloud.change(
                fn=model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
            mask_sky.change(
                fn=model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
            clean_depth.change(
                fn=model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
            transparent_cams.change(
                model_from_scene_fun,
                inputs=[
                    scene,
                    min_conf_thr,
                    as_pointcloud,
                    mask_sky,
                    clean_depth,
                    transparent_cams,
                    cam_size,
                    TSDF_thresh,
                ],
                outputs=outmodel,
            )
    demo.launch(share=share, server_name=server_name, server_port=server_port)
