# coding=utf-8

# Adapted from

# https://github.com/nerfstudio-project/gsplat/blob/main/examples/utils.py

# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.



import math

from math import exp

from typing import Dict, Optional, Tuple

import random

import numpy as np

import torch

from sklearn.neighbors import NearestNeighbors

from torch import Tensor

import torch.nn.functional as F

from torch.autograd import Variable





from datasets.colmap import Parser





def knn(x: Tensor, K: int = 4) -> Tensor:

    x_np = x.cpu().numpy()

    model = NearestNeighbors(n_neighbors=K, metric="euclidean").fit(x_np)

    distances, _ = model.kneighbors(x_np)

    return torch.from_numpy(distances).to(x)





def rgb_to_sh(rgb: Tensor) -> Tensor:

    C0 = 0.28209479177387814

    return (rgb - 0.5) / C0





def set_random_seed(seed: int):

    random.seed(seed)

    np.random.seed(seed)

    torch.manual_seed(seed)





class SSIM:

    window_cache = None



    @staticmethod

    def gaussian(window_size, sigma):

        gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])

        return gauss / gauss.sum()



    @staticmethod

    def create_window(window_size, channel):

        _1D_window = SSIM.gaussian(window_size, 1.5 / 2).unsqueeze(1)

        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)

        window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())

        return window



    @staticmethod

    def ssim(img1, img2, window_size=7, size_average=True):

        channel = img1.size(-3)

        if SSIM.window_cache is not None:

            window = SSIM.window_cache

        else:

            window = SSIM.create_window(window_size, channel)

            SSIM.window_cache = window

        window.to(img1.device)

        window = window.type_as(img1)

        return SSIM._ssim(img1, img2, window, window_size, channel, size_average)



    @staticmethod

    def _ssim(img1, img2, window, window_size, channel, size_average=True):

        padding = window_size // 2

        mu1 = F.conv2d(img1, window, padding=padding, groups=channel)

        mu2 = F.conv2d(img2, window, padding=padding, groups=channel)

        

        mu1_sq = mu1.pow(2)

        mu2_sq = mu2.pow(2)

        mu1_mu2 = mu1 * mu2

        sigma1_sq = F.conv2d(img1 * img1, window, padding=padding, groups=channel) - mu1_sq

        sigma2_sq = F.conv2d(img2 * img2, window, padding=padding, groups=channel) - mu2_sq

        sigma12 = F.conv2d(img1 * img2, window, padding=padding, groups=channel) - mu1_mu2



        C1 = 0.01 ** 2

        C2 = 0.03 ** 2

        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))

        if size_average:

            return ssim_map.mean()

        else:

            return ssim_map.mean(1).mean(1).mean(1)





