import torch


def dis_primal(
    input_view: torch.Tensor,
    param: torch.Tensor,
    n_sample: int,
) -> torch.Tensor:
  """Computes distortion penalty for the primal formulation.

  Let n be the number of samples in the view of interest and p the number of
  features. In the primal formulation, the 'param' matrix is the p*low_dim model
  parameter and input_view corresponds to the input data, of shape n*p. The
  distortion penalty can be written as

  distortion = ||input_view*input_view.T
                 - input_view*param*param.T*input_view.T||_2.

  The distortion is computed as is when n < p. However, if n > p, we
  compute the following formulation:

  distortion = torch.sqrt(Tr((I - param*param.T)*input_view.T*input_view
  *(I - param*param.T)*input_view.T*input_view))

  to avoid computing terms that are O(n**2) in memory or runtime.

  Arguments:
    input_view: torch.Tensor, one of the two views.
    param: torch.Tensor, model parameters.
    n_sample: int, sample size of entire dataset.
  Returns:
    distortion_value: torch.Tensor, scalar value.
  """
  n_sample, p_feature = input_view.shape
  if n_sample < p_feature:
    inner_prod = torch.matmul(input_view, input_view.t())
    tmp = torch.matmul(torch.matmul(
        torch.matmul(input_view, param), param.t()), input_view.t())
    tmp = (inner_prod - tmp)**2
    distortion_value = torch.sqrt(torch.sum(tmp))
  else:
    gram = torch.matmul(input_view.t(), input_view)
    tmp = torch.matmul(param, torch.matmul(param.t(), gram))
    prod = gram - tmp
    distortion_value = torch.sqrt(torch.trace(torch.matmul(prod, prod)))
  return distortion_value