content: # -*- coding: utf-8 -*-
  
  # Copyright (c) 2012-2015, P. M. Neila
  # All rights reserved.
  
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions are met:
  
  # * Redistributions of source code must retain the above copyright notice, this
  #   list of conditions and the following disclaimer.
  
  # * Redistributions in binary form must reproduce the above copyright notice,
  #   this list of conditions and the following disclaimer in the documentation
  #   and/or other materials provided with the distribution.
  
  # * Neither the name of the copyright holder nor the names of its
  #   contributors may be used to endorse or promote products derived from
  #   this software without specific prior written permission.
  
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  """
  Utilities for smoothing the occ/sdf grids.
  """
  
  import logging
  from typing import Tuple
  
  import numpy as np
  import torch
  import torch.nn.functional as F
  from scipy import ndimage as ndi
  from scipy import sparse
  
  __all__ = [
      "smooth",
      "smooth_constrained",
      "smooth_gaussian",
      "signed_distance_function",
      "smooth_gpu",
      "smooth_constrained_gpu",
      "smooth_gaussian_gpu",
      "signed_distance_function_gpu",
  ]
  
  
  def _build_variable_indices(band: np.ndarray) -> np.ndarray:
      num_variables = np.count_nonzero(band)
      variable_indices = np.full(band.shape, -1, dtype=np.int_)
      variable_indices[band] = np.arange(num_variables)
      return variable_indices
  
  
  def _buildq3d(variable_indices: np.ndarray):
      """
      Builds the filterq matrix for the given variables.
      """
  
      num_variables = variable_indices.max() + 1
      filterq = sparse.lil_matrix((3 * num_variables, num_variables))
  
      # Pad variable_indices to simplify out-of-bounds accesses
      variable_indices = np.pad(
          variable_indices, [(0, 1), (0, 1), (0, 1)], mode="constant", constant_values=-1
      )
  
      coords = np.nonzero(variable_indices >= 0)
      for count, (i, j, k) in enumerate(zip(*coords)):
  
          assert variable_indices[i, j, k] == count
  
          filterq[3 * count, count] = -2
          neighbor = variable_indices[i - 1, j, k]
          if neighbor >= 0:
              filterq[3 * count, neighbor] = 1
          else:
              filterq[3 * count, count] += 1
  
          neighbor = variable_indices[i + 1, j, k]
          if neighbor >= 0:
              filterq[3 * count, neighbor] = 1
          else:
              filterq[3 * count, count] += 1
  
          filterq[3 * count + 1, count] = -2
          neighbor = variable_indices[i, j - 1, k]
          if neighbor >= 0:
              filterq[3 * count + 1, neighbor] = 1
          else:
              filterq[3 * count + 1, count] += 1
  
          neighbor = variable_indices[i, j + 1, k]
          if neighbor >= 0:
              filterq[3 * count + 1, neighbor] = 1
          else:
              filterq[3 * count + 1, count] += 1
  
          filterq[3 * count + 2, count] = -2
          neighbor = variable_indices[i, j, k - 1]
          if neighbor >= 0:
              filterq[3 * count + 2, neighbor] = 1
          else:
              filterq[3 * count + 2, count] += 1
  
          neighbor = variable_indices[i, j, k + 1]
          if neighbor >= 0:
              filterq[3 * count + 2, neighbor] = 1
          else:
              filterq[3 * count + 2, count] += 1
  
      filterq = filterq.tocsr()
      return filterq.T.dot(filterq)
  
  
  def _buildq3d_gpu(variable_indices: torch.Tensor, chunk_size=10000):
      """
      Builds the filterq matrix for the given variables on GPU, using chunking to reduce memory usage.
      """
      device = variable_indices.device
      num_variables = variable_indices.max().item() + 1
  
      # Pad variable_indices to simplify out-of-bounds accesses
      variable_indices = torch.nn.functional.pad(
          variable_indices, (0, 1, 0, 1, 0, 1), mode="constant", value=-1
      )
  
      coords = torch.nonzero(variable_indices >= 0)
      i, j, k = coords[:, 0], coords[:, 1], coords[:, 2]
  
      # Function to process a chunk of data
      def process_chunk(start, end):
          row_indices = []
          col_indices = []
          values = []
  
          for axis in range(3):
              row_indices.append(3 * torch.arange(start, end, device=device) + axis)
              col_indices.append(
                  variable_indices[i[start:end], j[start:end], k[start:end]]
              )
              values.append(torch.full((end - start,), -2, device=device))
  
              for offset in [-1, 1]:
                  if axis == 0:
                      neighbor = variable_indices[
                          i[start:end] + offset, j[start:end], k[start:end]
                      ]
                  elif axis == 1:
                      neighbor = variable_indices[
                          i[start:end], j[start:end] + offset, k[start:end]
                      ]
                  else:
                      neighbor = variable_indices[
                          i[start:end], j[start:end], k[start:end] + offset
                      ]
  
                  mask = neighbor >= 0
                  row_indices.append(
                      3 * torch.arange(start, end, device=device)[mask] + axis
                  )
                  col_indices.append(neighbor[mask])
                  values.append(torch.ones(mask.sum(), device=device))
  
                  # Add 1 to the diagonal for out-of-bounds neighbors
                  row_indices.append(
                      3 * torch.arange(start, end, device=device)[~mask] + axis
                  )
                  col_indices.append(
                      variable_indices[i[start:end], j[start:end], k[start:end]][~mask]
                  )
                  values.append(torch.ones((~mask).sum(), device=device))
  
          return torch.cat(row_indices), torch.cat(col_indices), torch.cat(values)
  
      # Process data in chunks
      all_row_indices = []
      all_col_indices = []
      all_values = []
  
      for start in range(0, coords.shape[0], chunk_size):
          end = min(start + chunk_size, coords.shape[0])
          row_indices, col_indices, values = process_chunk(start, end)
          all_row_indices.append(row_indices)
          all_col_indices.append(col_indices)
          all_values.append(values)
  
      # Concatenate all chunks
      row_indices = torch.cat(all_row_indices)
      col_indices = torch.cat(all_col_indices)
      values = torch.cat(all_values)
  
      # Create sparse tensor
      indices = torch.stack([row_indices, col_indices])
      filterq = torch.sparse_coo_tensor(
          indices, values, (3 * num_variables, num_variables)
      )
  
      # Compute filterq.T @ filterq
      return torch.sparse.mm(filterq.t(), filterq)
  
  
  # Usage example:
  # variable_indices = torch.tensor(...).cuda()  # Your input tensor on GPU
  # result = _buildq3d_gpu(variable_indices)
  
  
  def _buildq2d(variable_indices: np.ndarray):
      """
      Builds the filterq matrix for the given variables.
  
      Version for 2 dimensions.
      """
  
      num_variables = variable_indices.max() + 1
      filterq = sparse.lil_matrix((3 * num_variables, num_variables))
  
      # Pad variable_indices to simplify out-of-bounds accesses
      variable_indices = np.pad(
          variable_indices, [(0, 1), (0, 1)], mode="constant", constant_values=-1
      )
  
      coords = np.nonzero(variable_indices >= 0)
      for count, (i, j) in enumerate(zip(*coords)):
          assert variable_indices[i, j] == count
  
          filterq[2 * count, count] = -2
          neighbor = variable_indices[i - 1, j]
          if neighbor >= 0:
              filterq[2 * count, neighbor] = 1
          else:
              filterq[2 * count, count] += 1
  
          neighbor = variable_indices[i + 1, j]
          if neighbor >= 0:
              filterq[2 * count, neighbor] = 1
          else:
              filterq[2 * count, count] += 1
  
          filterq[2 * count + 1, count] = -2
          neighbor = variable_indices[i, j - 1]
          if neighbor >= 0:
              filterq[2 * count + 1, neighbor] = 1
          else:
              filterq[2 * count + 1, count] += 1
  
          neighbor = variable_indices[i, j + 1]
          if neighbor >= 0:
              filterq[2 * count + 1, neighbor] = 1
          else:
              filterq[2 * count + 1, count] += 1
  
      filterq = filterq.tocsr()
      return filterq.T.dot(filterq)
  
  
  def _jacobi(
      filterq,
      x0: np.ndarray,
      lower_bound: np.ndarray,
      upper_bound: np.ndarray,
      max_iters: int = 10,
      rel_tol: float = 1e-6,
      weight: float = 0.5,
  ):
      """Jacobi method with constraints."""
  
      jacobi_r = sparse.lil_matrix(filterq)
      shp = jacobi_r.shape
      jacobi_d = 1.0 / filterq.diagonal()
      jacobi_r.setdiag((0,) * shp[0])
      jacobi_r = jacobi_r.tocsr()
  
      x = x0
  
      # We check the stopping criterion each 10 iterations
      check_each = 10
      cum_rel_tol = 1 - (1 - rel_tol) ** check_each
  
      energy_now = np.dot(x, filterq.dot(x)) / 2
      logging.info("Energy at iter %d: %.6g", 0, energy_now)
      for i in range(max_iters):
  
          x_1 = -jacobi_d * jacobi_r.dot(x)
          x = weight * x_1 + (1 - weight) * x
  
          # Constraints.
          x = np.maximum(x, lower_bound)
          x = np.minimum(x, upper_bound)
  
          # Stopping criterion
          if (i + 1) % check_each == 0:
              # Update energy
              energy_before = energy_now
              energy_now = np.dot(x, filterq.dot(x)) / 2
  
              logging.info("Energy at iter %d: %.6g", i + 1, energy_now)
  
              # Check stopping criterion
              cum_rel_improvement = (energy_before - energy_now) / energy_before
              if cum_rel_improvement < cum_rel_tol:
                  break
  
      return x
  
  
  def signed_distance_function(
      levelset: np.ndarray, band_radius: int
  ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
      """
      Return the distance to the 0.5 levelset of a function, the mask of the
      border (i.e., the nearest cells to the 0.5 level-set) and the mask of the
      band (i.e., the cells of the function whose distance to the 0.5 level-set
      is less of equal to `band_radius`).
      """
  
      binary_array = np.where(levelset > 0, True, False)
  
      # Compute the band and the border.
      dist_func = ndi.distance_transform_edt
      distance = np.where(
          binary_array, dist_func(binary_array) - 0.5, -dist_func(~binary_array) + 0.5
      )
      border = np.abs(distance) < 1
      band = np.abs(distance) <= band_radius
  
      return distance, border, band
  
  
  def signed_distance_function_iso0(
      levelset: np.ndarray, band_radius: int
  ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
      """
      Return the distance to the 0 levelset of a function, the mask of the
      border (i.e., the nearest cells to the 0 level-set) and the mask of the
      band (i.e., the cells of the function whose distance to the 0 level-set
      is less of equal to `band_radius`).
      """
  
      binary_array = levelset > 0
  
      # Compute the band and the border.
      dist_func = ndi.distance_transform_edt
      distance = np.where(
          binary_array, dist_func(binary_array), -dist_func(~binary_array)
      )
      border = np.zeros_like(levelset, dtype=bool)
      border[:-1, :, :] |= levelset[:-1, :, :] * levelset[1:, :, :] <= 0
      border[:, :-1, :] |= levelset[:, :-1, :] * levelset[:, 1:, :] <= 0
      border[:, :, :-1] |= levelset[:, :, :-1] * levelset[:, :, 1:] <= 0
      band = np.abs(distance) <= band_radius
  
      return distance, border, band
  
  
  def signed_distance_function_gpu(levelset: torch.Tensor, band_radius: int):
      binary_array = (levelset > 0).float()
  
      # Compute distance transform
      dist_pos = (
          F.max_pool3d(
              -binary_array.unsqueeze(0).unsqueeze(0), kernel_size=3, stride=1, padding=1
          )
          .squeeze(0)
          .squeeze(0)
          + binary_array
      )
      dist_neg = F.max_pool3d(
          (binary_array - 1).unsqueeze(0).unsqueeze(0), kernel_size=3, stride=1, padding=1
      ).squeeze(0).squeeze(0) + (1 - binary_array)
  
      distance = torch.where(binary_array > 0, dist_pos - 0.5, -dist_neg + 0.5)
  
      # breakpoint()
  
      # Use levelset as distance directly
      # distance = levelset
      # print(distance.shape)
      # Compute border and band
      border = torch.abs(distance) < 1
      band = torch.abs(distance) <= band_radius
  
      return distance, border, band
  
  
  def smooth_constrained(
      binary_array: np.ndarray,
      band_radius: int = 4,
      max_iters: int = 250,
      rel_tol: float = 1e-6,
  ) -> np.ndarray:
      """
      Implementation of the smoothing method from
  
      "Surface Extraction from Binary Volumes with Higher-Order Smoothness"
      Victor Lempitsky, CVPR10
      """
  
      # # Compute the distance map, the border and the band.
      logging.info("Computing distance transform...")
      # distance, _, band = signed_distance_function(binary_array, band_radius)
      binary_array_gpu = torch.from_numpy(binary_array).cuda()
      distance, _, band = signed_distance_function_gpu(binary_array_gpu, band_radius)
      distance = distance.cpu().numpy()
      band = band.cpu().numpy()
  
      variable_indices = _build_variable_indices(band)
  
      # Compute filterq.
      logging.info("Building matrix filterq...")
      if binary_array.ndim == 3:
          filterq = _buildq3d(variable_indices)
          # variable_indices_gpu = torch.from_numpy(variable_indices).cuda()
          # filterq_gpu = _buildq3d_gpu(variable_indices_gpu)
          # filterq = filterq_gpu.cpu().numpy()
      elif binary_array.ndim == 2:
          filterq = _buildq2d(variable_indices)
      else:
          raise ValueError("binary_array.ndim not in [2, 3]")
  
      # Initialize the variables.
      res = np.asarray(distance, dtype=np.double)
      x = res[band]
      upper_bound = np.where(x < 0, x, np.inf)
      lower_bound = np.where(x > 0, x, -np.inf)
  
      upper_bound[np.abs(upper_bound) < 1] = 0
      lower_bound[np.abs(lower_bound) < 1] = 0
  
      # Solve.
      logging.info("Minimizing energy...")
      x = _jacobi(
          filterq=filterq,
          x0=x,
          lower_bound=lower_bound,
          upper_bound=upper_bound,
          max_iters=max_iters,
          rel_tol=rel_tol,
      )
  
      res[band] = x
      return res
  
  
  def total_variation_denoising(x, weight=0.1, num_iterations=5, eps=1e-8):
      diff_x = torch.diff(x, dim=0, prepend=x[:1])
      diff_y = torch.diff(x, dim=1, prepend=x[:, :1])
      diff_z = torch.diff(x, dim=2, prepend=x[:, :, :1])
  
      norm = torch.sqrt(diff_x**2 + diff_y**2 + diff_z**2 + eps)
  
      div_x = torch.diff(diff_x / norm, dim=0, append=diff_x[-1:] / norm[-1:])
      div_y = torch.diff(diff_y / norm, dim=1, append=diff_y[:, -1:] / norm[:, -1:])
      div_z = torch.diff(diff_z / norm, dim=2, append=diff_z[:, :, -1:] / norm[:, :, -1:])
  
      return x - weight * (div_x + div_y + div_z)
  
  
  def smooth_constrained_gpu(
      binary_array: torch.Tensor,
      band_radius: int = 4,
      max_iters: int = 250,
      rel_tol: float = 1e-4,
  ):
      distance, _, band = signed_distance_function_gpu(binary_array, band_radius)
  
      # Initialize variables
      x = distance[band]
      upper_bound = torch.where(x < 0, x, torch.tensor(float("inf"), device=x.device))
      lower_bound = torch.where(x > 0, x, torch.tensor(float("-inf"), device=x.device))
  
      upper_bound[torch.abs(upper_bound) < 1] = 0
      lower_bound[torch.abs(lower_bound) < 1] = 0
  
      # Define the 3D Laplacian kernel
      laplacian_kernel = torch.tensor(
          [
              [
                  [
                      [[0, 1, 0], [1, -6, 1], [0, 1, 0]],
                      [[1, 0, 1], [0, 0, 0], [1, 0, 1]],
                      [[0, 1, 0], [1, 0, 1], [0, 1, 0]],
                  ]
              ]
          ],
          device=x.device,
      ).float()
  
      laplacian_kernel = laplacian_kernel / laplacian_kernel.abs().sum()
  
      breakpoint()
  
      # Simplified Jacobi iteration
      for i in range(max_iters):
          # Reshape x to 5D tensor (batch, channel, depth, height, width)
          x_5d = x.view(1, 1, *band.shape)
          x_3d = x.view(*band.shape)
  
          # Apply 3D convolution
          laplacian = F.conv3d(x_5d, laplacian_kernel, padding=1)
  
          # Reshape back to original dimensions
          laplacian = laplacian.view(x.shape)
  
          # Use a small relaxation factor to improve stability
          relaxation_factor = 0.1
          tv_weight = 0.1
          # x_new = x + relaxation_factor * laplacian
          x_new = total_variation_denoising(x_3d, weight=tv_weight)
          # Print laplacian min and max
          # print(f"Laplacian min: {laplacian.min().item():.4f}, max: {laplacian.max().item():.4f}")
  
          # Apply constraints
          # Reshape x_new to match the dimensions of lower_bound and upper_bound
          x_new = x_new.view(x.shape)
          x_new = torch.clamp(x_new, min=lower_bound, max=upper_bound)
  
          # Check for convergence
          diff_norm = torch.norm(x_new - x)
          print(diff_norm)
          x_norm = torch.norm(x)
  
          if x_norm > 1e-8:  # Avoid division by very small numbers
              relative_change = diff_norm / x_norm
              if relative_change < rel_tol:
                  break
          elif diff_norm < rel_tol:  # If x_norm is very small, check absolute change
              break
  
          x = x_new
  
          # Check for NaN and break if found, also check for inf
          if torch.isnan(x).any() or torch.isinf(x).any():
              print(f"NaN or Inf detected at iteration {i}")
              breakpoint()
              break
  
      result = distance.clone()
      result[band] = x
      return result
  
  
  def smooth_gaussian(binary_array: np.ndarray, sigma: float = 3) -> np.ndarray:
      vol = np.float_(binary_array) - 0.5
      return ndi.gaussian_filter(vol, sigma=sigma)
  
  
  def smooth_gaussian_gpu(binary_array: torch.Tensor, sigma: float = 3):
      # vol = binary_array.float()
      vol = binary_array
      kernel_size = int(2 * sigma + 1)
      kernel = torch.ones(
          1,
          1,
          kernel_size,
          kernel_size,
          kernel_size,
          device=binary_array.device,
          dtype=vol.dtype,
      ) / (kernel_size**3)
      return F.conv3d(
          vol.unsqueeze(0).unsqueeze(0), kernel, padding=kernel_size // 2
      ).squeeze()
  
  
  def smooth(binary_array: np.ndarray, method: str = "auto", **kwargs) -> np.ndarray:
      """
      Smooths the 0.5 level-set of a binary array. Returns a floating-point
      array with a smoothed version of the original level-set in the 0 isovalue.
  
      This function can apply two different methods:
  
      - A constrained smoothing method which preserves details and fine
        structures, but it is slow and requires a large amount of memory. This
        method is recommended when the input array is small (smaller than
        (500, 500, 500)).
      - A Gaussian filter applied over the binary array. This method is fast, but
        not very precise, as it can destroy fine details. It is only recommended
        when the input array is large and the 0.5 level-set does not contain
        thin structures.
  
      Parameters
      ----------
      binary_array : ndarray
          Input binary array with the 0.5 level-set to smooth.
      method : str, one of ['auto', 'gaussian', 'constrained']
          Smoothing method. If 'auto' is given, the method will be automatically
          chosen based on the size of `binary_array`.
  
      Parameters for 'gaussian'
      -------------------------
      sigma : float
          Size of the Gaussian filter (default 3).
  
      Parameters for 'constrained'
      ----------------------------
      max_iters : positive integer
          Number of iterations of the constrained optimization method
          (default 250).
      rel_tol: float
          Relative tolerance as a stopping criterion (default 1e-6).
  
      Output
      ------
      res : ndarray
          Floating-point array with a smoothed 0 level-set.
      """
  
      binary_array = np.asarray(binary_array)
  
      if method == "auto":
          if binary_array.size > 512**3:
              method = "gaussian"
          else:
              method = "constrained"
  
      if method == "gaussian":
          return smooth_gaussian(binary_array, **kwargs)
  
      if method == "constrained":
          return smooth_constrained(binary_array, **kwargs)
  
      raise ValueError("Unknown method '{}'".format(method))
  
  
  def smooth_gpu(binary_array: torch.Tensor, method: str = "auto", **kwargs):
      if method == "auto":
          method = "gaussian" if binary_array.numel() > 512**3 else "constrained"
  
      if method == "gaussian":
          return smooth_gaussian_gpu(binary_array, **kwargs)
      elif method == "constrained":
          return smooth_constrained_gpu(binary_array, **kwargs)
      else:
          raise ValueError(f"Unknown method '{method}'")
  