# py
class PositionalEncoder(nn.Module):
  # sine-cosine positional encoder for input points.
  def __init__( self,
                d_input: int,
                n_freqs: int,
                log_space: bool = False ):
    super().__init__()
    self.d_input = d_input
    self.n_freqs = n_freqs         # 是不是視線上的采樣頻率？
    self.log_space = log_space
    self.d_output = d_input * (1 + 2 * self.n_freqs)
    self.embed_fns = [lambda x: x] # 冒號(hào)前面的x表示函數(shù)參數(shù)，后面的表示匿名函數(shù)運(yùn)算

    # Define frequencies in either linear or log scale
    if self.log_space:
      freq_bands = 2.**torch.linspace(0., self.n_freqs - 1, self.n_freqs)
    else:
      freq_bands = torch.linspace(2.**0., 2.**(self.n_freqs - 1), self.n_freqs)

    # Alternate sin and cos
    for freq in freq_bands:
      self.embed_fns.append(lambda x, freq=freq: torch.sin(x * freq))
      self.embed_fns.append(lambda x, freq=freq: torch.cos(x * freq))
  
  def forward(self, x) -> torch.Tensor:
    # Apply positional encoding to input.
    return torch.concat([fn(x) for fn in self.embed_fns], dim=-1)

class NeRF(nn.Module):
  # Neural radiance fields module.
  def __init__( self,
                d_input: int = 3,                  
                n_layers: int = 8,
                d_filter: int = 256,
                skip: Tuple[int] = (4,), # (4,)只有一個(gè)元素4的元組       
                d_viewdirs: Optional[int] = None): 

    super().__init__()
    self.d_input = d_input        # 這里是3D XYZ,？
    self.skip = skip              # 是要跳過什么？為啥要跳過？被遮擋？
    self.act = nn.functional.relu
    self.d_viewdirs = d_viewdirs  # d_viewdirs 是2D方向？

    # Create model layers
    # [if_true 就執(zhí)行的指令] if [if_true條件] else [if_false]
    # 是否skip的區(qū)別是，訓(xùn)練輸入維度是否多3維，
    # if i in skip =  if i in (4,),似乎是判斷i是否等于4
    # self.d_input=3 ：如果層id=4,網(wǎng)絡(luò)輸入要加3維，這是為什么？第4層有何特殊的？
    self.layers = nn.ModuleList(
      [nn.Linear(self.d_input, d_filter)] +
      [nn.Linear(d_filter + self.d_input, d_filter) if i in skip else \
       nn.Linear(d_filter  , d_filter) for i in range(n_layers - 1)]
    )

    # Bottleneck layers
    if self.d_viewdirs is not None:
      # If using viewdirs, split alpha and RGB
      self.alpha_out = nn.Linear(d_filter, 1)
      self.rgb_filters = nn.Linear(d_filter, d_filter)
      self.branch = nn.Linear(d_filter + self.d_viewdirs, d_filter // 2)
      self.output = nn.Linear(d_filter // 2, 3) # 為啥要取一半？
    else:
      # If no viewdirs, use simpler output
      self.output = nn.Linear(d_filter, 4) # d_filter=256，輸出是4維RGBA
  
  def forward(self,
              x: torch.Tensor, # ？
              viewdirs: Optional[torch.Tensor] = None) -> torch.Tensor:
   
    # Forward pass with optional view direction.
    if self.d_viewdirs is None and viewdirs is not None:
      raise ValueError('Cannot input x_direction')

    # Apply forward pass up to bottleneck
    x_input = x    # 這里的x是幾維？從下面的分離RGB和A看，應(yīng)該是4D
    # 下面通過8層MLP訓(xùn)練RGBA
    for i, layer in enumerate(self.layers):  # 8層，每一層進(jìn)行運(yùn)算
      x = self.act(layer(x)) 
      if i in self.skip:
        x = torch.cat([x, x_input], dim=-1)

    # Apply bottleneck  bottleneck 瓶頸是啥？是不是最費(fèi)算力的模塊？
    if self.d_viewdirs is not None:
      # 從網(wǎng)絡(luò)輸出分離A,RGB還需要經(jīng)過更多訓(xùn)練
      alpha = self.alpha_out(x)  
      
      # Pass through bottleneck to get RGB
      x = self.rgb_filters(x) 
      x = torch.concat([x, viewdirs], dim=-1)
      x = self.act(self.branch(x)) # self.branch shape: (d_filter // 2)
      x = self.output(x)           # self.output shape: (3)

      # Concatenate alphas to output
      x = torch.concat([x, alpha], dim=-1)
    else:
      # Simple output
      x = self.output(x)
    return x

def raw2outputs(raw: torch.Tensor,
                z_vals: torch.Tensor,
                rays_d: torch.Tensor,
                raw_noise_std: float = 0.0,
                white_bkgd: bool = False) 
  -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

  # 將原始的NeRF輸出轉(zhuǎn)為RGB和其他映射
  # Difference between consecutive elements of `z_vals`. [n_rays, n_samples]
  dists = z_vals[..., 1:] - z_vals[..., :-1]# ？這里減法的意義是啥？
  dists = torch.cat([dists, 1e10 * torch.ones_like(dists[..., :1])], dim=-1)

  # 將每個(gè)距離乘以其對(duì)應(yīng)方向光線的范數(shù)，以轉(zhuǎn)換為真實(shí)世界的距離（考慮非單位方向）
  dists = dists * torch.norm(rays_d[..., None, :], dim=-1)

  # 將噪聲添加到模型對(duì)密度的預(yù)測(cè)中，用于在訓(xùn)練期間規(guī)范網(wǎng)絡(luò)(防止漂浮物偽影)
  noise = 0.
  if raw_noise_std > 0.:
    noise = torch.randn(raw[..., 3].shape) * raw_noise_std

  # Predict density of each sample along each ray. Higher values imply
  # higher likelihood of being absorbed at this point. [n_rays, n_samples]
  alpha = 1.0 - torch.exp(-nn.functional.relu(raw[..., 3] + noise) * dists)

  # Compute weight for RGB of each sample along each ray. [n_rays, n_samples]
  # The higher the alpha, the lower subsequent weights are driven.
  weights = alpha * cumprod_exclusive(1. - alpha + 1e-10)

  # Compute weighted RGB map.
  rgb = torch.sigmoid(raw[..., :3])  # [n_rays, n_samples, 3]
  rgb_map = torch.sum(weights[..., None] * rgb, dim=-2)  # [n_rays, 3]

  # Estimated depth map is predicted distance.
  depth_map = torch.sum(weights * z_vals, dim=-1)

  # Disparity map is inverse depth.
  disp_map = 1. / torch.max(1e-10 * torch.ones_like(depth_map),
                            depth_map / torch.sum(weights, -1))

  # Sum of weights along each ray. In [0, 1] up to numerical error.
  acc_map = torch.sum(weights, dim=-1)

  # To composite onto a white background, use the accumulated alpha map.
  if white_bkgd:
    rgb_map = rgb_map + (1. - acc_map[..., None])

  return rgb_map, depth_map, acc_map, weights


def cumprod_exclusive(tensor: torch.Tensor) -> torch.Tensor:
  # (Courtesy of https://github.com/krrish94/nerf-pytorch)
  # Compute regular cumprod first.
  cumprod = torch.cumprod(tensor, -1)
  # "Roll" the elements along dimension 'dim' by 1 element.
  cumprod = torch.roll(cumprod, 1, -1)
  # Replace the first element by "1" as this is what tf.cumprod(..., exclusive=True) does.
  cumprod[..., 0] = 1.
  return cumprod

def sample_stratified(rays_o: torch.Tensor,
                      rays_d: torch.Tensor,
                      near: float,
                      far: float,
                      n_samples: int,
                      perturb: Optional[bool] = True,
                      inverse_depth: bool = False)
   -> Tuple[torch.Tensor, torch.Tensor]:
  # Sample along ray from regularly-spaced bins.
  # Grab samples for space integration along ray
  t_vals = torch.linspace(0., 1., n_samples, device=rays_o.device)
  if not inverse_depth:
    # Sample linearly between `near` and `far`
    z_vals = near * (1.-t_vals) + far * (t_vals)
  else:
    # Sample linearly in inverse depth (disparity)
    z_vals = 1./(1./near * (1.-t_vals) + 1./far * (t_vals))

  # Draw uniform samples from bins along ray
  if perturb:
    mids = .5 * (z_vals[1:] + z_vals[:-1])
    upper = torch.concat([mids, z_vals[-1:]], dim=-1)
    lower = torch.concat([z_vals[:1], mids], dim=-1)
    t_rand = torch.rand([n_samples], device=z_vals.device)
    z_vals = lower + (upper - lower) * t_rand
  z_vals = z_vals.expand(list(rays_o.shape[:-1]) + [n_samples])

  # Apply scale from `rays_d` and offset from `rays_o` to samples
  # pts: (width, height, n_samples, 3)
  pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals[..., :, None]
  return pts, z_vals

def sample_hierarchical(rays_o: torch.Tensor,
                        rays_d: torch.Tensor,
                        z_vals: torch.Tensor,
                        weights: torch.Tensor,
                        n_samples: int,
                        perturb: bool = False) 
  -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
  # Apply hierarchical sampling to the rays.
  # Draw samples from PDF using z_vals as bins and weights as probabilities.
  z_vals_mid = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
  new_z_samples = sample_pdf(z_vals_mid, weights[..., 1:-1], n_samples, perturb=perturb)
  new_z_samples = new_z_samples.detach()

  # Resample points from ray based on PDF.
  z_vals_combined, _ = torch.sort(torch.cat([z_vals, new_z_samples], dim=-1), dim=-1)
  # [N_rays, N_samples + n_samples, 3]
  pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals_combined[..., :, None]  
  return pts, z_vals_combined, new_z_samples

def sample_pdf(bins: torch.Tensor,
               weights: torch.Tensor,
               n_samples: int,
               perturb: bool = False) -> torch.Tensor:
  # Apply inverse transform sampling to a weighted set of points.
  # Normalize weights to get PDF.
  # [n_rays, weights.shape[-1]]
  pdf = (weights + 1e-5) / torch.sum(weights + 1e-5, -1, keepdims=True) 

  # Convert PDF to CDF.
  cdf = torch.cumsum(pdf, dim=-1) # [n_rays, weights.shape[-1]]
  # [n_rays, weights.shape[-1] + 1]
  cdf = torch.concat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1) 

  # Take sample positions to grab from CDF. Linear when perturb == 0.
  if not perturb:
    u = torch.linspace(0., 1., n_samples, device=cdf.device)
    u = u.expand(list(cdf.shape[:-1]) + [n_samples]) # [n_rays, n_samples]
  else:
    # [n_rays, n_samples]
    u = torch.rand(list(cdf.shape[:-1]) + [n_samples], device=cdf.device) 

  # Find indices along CDF where values in u would be placed.
  u = u.contiguous() # Returns contiguous tensor with same values.
  inds = torch.searchsorted(cdf, u, right=True) # [n_rays, n_samples]

  # Clamp indices that are out of bounds.
  below = torch.clamp(inds - 1, min=0)
  above = torch.clamp(inds, max=cdf.shape[-1] - 1)
  inds_g = torch.stack([below, above], dim=-1) # [n_rays, n_samples, 2]

  # Sample from cdf and the corresponding bin centers.
  matched_shape = list(inds_g.shape[:-1]) + [cdf.shape[-1]]
  cdf_g = torch.gather(cdf.unsqueeze(-2).expand(matched_shape), dim=-1,index=inds_g)
  bins_g = torch.gather(bins.unsqueeze(-2).expand(matched_shape), dim=-1, index=inds_g)

  # Convert samples to ray length.
  denom = (cdf_g[..., 1] - cdf_g[..., 0])
  denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
  t = (u - cdf_g[..., 0]) / denom
  samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])

  return samples # [n_rays, n_samples]