def create_splats_with_optimizers(

    parser: Parser,

    init_type: str = "sfm",

    init_num_pts: int = 100_000,

    init_extent: float = 3.0,

    init_opacity: float = 0.1,

    init_scale: float = 1.0,

    scene_scale: float = 1.0,

    sh_degree: int = 3,

    sparse_grad: bool = False,

    batch_size: int = 1,

    feature_dim: Optional[int] = None,

    device: str = "npu",

    world_rank: int = 0,

    world_size: int = 1,

) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]:

    if init_type == "sfm":

        points = torch.from_numpy(parser.points).float()

        rgbs = torch.from_numpy(parser.points_rgb / 255.0).float()

    elif init_type == "random":

        points = init_extent * scene_scale * (torch.rand((init_num_pts, 3)) * 2 - 1)

        rgbs = torch.rand((init_num_pts, 3))

    else:

        raise ValueError("Please specify a correct init_type: sfm or random")



    # Initialize the GS size to be the average dist of the 3 nearest neighbors

    dist2_avg = (knn(points, 4)[:, 1:] ** 2).mean(dim=-1)  # [N,]

    dist_avg = torch.sqrt(dist2_avg)

    scales = torch.log(dist_avg * init_scale).unsqueeze(-1).repeat(1, 3)  # [N, 3]



    # Distribute the GSs to different ranks (also works for single rank)

    points = points[world_rank::world_size]

    rgbs = rgbs[world_rank::world_size]

    scales = scales[world_rank::world_size]



    N = points.shape[0]

    quats = torch.rand((N, 4))  # [N, 4]

    opacities = torch.logit(torch.full((N,), init_opacity))  # [N,]



    params = [

        # name, value, lr

        ("means", torch.nn.Parameter(points), 1.6e-4 * scene_scale),

        ("scales", torch.nn.Parameter(scales), 5e-3),

        ("quats", torch.nn.Parameter(quats), 1e-3),

        ("opacities", torch.nn.Parameter(opacities), 5e-2),

    ]



    if feature_dim is None:

        # color is SH coefficients.

        colors = torch.zeros((N, (sh_degree + 1) ** 2, 3))  # [N, K, 3]

        colors[:, 0, :] = rgb_to_sh(rgbs)

        params.append(("sh0", torch.nn.Parameter(colors[:, :1, :]), 2.5e-3))

        params.append(("shN", torch.nn.Parameter(colors[:, 1:, :]), 2.5e-3 / 20))

    else:

        # features will be used for appearance and view-dependent shading

        features = torch.rand(N, feature_dim)  # [N, feature_dim]

        params.append(("features", torch.nn.Parameter(features), 2.5e-3))

        colors = torch.logit(rgbs)  # [N, 3]

        params.append(("colors", torch.nn.Parameter(colors), 2.5e-3))



    splats = torch.nn.ParameterDict({n: v for n, v, _ in params}).to(device)

    # Scale learning rate based on batch size, reference:

    # https://www.cs.princeton.edu/~smalladi/blog/2024/01/22/SDEs-ScalingRules/

    # Note that this would not make the training exactly equivalent, see

    # https://arxiv.org/pdf/2402.18824v1

    BS = batch_size * world_size

    optimizer_class = None

    if sparse_grad:

        optimizer_class = torch.optim.SparseAdam

    else:

        optimizer_class = torch.optim.Adam

    optimizers = {

        name: optimizer_class(

            [{"params": splats[name], "lr": lr * math.sqrt(BS), "name": name}],

            eps=1e-15 / math.sqrt(BS),

            betas=(1 - BS * (1 - 0.9), 1 - BS * (1 - 0.999)),

        )

        for name, _, lr in params

    }

    return splats, optimizers





def assert_check(condition, message):

    if not condition:

        raise ValueError(message)





def validate_inputs(

    gs: Tuple,

    camera: Tuple,

    sh_degree: Optional[int],

    N: int,

    C: int

) -> None:

    means, quats, scales, opacities, colors = gs

    viewmats, Ks, camera_model, render_mode = camera

    assert_check(means.shape == (N, 3), f"Invalid shape for means: {means.shape}")

    assert_check(quats.shape == (N, 4), f"Invalid shape for quats: {quats.shape}")

    assert_check(scales.shape == (N, 3), f"Invalid shape for scales: {scales.shape}")

    assert_check(opacities.shape == (N,), f"Invalid shape for opacities: {opacities.shape}")

    assert_check(viewmats.shape == (C, 4, 4), f"Invalid shape for viewmats: {viewmats.shape}")

    assert_check(Ks.shape == (C, 3, 3), f"Invalid shape for Ks: {Ks.shape}")

    assert_check(camera_model in ["pinhole", "ortho", "fisheye"], f"Invalid camera_model: {camera_model}")

    assert_check(render_mode in ["RGB", "D", "ED", "RGB+D", "RGB+ED"], f"Invalid render_mode: {render_mode}")



    if sh_degree is None:

        # treat colors as post-activation values, should be in shape [ N, D] or [C, N, D]

        assert_check((colors.dim() == 2 and colors.shape[0] == N) or 

                    (colors.dim() == 3 and colors.shape[:2] == (C, N)),

                    f"Invalid shape for colors: {colors.shape}")

    else:

        # treat colors as SH coefficients, should be in shape [N, K, 3] or [C, N, K, 3]

        # Allowing for activating partial SH bands

        assert_check((colors.dim() == 3 and colors.shape[0] == N and colors.shape[2] == 3) or

                    (colors.dim() == 4 and colors.shape[:2] == (C, N) and colors.shape[3] == 3),

                    f"Invalid shape for colors: {colors.shape}")



        assert_check((sh_degree + 1) ** 2 <= colors.shape[-2], f"Invalid sh_degree for colors shape: {colors.shape}")