# SPDX-FileCopyrightText: Copyright 2025 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

import argparse
import math
import os
import time
from typing import Tuple

import imageio
import nerfview
import numpy as np
import torch
import torch.nn.functional as F
import tqdm
import viser

from gsplat._helper import load_test_data
from gsplat.distributed import cli
from gsplat.rendering import rasterization


def main(local_rank: int, world_rank, world_size: int, args):
    torch.manual_seed(42)
    device = torch.device("cuda", local_rank)

    if args.ckpt is None:
        (
            means,
            quats,
            scales,
            opacities,
            colors,
            viewmats,
            Ks,
            width,
            height,
        ) = load_test_data(device=device, scene_grid=args.scene_grid)

        assert world_size <= 2
        means = means[world_rank::world_size].contiguous()
        means.requires_grad = True
        quats = quats[world_rank::world_size].contiguous()
        quats.requires_grad = True
        scales = scales[world_rank::world_size].contiguous()
        scales.requires_grad = True
        opacities = opacities[world_rank::world_size].contiguous()
        opacities.requires_grad = True
        colors = colors[world_rank::world_size].contiguous()
        colors.requires_grad = True

        viewmats = viewmats[world_rank::world_size][:1].contiguous()
        Ks = Ks[world_rank::world_size][:1].contiguous()

        sh_degree = None
        C = len(viewmats)
        N = len(means)
        print("rank", world_rank, "Number of Gaussians:", N, "Number of Cameras:", C)

        # batched render
        for _ in tqdm.trange(1):
            render_colors, render_alphas, meta = rasterization(
                means,  # [N, 3]
                quats,  # [N, 4]
                scales,  # [N, 3]
                opacities,  # [N]
                colors,  # [N, S, 3]
                viewmats,  # [C, 4, 4]
                Ks,  # [C, 3, 3]
                width,
                height,
                render_mode="RGB+D",
                packed=False,
                distributed=world_size > 1,
            )
        C = render_colors.shape[0]
        assert render_colors.shape == (C, height, width, 4)
        assert render_alphas.shape == (C, height, width, 1)
        render_colors.sum().backward()

        render_rgbs = render_colors[..., 0:3]
        render_depths = render_colors[..., 3:4]
        render_depths = render_depths / render_depths.max()

        # dump batch images
        os.makedirs(args.output_dir, exist_ok=True)
        canvas = (
            torch.cat(
                [
                    render_rgbs.reshape(C * height, width, 3),
                    render_depths.reshape(C * height, width, 1).expand(-1, -1, 3),
                    render_alphas.reshape(C * height, width, 1).expand(-1, -1, 3),
                ],
                dim=1,
            )
            .detach()
            .cpu()
            .numpy()
        )
        imageio.imsave(
            f"{args.output_dir}/render_rank{world_rank}.png",
            (canvas * 255).astype(np.uint8),
        )
    else:
        means, quats, scales, opacities, sh0, shN = [], [], [], [], [], []
        for ckpt_path in args.ckpt:
            ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)[
                "splats"
            ]
            means.append(ckpt["means"])
            quats.append(F.normalize(ckpt["quats"], p=2, dim=-1))
            scales.append(torch.exp(ckpt["scales"]))
            opacities.append(torch.sigmoid(ckpt["opacities"]))
            sh0.append(ckpt["sh0"])
            shN.append(ckpt["shN"])
        means = torch.cat(means, dim=0)
        quats = torch.cat(quats, dim=0)
        scales = torch.cat(scales, dim=0)
        opacities = torch.cat(opacities, dim=0)
        sh0 = torch.cat(sh0, dim=0)
        shN = torch.cat(shN, dim=0)
        colors = torch.cat([sh0, shN], dim=-2)
        sh_degree = int(math.sqrt(colors.shape[-2]) - 1)
        print("Number of Gaussians:", len(means))

    # register and open viewer
    @torch.no_grad()
    def viewer_render_fn(camera_state: nerfview.CameraState, img_wh: Tuple[int, int]):
        width, height = img_wh
        c2w = camera_state.c2w
        K = camera_state.get_K(img_wh)
        c2w = torch.from_numpy(c2w).float().to(device)
        K = torch.from_numpy(K).float().to(device)
        viewmat = c2w.inverse()

        if args.backend == "gsplat":
            rasterization_fn = rasterization
        elif args.backend == "inria":
            from gsplat import rasterization_inria_wrapper

            rasterization_fn = rasterization_inria_wrapper
        else:
            raise ValueError

        if gui_ckeckbox_3dgut.value:
            radial_coeffs = torch.tensor(
                [
                    [
                        gui_slider_radial_coeffs1.value,  # k1
                        gui_slider_radial_coeffs2.value,  # k2
                        0.0,
                        0.0,
                        0.0,
                        0.0,
                    ]
                ],
                device=device,
            )
            tangential_coeffs = torch.tensor(
                [
                    [
                        gui_slider_tangential_coeffs1.value,  # p1
                        gui_slider_tangential_coeffs2.value,  # p2
                    ]
                ],
                device=device,
            )
            thin_prism_coeffs = torch.tensor(
                [
                    [
                        gui_slider_thin_prism_coeffs1.value,  # s1
                        gui_slider_thin_prism_coeffs2.value,  # s2
                        0.0,  # s3
                        0.0,  # s4
                    ]
                ],
                device=device,
            )
        else:
            radial_coeffs = None
            tangential_coeffs = None
            thin_prism_coeffs = None

        render_colors, render_alphas, meta = rasterization_fn(
            means,  # [N, 3]
            quats,  # [N, 4]
            scales,  # [N, 3]
            opacities,  # [N]
            colors,  # [N, S, 3]
            viewmat[None],  # [1, 4, 4]
            K[None],  # [1, 3, 3]
            width,
            height,
            sh_degree=sh_degree,
            packed=False,
            render_mode="RGB",
            # this is to speedup large-scale rendering by skipping far-away Gaussians.
            radius_clip=3,
            camera_model=gui_dropdown_camera_type.value,  # "pinhole" or "fisheye"
            with_ut=gui_ckeckbox_3dgut.value,
            with_eval3d=gui_ckeckbox_3dgut.value,
            radial_coeffs=radial_coeffs,
            tangential_coeffs=tangential_coeffs,
            thin_prism_coeffs=thin_prism_coeffs,
        )
        render_rgbs = render_colors[0, ..., 0:3].cpu().numpy()
        return render_rgbs

    server = viser.ViserServer(port=args.port, verbose=False)

    # add controls
    gui_ckeckbox_3dgut = server.gui.add_checkbox(
        "3DGUT (UT + Eval3D)", initial_value=False
    )
    gui_dropdown_camera_type = server.gui.add_dropdown(
        "Camera Type", ("pinhole", "fisheye"), initial_value="pinhole"
    )
    gui_slider_radial_coeffs1 = server.gui.add_slider(
        "radial coeffs: k1", min=-1.0, max=1.0, step=0.01, initial_value=0
    )
    gui_slider_radial_coeffs2 = server.gui.add_slider(
        "radial coeffs: k2", min=-1.0, max=1.0, step=0.01, initial_value=0
    )
    gui_slider_tangential_coeffs1 = server.gui.add_slider(
        "tangential coeffs: p1", min=-1.0, max=1.0, step=0.01, initial_value=0
    )
    gui_slider_tangential_coeffs2 = server.gui.add_slider(
        "tangential coeffs: p2", min=-1.0, max=1.0, step=0.01, initial_value=0
    )
    gui_slider_thin_prism_coeffs1 = server.gui.add_slider(
        "thin prism coeffs: s1", min=-1.0, max=1.0, step=0.01, initial_value=0
    )
    gui_slider_thin_prism_coeffs2 = server.gui.add_slider(
        "thin prism coeffs: s2", min=-1.0, max=1.0, step=0.01, initial_value=0
    )

    @gui_ckeckbox_3dgut.on_update
    @gui_dropdown_camera_type.on_update
    @gui_slider_radial_coeffs1.on_update
    @gui_slider_radial_coeffs2.on_update
    @gui_slider_tangential_coeffs1.on_update
    @gui_slider_tangential_coeffs2.on_update
    @gui_slider_thin_prism_coeffs1.on_update
    @gui_slider_thin_prism_coeffs2.on_update
    def _(event: viser.GuiEvent) -> None:
        viewer.rerender(None)

    viewer = nerfview.Viewer(
        server=server,
        render_fn=viewer_render_fn,
        mode="rendering",
    )
    print("Viewer running... Ctrl+C to exit.")
    time.sleep(100000)


if __name__ == "__main__":
    """
    # Use single GPU to view the scene
    CUDA_VISIBLE_DEVICES=0 python simple_viewer.py \
        --ckpt results/garden/ckpts/ckpt_3499_rank0.pt results/garden/ckpts/ckpt_3499_rank1.pt \
        --port 8081
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output_dir", type=str, default="results/", help="where to dump outputs"
    )
    parser.add_argument(
        "--scene_grid", type=int, default=1, help="repeat the scene into a grid of NxN"
    )
    parser.add_argument(
        "--ckpt", type=str, nargs="+", default=None, help="path to the .pt file"
    )
    parser.add_argument(
        "--port", type=int, default=8080, help="port for the viewer server"
    )
    parser.add_argument("--backend", type=str, default="gsplat", help="gsplat, inria")
    args = parser.parse_args()
    assert args.scene_grid % 2 == 1, "scene_grid must be odd"

    cli(main, args, verbose=True)