def init_models():
  # Initialize models, encoders, and optimizer for NeRF training.
  encoder = PositionalEncoder(d_input, n_freqs, log_space=log_space)
  encode = lambda x: encoder(x)
  # View direction encoders
  if use_viewdirs:
    encoder_viewdirs = PositionalEncoder(d_input, n_freqs_views,log_space=log_space)
    encode_viewdirs  = lambda x: encoder_viewdirs(x)
    d_viewdirs       = encoder_viewdirs.d_output
  else:
    encode_viewdirs = None
    d_viewdirs = None

  model = NeRF(encoder.d_output, 
               n_layers=n_layers, 
               d_filter=d_filter, skip=skip,d_viewdirs=d_viewdirs)
  model.to(device)
  model_params = list(model.parameters())
  if use_fine_model:
    fine_model = NeRF(encoder.d_output, 
                      n_layers=n_layers, 
                      d_filter=d_filter, skip=skip,d_viewdirs=d_viewdirs)
    fine_model.to(device)
    model_params = model_params + list(fine_model.parameters())
  else:
    fine_model = None

  optimizer      = torch.optim.Adam(model_params, lr=lr)
  warmup_stopper = EarlyStopping(patience=50)
  return model, fine_model, encode, encode_viewdirs, optimizer, warmup_stopper

