content: from dataclasses import dataclass, field
  from collections import defaultdict
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  import math
  
  from ..utils.typing import *
  from ..utils.base import BaseModule
  from ..utils.ops import trunc_exp
  from ..utils.ops import scale_tensor
  from .networks import MLP
  from einops import rearrange, reduce
  
  from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply
  
  inverse_sigmoid = lambda x: np.log(x / (1 - x))
  
  def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):
      Rt = np.zeros((4, 4))
      Rt[:3, :3] = R.transpose()
      Rt[:3, 3] = t
      Rt[3, 3] = 1.0
  
      C2W = np.linalg.inv(Rt)
      cam_center = C2W[:3, 3]
      cam_center = (cam_center + translate) * scale
      C2W[:3, 3] = cam_center
      Rt = np.linalg.inv(C2W)
      return np.float32(Rt)
  
  def getProjectionMatrix(znear, zfar, fovX, fovY):
      tanHalfFovY = math.tan((fovY / 2))
      tanHalfFovX = math.tan((fovX / 2))
  
      top = tanHalfFovY * znear
      bottom = -top
      right = tanHalfFovX * znear
      left = -right
  
      P = torch.zeros(4, 4)
  
      z_sign = 1.0
  
      P[0, 0] = 2.0 * znear / (right - left)
      P[1, 1] = 2.0 * znear / (top - bottom)
      P[0, 2] = (right + left) / (right - left)
      P[1, 2] = (top + bottom) / (top - bottom)
      P[3, 2] = z_sign
      P[2, 2] = z_sign * zfar / (zfar - znear)
      P[2, 3] = -(zfar * znear) / (zfar - znear)
      return P
  
  def intrinsic_to_fov(intrinsic, w, h):
      fx, fy = intrinsic[0, 0], intrinsic[1, 1]
      fov_x = 2 * torch.arctan2(w, 2 * fx)
      fov_y = 2 * torch.arctan2(h, 2 * fy)
      return fov_x, fov_y
  
  
  class Camera:
      def __init__(self, w2c, intrinsic, FoVx, FoVy, height, width, trans=np.array([0.0, 0.0, 0.0]), scale=1.0) -> None:
          self.FoVx = FoVx
          self.FoVy = FoVy
          self.height = height
          self.width = width
          self.world_view_transform = w2c.transpose(0, 1)
  
          self.zfar = 100.0
          self.znear = 0.01
  
          self.trans = trans
          self.scale = scale
  
          self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).to(w2c.device)
          self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)
          self.camera_center = self.world_view_transform.inverse()[3, :3]
  
      @staticmethod
      def from_c2w(c2w, intrinsic, height, width):
          w2c = torch.inverse(c2w)
          FoVx, FoVy = intrinsic_to_fov(intrinsic, w=torch.tensor(width, device=w2c.device), h=torch.tensor(height, device=w2c.device))
          return Camera(w2c=w2c, intrinsic=intrinsic, FoVx=FoVx, FoVy=FoVy, height=height, width=width)
  
  class GaussianModel(NamedTuple):
      xyz: Tensor
      opacity: Tensor
      rotation: Tensor
      scaling: Tensor
      shs: Tensor
  
      def to_ply(self):
          
          xyz = self.xyz.detach().cpu().numpy()
          normals = np.zeros_like(xyz)
          features_dc = self.shs[:, :1]
          features_rest = self.shs[:, 1:]
          f_dc = features_dc.detach().flatten(start_dim=1).contiguous().cpu().numpy()
          f_rest = features_rest.detach().flatten(start_dim=1).contiguous().cpu().numpy()
          opacities = inverse_sigmoid(torch.clamp(self.opacity, 1e-3, 1 - 1e-3).detach().cpu().numpy())
          scale = np.log(self.scaling.detach().cpu().numpy())
          rotation = self.rotation.detach().cpu().numpy()
  
          return write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, construct_list_of_gs_attributes(features_dc, features_rest, self.scaling, self.rotation))
  
  class GSLayer(BaseModule):
      @dataclass
      class Config(BaseModule.Config):
          in_channels: int = 128
          feature_channels: dict = field(default_factory=dict)
          xyz_offset: bool = True
          restrict_offset: bool = False
          use_rgb: bool = False
          clip_scaling: Optional[float] = None
          init_scaling: float = -5.0
          init_density: float = 0.1
  
      cfg: Config
  
      def configure(self, *args, **kwargs) -> None:
          self.out_layers = nn.ModuleList()
          for key, out_ch in self.cfg.feature_channels.items():
              if key == "shs" and self.cfg.use_rgb:
                  out_ch = 3
              layer = nn.Linear(self.cfg.in_channels, out_ch)
  
              # initialize
              if not (key == "shs" and self.cfg.use_rgb):
                  nn.init.constant_(layer.weight, 0)
                  nn.init.constant_(layer.bias, 0)
              if key == "scaling":
                  nn.init.constant_(layer.bias, self.cfg.init_scaling)
              elif key == "rotation":
                  nn.init.constant_(layer.bias, 0)
                  nn.init.constant_(layer.bias[0], 1.0)
              elif key == "opacity":
                  nn.init.constant_(layer.bias, inverse_sigmoid(self.cfg.init_density))
  
              self.out_layers.append(layer)
  
      def forward(self, x, pts):
          ret = {}
          for k, layer in zip(self.cfg.feature_channels.keys(), self.out_layers):
              v = layer(x)
              if k == "rotation":
                  v = torch.nn.functional.normalize(v)
              elif k == "scaling":
                  v = trunc_exp(v)
                  if self.cfg.clip_scaling is not None:
                      v = torch.clamp(v, min=0, max=self.cfg.clip_scaling)
              elif k == "opacity":
                  v = torch.sigmoid(v)
              elif k == "shs":
                  if self.cfg.use_rgb:
                      v = torch.sigmoid(v)
                  v = torch.reshape(v, (v.shape[0], -1, 3))
              elif k == "xyz":
                  if self.cfg.restrict_offset:
                      max_step = 1.2 / 32
                      v = (torch.sigmoid(v) - 0.5) * max_step
                  v = v + pts if self.cfg.xyz_offset else pts
              ret[k] = v
  
          return GaussianModel(**ret)
  
  class GS3DRenderer(BaseModule):
      @dataclass
      class Config(BaseModule.Config):
          mlp_network_config: Optional[dict] = None
          gs_out: dict = field(default_factory=dict)
          sh_degree: int = 3
          scaling_modifier: float = 1.0
          random_background: bool = False
          radius: float = 1.0
          feature_reduction: str = "concat"
          projection_feature_dim: int = 773
          background_color: Tuple[float, float, float] = field(
              default_factory=lambda: (1.0, 1.0, 1.0)
          )
  
      cfg: Config
  
      def configure(self, *args, **kwargs) -> None:
          if self.cfg.feature_reduction == "mean":
              mlp_in = 80
          elif self.cfg.feature_reduction == "concat":
              mlp_in = 80 * 3
          else:
              raise NotImplementedError
          mlp_in = mlp_in + self.cfg.projection_feature_dim
          if self.cfg.mlp_network_config is not None:
              self.mlp_net = MLP(mlp_in, self.cfg.gs_out.in_channels, **self.cfg.mlp_network_config)
          else:
              self.cfg.gs_out.in_channels = mlp_in
          self.gs_net = GSLayer(self.cfg.gs_out)
  
      def forward_gs(self, x, p):
          if self.cfg.mlp_network_config is not None:
              x = self.mlp_net(x)
          return self.gs_net(x, p)
  
      def forward_single_view(self,
          gs: GaussianModel,
          viewpoint_camera: Camera,
          background_color: Optional[Float[Tensor, "3"]],
          ret_mask: bool = True,
          ):
          from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
          # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
          screenspace_points = torch.zeros_like(gs.xyz, dtype=gs.xyz.dtype, requires_grad=True, device=self.device) + 0
          try:
              screenspace_points.retain_grad()
          except:
              pass
          
          bg_color = background_color
          # Set up rasterization configuration
          tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
          tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
  
          raster_settings = GaussianRasterizationSettings(
              image_height=int(viewpoint_camera.height),
              image_width=int(viewpoint_camera.width),
              tanfovx=tanfovx,
              tanfovy=tanfovy,
              bg=bg_color,
              scale_modifier=self.cfg.scaling_modifier,
              viewmatrix=viewpoint_camera.world_view_transform,
              projmatrix=viewpoint_camera.full_proj_transform.float(),
              sh_degree=self.cfg.sh_degree,
              campos=viewpoint_camera.camera_center,
              prefiltered=False,
              debug=False
          )
  
          rasterizer = GaussianRasterizer(raster_settings=raster_settings)
  
          means3D = gs.xyz
          means2D = screenspace_points
          opacity = gs.opacity
  
          # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
          # scaling / rotation by the rasterizer.
          scales = None
          rotations = None
          cov3D_precomp = None
          scales = gs.scaling
          rotations = gs.rotation
  
          # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
          # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
          shs = None
          colors_precomp = None
          if self.gs_net.cfg.use_rgb:
              colors_precomp = gs.shs.squeeze(1)
          else:
              shs = gs.shs
  
          # Rasterize visible Gaussians to image, obtain their radii (on screen). 
          with torch.autocast(device_type=self.device.type, dtype=torch.float32):
              rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
                  means3D = means3D,
                  means2D = means2D,
                  shs = shs,
                  colors_precomp = colors_precomp,
                  opacities = opacity,
                  scales = scales,
                  rotations = rotations,
                  cov3D_precomp = cov3D_precomp)
          
          ret = {
              "comp_rgb": rendered_image.permute(1, 2, 0),
              "comp_rgb_bg": bg_color
          }
          
          if ret_mask:
              mask_bg_color = torch.zeros(3, dtype=torch.float32, device=self.device)
              raster_settings = GaussianRasterizationSettings(
                  image_height=int(viewpoint_camera.height),
                  image_width=int(viewpoint_camera.width),
                  tanfovx=tanfovx,
                  tanfovy=tanfovy,
                  bg=mask_bg_color,
                  scale_modifier=self.cfg.scaling_modifier,
                  viewmatrix=viewpoint_camera.world_view_transform,
                  projmatrix=viewpoint_camera.full_proj_transform.float(),
                  sh_degree=0,
                  campos=viewpoint_camera.camera_center,
                  prefiltered=False,
                  debug=False
              )
              rasterizer = GaussianRasterizer(raster_settings=raster_settings)
              
              with torch.autocast(device_type=self.device.type, dtype=torch.float32):
                  rendered_mask, radii, rendered_depth, rendered_alpha = rasterizer(
                      means3D = means3D,
                      means2D = means2D,
                      # shs = ,
                      colors_precomp = torch.ones_like(means3D),
                      opacities = opacity,
                      scales = scales,
                      rotations = rotations,
                      cov3D_precomp = cov3D_precomp)
                  ret["comp_mask"] = rendered_mask.permute(1, 2, 0)
  
          return ret
      
      def query_triplane(
          self,
          positions: Float[Tensor, "*B N 3"],
          triplanes: Float[Tensor, "*B 3 Cp Hp Wp"],
      ) -> Dict[str, Tensor]:
          batched = positions.ndim == 3
          if not batched:
              # no batch dimension
              triplanes = triplanes[None, ...]
              positions = positions[None, ...]
  
          positions = scale_tensor(positions, (-self.cfg.radius, self.cfg.radius), (-1, 1))
          indices2D: Float[Tensor, "B 3 N 2"] = torch.stack(
                  (positions[..., [0, 1]], positions[..., [0, 2]], positions[..., [1, 2]]),
                  dim=-3,
              )
          out: Float[Tensor, "B3 Cp 1 N"] = F.grid_sample(
              rearrange(triplanes, "B Np Cp Hp Wp -> (B Np) Cp Hp Wp", Np=3),
              rearrange(indices2D, "B Np N Nd -> (B Np) () N Nd", Np=3),
              align_corners=False,
              mode="bilinear",
          )
          if self.cfg.feature_reduction == "concat":
              out = rearrange(out, "(B Np) Cp () N -> B N (Np Cp)", Np=3)
          elif self.cfg.feature_reduction == "mean":
              out = reduce(out, "(B Np) Cp () N -> B N Cp", Np=3, reduction="mean")
          else:
              raise NotImplementedError
          
          if not batched:
              out = out.squeeze(0)
  
          return out
  
      def forward_single_batch(
          self,
          gs_hidden_features: Float[Tensor, "Np Cp"],
          query_points: Float[Tensor, "Np 3"],
          c2ws: Float[Tensor, "Nv 4 4"],
          intrinsics: Float[Tensor, "Nv 4 4"],
          height: int,
          width: int,
          background_color: Optional[Float[Tensor, "3"]],
      ):
          gs: GaussianModel = self.forward_gs(gs_hidden_features, query_points)
          out_list = []
         
          for c2w, intrinsic in zip(c2ws, intrinsics):
              out_list.append(self.forward_single_view(
                                  gs, 
                                  Camera.from_c2w(c2w, intrinsic, height, width),
                                  background_color
                              ))
          
          out = defaultdict(list)
          for out_ in out_list:
              for k, v in out_.items():
                  out[k].append(v)
          out = {k: torch.stack(v, dim=0) for k, v in out.items()}
          out["3dgs"] = gs
  
          return out
  
      def forward(self, 
          gs_hidden_features: Float[Tensor, "B Np Cp"],
          query_points: Float[Tensor, "B Np 3"],
          c2w: Float[Tensor, "B Nv 4 4"],
          intrinsic: Float[Tensor, "B Nv 4 4"],
          height,
          width,
          additional_features: Optional[Float[Tensor, "B C H W"]] = None,
          background_color: Optional[Float[Tensor, "B 3"]] = None,
          **kwargs):
          batch_size = gs_hidden_features.shape[0]
          out_list = []
          gs_hidden_features = self.query_triplane(query_points, gs_hidden_features)
          if additional_features is not None:
              gs_hidden_features = torch.cat([gs_hidden_features, additional_features], dim=-1)
  
          for b in range(batch_size):
              out_list.append(self.forward_single_batch(
                  gs_hidden_features[b],
                  query_points[b],
                  c2w[b],
                  intrinsic[b],
                  height, width,
                  background_color[b] if background_color is not None else None))
  
          out = defaultdict(list)
          for out_ in out_list:
              for k, v in out_.items():
                  out[k].append(v)
          for k, v in out.items():
              if isinstance(v[0], torch.Tensor):
                  out[k] = torch.stack(v, dim=0)
              else:
                  out[k] = v
          return out
          