import einops
import gradio as gr
import matplotlib.cm as cm
import numpy as np
import plotly.graph_objects as go
import torch
import torch.nn.functional as F

from rendering import estimate_surface_normal

DESCRIPTION = """
<div class="head">
<div class="title">LiDAR Data Synthesis with Denoising Diffusion Probabilistic Models</div>
<div class="authors">
<a href="https://kazuto1011.github.io/" target="_blank" rel="noopener"> Kazuto Nakashima</a>
&nbsp;&nbsp;&nbsp;
<a href="https://robotics.ait.kyushu-u.ac.jp/kurazume/en/" target="_blank" rel="noopener"> Ryo Kurazume</a>
</div>
<div class="affiliations">Kyushu University</div>
<div class="conference">ICRA 2024</div>
<div class="materials">
<a href="https://kazuto1011.github.io/r2dm">Project</a> |
<a href="https://arxiv.org/abs/2309.09256">Paper</a> |
<a href="https://github.com/kazuto1011/r2dm">Code</a>
</div>
<br>
<div class="description">
This is a demo of our paper "LiDAR Data Synthesis with Denoising Diffusion Probabilistic Models" presented at ICRA 2024.<br>
We propose <strong>R2DM</strong>, a continuous-time diffusion model for LiDAR data generation based on the equirectangular range/reflectance image representation.<br>
</div>
<br>
</div>
"""

RUN_LOCALLY = """
To run this demo locally:

```bash
git clone https://huggingface.co/spaces/kazuto1011/r2dm
```
```bash
cd r2dm
```
```bash
pip install -r requirements.txt
```
```bash
pip install gradio
```
```bash
gradio app.py
```
"""

THEME = gr.themes.Default(font=gr.themes.GoogleFont("Titillium Web"))


if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

torch.set_grad_enabled(False)
torch.backends.cudnn.benchmark = True
device = torch.device(device)

model_dict = {
    "KITTI Raw (64x512)": torch.hub.load(
        "kazuto1011/r2dm",
        "pretrained_r2dm",
        config="r2dm-h-kittiraw-300k",
        device="cpu",
        show_info=False,
    ),
    "KITTI-360 (64x1024)": torch.hub.load(
        "kazuto1011/r2dm",
        "pretrained_r2dm",
        config="r2dm-h-kitti360-300k",
        device="cpu",
        show_info=False,
    ),
}


def colorize(tensor: torch.Tensor, cmap_fn=cm.turbo):
    colors = cmap_fn(np.linspace(0, 1, 256))[:, :3]
    colors = torch.from_numpy(colors).to(tensor)
    tensor = tensor.squeeze(1) if tensor.ndim == 4 else tensor
    ids = (tensor * 256).clamp(0, 255).long()
    tensor = F.embedding(ids, colors).permute(0, 3, 1, 2)
    tensor = tensor.mul(255).clamp(0, 255).byte()
    return tensor


def generate(num_steps: int, sampling_mode: str, dataset: str, progress=gr.Progress()):
    # model setup
    model, lidar_utils, _ = model_dict[dataset]
    model.to(device)
    lidar_utils.to(device)

    # sampling
    num_steps = int(num_steps)
    x = model.randn(1, *model.sampling_shape, device=model.device)
    steps = torch.linspace(1.0, 0.0, num_steps + 1, device=model.device)[None]
    for i in progress.tqdm(range(num_steps), desc="Generating LiDAR data"):
        step_t = steps[:, i]
        step_s = steps[:, i + 1]
        x = model.p_step(x, step_t, step_s, mode=sampling_mode.lower())

    # rendering point cloud
    x = lidar_utils.denormalize(x.clamp(-1, 1))
    depth = lidar_utils.revert_depth(x[:, [0]])
    rflct = x[:, [1]]
    point = lidar_utils.to_xyz(depth)
    color = (-estimate_surface_normal(point) + 1) / 2
    point = einops.rearrange(point, "1 c h w -> (h w) c").cpu().numpy()
    color = einops.rearrange(color, "1 c h w -> (h w) c").cpu().numpy()
    fig = go.Figure(
        data=[
            go.Scatter3d(
                x=-point[..., 0],
                y=-point[..., 1],
                z=point[..., 2],
                mode="markers",
                marker=dict(size=1, color=color),
            )
        ],
        layout=dict(
            scene=dict(
                xaxis=dict(showticklabels=False, visible=False),
                yaxis=dict(showticklabels=False, visible=False),
                zaxis=dict(showticklabels=False, visible=False),
                aspectmode="data",
            ),
            margin=dict(l=0, r=0, b=0, t=0),
            paper_bgcolor="rgba(0,0,0,0)",
            plot_bgcolor="rgba(0,0,0,0)",
        ),
    )
    depth = depth / lidar_utils.max_depth
    depth = colorize(depth, cm.turbo)[0].permute(1, 2, 0).cpu().numpy()
    rflct = colorize(rflct, cm.turbo)[0].permute(1, 2, 0).cpu().numpy()

    model.cpu()
    lidar_utils.cpu()
    return depth, rflct, fig


with gr.Blocks(css="./style.css", theme=THEME) as demo:
    gr.HTML(DESCRIPTION)

    with gr.Row(variant="panel"):
        with gr.Column():
            gr.Textbox(device, label="Running device")
            dataset = gr.Dropdown(
                choices=list(model_dict.keys()),
                value=list(model_dict.keys())[0],
                label="Dataset",
            )
            sampling_mode = gr.Dropdown(
                choices=["DDPM", "DDIM"],
                value="DDPM",
                label="Sampler",
            )
            num_steps = gr.Dropdown(
                choices=[2**i for i in range(5, 11)],
                value=32,
                label="Number of sampling steps (>256 is recommended)",
            )
            btn = gr.Button(value="Generate")

        with gr.Column():
            range_view = gr.Image(type="numpy", label="Range image")
            rflct_view = gr.Image(type="numpy", label="Reflectance image")
            point_view = gr.Plot(label="Point cloud")

    with gr.Row(variant="panel"):
        gr.Markdown(RUN_LOCALLY)

    btn.click(
        generate,
        inputs=[num_steps, sampling_mode, dataset],
        outputs=[range_view, rflct_view, point_view],
        concurrency_limit=1,
    )


demo.queue()
demo.launch()