def train():
  # Launch training session for NeRF.
  # Shuffle rays across all images.
  if not one_image_per_step:
    height, width = images.shape[1:3]
    all_rays = torch.stack([torch.stack(get_rays(height, width, focal, p), 0)
                           for p in poses[:n_training]], 0)
    rays_rgb = torch.cat([all_rays, images[:, None]], 1)
    rays_rgb = torch.permute(rays_rgb, [0, 2, 3, 1, 4])
    rays_rgb = rays_rgb.reshape([-1, 3, 3])
    rays_rgb = rays_rgb.type(torch.float32)
    rays_rgb = rays_rgb[torch.randperm(rays_rgb.shape[0])]
    i_batch = 0

  train_psnrs = []
  val_psnrs = []
  iternums = []
  for i in trange(n_iters):
    model.train()
    if one_image_per_step:
      # Randomly pick an image as the target.
      target_img_idx = np.random.randint(images.shape[0])
      target_img     = images[target_img_idx].to(device)
      if center_crop and i < center_crop_iters:
        target_img = crop_center(target_img)
      height, width = target_img.shape[:2]
      target_pose = poses[target_img_idx].to(device)
      rays_o, rays_d = get_rays(height, width, focal, target_pose)
      rays_o = rays_o.reshape([-1, 3])
      rays_d = rays_d.reshape([-1, 3])
    else:
      # Random over all images.
      batch = rays_rgb[i_batch:i_batch + batch_size]
      batch = torch.transpose(batch, 0, 1)
      rays_o, rays_d, target_img = batch
      height, width = target_img.shape[:2]
      i_batch += batch_size
      # Shuffle after one epoch
      if i_batch >= rays_rgb.shape[0]:
          rays_rgb = rays_rgb[torch.randperm(rays_rgb.shape[0])]
          i_batch = 0
    target_img = target_img.reshape([-1, 3])

    # Run one iteration of TinyNeRF and get the rendered RGB image.
    outputs = nerf_forward(rays_o, rays_d,
                           near, far, encode, model,
                           kwargs_sample_stratified=kwargs_sample_stratified,
                           n_samples_hierarchical=n_samples_hierarchical,
                           kwargs_sample_hierarchical=kwargs_sample_hierarchical,
                           fine_model=fine_model,
                           viewdirs_encoding_fn=encode_viewdirs,
                           chunksize=chunksize)
    # Backprop!
    rgb_predicted = outputs['rgb_map']
    loss = torch.nn.functional.mse_loss(rgb_predicted, target_img)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    psnr = -10. * torch.log10(loss)
    train_psnrs.append(psnr.item())

    # Evaluate testimg at given display rate.
    if i % display_rate == 0:
      model.eval()
      height, width = testimg.shape[:2]
      rays_o, rays_d = get_rays(height, width, focal, testpose)
      rays_o = rays_o.reshape([-1, 3])
      rays_d = rays_d.reshape([-1, 3])
      outputs = nerf_forward(rays_o, rays_d,
                             near, far, encode, model,
                             kwargs_sample_stratified=kwargs_sample_stratified,
                             n_samples_hierarchical=n_samples_hierarchical,
                             kwargs_sample_hierarchical=kwargs_sample_hierarchical,
                             fine_model=fine_model,
                             viewdirs_encoding_fn=encode_viewdirs,
                             chunksize=chunksize)
      rgb_predicted = outputs['rgb_map']
      loss = torch.nn.functional.mse_loss(rgb_predicted, testimg.reshape(-1, 3))
      val_psnr = -10. * torch.log10(loss)
      val_psnrs.append(val_psnr.item())
      iternums.append(i)

    # Check PSNR for issues and stop if any are found.
    if i == warmup_iters - 1:
      if val_psnr < warmup_min_fitness:
        return False, train_psnrs, val_psnrs
    elif i < warmup_iters:
      if warmup_stopper is not None and warmup_stopper(i, psnr):
        return False, train_psnrs, val_psnrs

  return True, train_psnrs, val_psnrs

# Run training session(s)
for _ in range(n_restarts):
  model, fine_model, encode, encode_viewdirs, optimizer, warmup_stopper = init_models()
  success, train_psnrs, val_psnrs = train()
  if success and val_psnrs[-1] >= warmup_min_fitness:
    print('Training successful!')
    break
print(f'Done!')