text: from abc import ABC, abstractmethod
  
  import torch
  from fvcore.common.registry import Registry
  
  from meddlr.data.transforms.transform import build_normalizer
  from meddlr.ops import complex as cplx
  from meddlr.utils import transforms as T
  
  LOSS_COMPUTER_REGISTRY = Registry("LOSS_COMPUTER")  # noqa F401 isort:skip
  LOSS_COMPUTER_REGISTRY.__doc__ = """
  Registry for loss computers.
  
  The registered object will be called with `obj(cfg)`
  and expected to return a LossComputer object.
  """
  
  EPS = 1e-11
  IMAGE_LOSSES = ["l1", "l2", "psnr", "nrmse", "mag_l1", "perp_loss"]
  KSPACE_LOSSES = ["k_l1", "k_l1_normalized", "k_l1_l2_sum_normalized"]
  
  
  def build_loss_computer(cfg, name, **kwargs):
      return LOSS_COMPUTER_REGISTRY.get(name)(cfg, **kwargs)
  
  
  class LossComputer(ABC):
      def __init__(self, cfg):
          self._normalizer = build_normalizer(cfg)
  
      @abstractmethod
      def __call__(self, input, output):
          pass
  
      def _get_metrics(self, target: torch.Tensor, output: torch.Tensor, loss_name):
          # Compute metrics
          if loss_name == "mag_l1":
              abs_error = torch.abs(output - target)
              abs_mag_error = abs_error
          else:
              abs_error = cplx.abs(output - target)
              abs_mag_error = torch.abs(cplx.abs(output) - cplx.abs(target))
          l1 = torch.mean(abs_error)
          mag_l1 = torch.mean(abs_mag_error)
          N = target.shape[0]
  
          abs_error = abs_error.view(N, -1)
          if loss_name == "mag_l1":
              tgt_mag = torch.abs(target).view(N, -1)
          else:
              tgt_mag = cplx.abs(target).view(N, -1)
          l2 = torch.sqrt(torch.mean(abs_error**2, dim=1))
          psnr = 20 * torch.log10(tgt_mag.max(dim=1)[0] / (l2 + EPS))
          nrmse = l2 / torch.sqrt(torch.mean(tgt_mag**2, dim=1))
  
          metrics_dict = {
              "l1": l1,
              "l2": l2.mean(),
              "psnr": psnr.mean(),
              "nrmse": nrmse.mean(),
              "mag_l1": mag_l1,
          }
  
          if loss_name == "perp_loss":
              metrics_dict.update(perp_loss(output, target))
  
          if loss_name in KSPACE_LOSSES:
              target, output = T.fft2(target), T.fft2(output)
              abs_error = cplx.abs(target - output)
              if loss_name == "k_l1":
                  metrics_dict["loss"] = torch.mean(abs_error)
              elif loss_name == "k_l1_normalized":
                  metrics_dict["loss"] = torch.mean(abs_error / (cplx.abs(target) + EPS))
              elif loss_name == "k_l1_l2_sum_normalized":
                  kl1_norm = torch.sum(abs_error) / torch.sum(cplx.abs(target))
                  kl2_norm = torch.sqrt(torch.sum(abs_error**2)) / torch.sqrt(
                      torch.sum(cplx.abs(target) ** 2)
                  )  # noqa: E501
                  metrics_dict["loss"] = 0.5 * kl1_norm + 0.5 * kl2_norm
              else:
                  assert False  # should not reach here
          else:
              loss = metrics_dict[loss_name]
              metrics_dict["loss"] = loss
  
          return metrics_dict
  
  
  @LOSS_COMPUTER_REGISTRY.register()
  class BasicLossComputer(LossComputer):
      def __init__(self, cfg):
          super().__init__(cfg)
          loss_name = cfg.MODEL.RECON_LOSS.NAME
          assert loss_name in IMAGE_LOSSES or loss_name in KSPACE_LOSSES
          self.loss = loss_name
          self.renormalize_data = cfg.MODEL.RECON_LOSS.RENORMALIZE_DATA
  
      def __call__(self, input, output):
          pred: torch.Tensor = output["pred"]
          target = output["target"].to(pred.device)
  
          if self.renormalize_data:
              normalization_args = {k: input.get(k, output.get(k, None)) for k in ["mean", "std"]}
              normalized = self._normalizer.undo(
                  image=pred,
                  target=target,
                  mean=normalization_args["mean"],
                  std=normalization_args["std"],
              )
              output = normalized["image"]
              target = normalized["target"]
          else:
              output = pred
  
          metrics_dict = self._get_metrics(target, output, self.loss)
          return metrics_dict
  
  
  @LOSS_COMPUTER_REGISTRY.register()
  class N2RLossComputer(LossComputer):
      def __init__(self, cfg):
          super().__init__(cfg)
          recon_loss = cfg.MODEL.RECON_LOSS.NAME
          consistency_loss = cfg.MODEL.CONSISTENCY.LOSS_NAME
          latent_loss = cfg.MODEL.CONSISTENCY.LATENT_LOSS_NAME
  
          assert recon_loss in IMAGE_LOSSES or recon_loss in KSPACE_LOSSES
          assert consistency_loss in IMAGE_LOSSES or consistency_loss in KSPACE_LOSSES
  
          self.recon_loss = recon_loss
          self.consistency_loss = consistency_loss
          self.latent_loss = latent_loss
          self.renormalize_data = cfg.MODEL.RECON_LOSS.RENORMALIZE_DATA
          self.consistency_weight = cfg.MODEL.CONSISTENCY.LOSS_WEIGHT
          self.latent_weight = cfg.MODEL.CONSISTENCY.LATENT_LOSS_WEIGHT
          self.use_latent = cfg.MODEL.CONSISTENCY.USE_LATENT
          self.use_consistency = cfg.MODEL.CONSISTENCY.USE_CONSISTENCY
          self.num_latent_layers = cfg.MODEL.CONSISTENCY.NUM_LATENT_LAYERS
          self.latent_keys = ["E4", "E3", "D3", "E2", "D2", "E1", "D1"]
          # self.use_robust = cfg.MODEL.LOSS.USE_ROBUST
          # self.beta = cfg.MODEL.LOSS.BETA
          # self.robust_step_size = cfg.MODEL.LOSS.ROBUST_STEP_SIZE
  
      def _compute_metrics(self, input, output, loss):
          """Computes image metrics on prediction and target data."""
          if output is None or len(output) == 0:
              return {k: torch.Tensor([0.0]).detach() for k in ["l1", "l2", "psnr", "loss"]}
  
          pred: torch.Tensor = output["pred"]
          target = output["target"].to(pred.device)
          if self.renormalize_data:
              normalized = self._normalizer.undo(
                  image=pred, target=target, mean=input["mean"], std=input["std"]
              )
              output = normalized["image"]
              target = normalized["target"]
          else:
              output = pred
  
          # Compute metrics
          metrics_dict = self._get_metrics(target, output, loss)
          return metrics_dict
  
      # def compute_robust_loss(self, group_loss):
      #     if torch.is_grad_enabled():  # update adv_probs if in training mode
      #         adjusted_loss = group_loss
      #         if self.do_adj:
      #             adjusted_loss += self.loss_adjustment
      #         logit_step = self.robust_step_size * adjusted_loss.data
      #         if self.stable:
      #             self.adv_probs_logits = self.adv_probs_logits + logit_step
      #         else:
      #             self.adv_probs = self.adv_probs * torch.exp(logit_step)
      #             self.adv_probs = self.adv_probs / self.adv_probs.sum()
      #
      #     if self.stable:
      #         adv_probs = torch.softmax(self.adv_probs_logits, dim=-1)
      #     else:
      #         adv_probs = self.adv_probs
      #     robust_loss = group_loss @ adv_probs
      #     return robust_loss, adv_probs
  
      def __call__(self, input, output):
          output_recon = output.get("recon", None)
          output_consistency = output.get("consistency", None)
  
          loss = 0
          metrics_recon = {
              "recon_{}".format(k): v
              for k, v in self._compute_metrics(
                  input.get("supervised", None), output_recon, self.recon_loss
              ).items()
          }
          if output_recon is not None:
              loss += metrics_recon["recon_loss"]
  
          metrics_consistency = {
              "cons_{}".format(k): v
              for k, v in self._compute_metrics(
                  input.get("unsupervised", None), output_consistency, self.consistency_loss
              ).items()  # noqa
          }
          if output_consistency is not None and self.use_consistency:
              loss += self.consistency_weight * metrics_consistency["cons_loss"]
  
          if self.use_latent:
              num_losses = self.num_latent_layers * 2 - 1
              all_metrics_latent = []
              for i in range(num_losses):
                  output_latent = {}
                  output_latent["target"] = output_recon["latent"][self.latent_keys[i]]
                  output_latent["pred"] = output_consistency["latent"][self.latent_keys[i]]
  
                  metrics_latent = {
                      "latent_" + self.latent_keys[i] + "_{}".format(k): v
                      for k, v in self._compute_metrics(
                          None, output_latent, self.latent_loss
                      ).items()  # noqa
                  }
  
                  all_metrics_latent.append(metrics_latent)
                  loss += (
                      self.latent_weight * metrics_latent["latent_" + self.latent_keys[i] + "_loss"]
                  )
  
          metrics = {}
          if output_consistency is not None and self.use_consistency:
              metrics.update(metrics_consistency)
          if output_recon is not None:
              metrics.update(metrics_recon)
          if self.use_latent:
              for i in range(num_losses):
                  metrics.update(all_metrics_latent[i])
  
          metrics["loss"] = loss
          return metrics
  
  
  def perp_loss(yhat, y):
      """Implementation of the perpendicular loss.
  
      Args:
          yhat: Predicted reconstruction. Must be complex.
          y: Target reconstruction. Must be complex.
  
      Returns:
          Dict[str, scalar]:
  
      References:
          Terpstra, et al. "Rethinking complex image reconstruction:
          ⟂-loss for improved complex image reconstruction with deep learning."
          International Society of Magnetic Resonance in Medicine Annual Meeting
          2021.
      """
      if cplx.is_complex(yhat):
          yhat = torch.view_as_real(yhat)
      if cplx.is_complex(y):
          y = torch.view_as_real(y)
  
      P = torch.abs(yhat[..., 0] * y[..., 1] - yhat[..., 1] * y[..., 0]) / cplx.abs(y)
      l1 = torch.abs(cplx.abs(y) - cplx.abs(yhat))
  
      return {"p_perp_loss": torch.mean(P), "perp_loss": torch.mean(P + l1)}
  