text: # Copyright 2019-2020 Xanadu Quantum Technologies Inc.
  
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  
  #     http://www.apache.org/licenses/LICENSE-2.0
  
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """
  Functions for constructing/calculating the means, variances and covariances of
  Gaussian states.
  """
  
  from itertools import product
  
  from scipy.special import factorial
  
  import numpy as np
  
  from .._hafnian import hafnian, reduction
  from .._torontonian import threshold_detection_prob
  
  from .conversions import reduced_gaussian, Qmat, Xmat, complex_to_real_displacements
  
  
  def photon_number_mean(mu, cov, j, hbar=2):
      r"""Calculate the mean photon number of mode j of a Gaussian state.
  
      Args:
          mu (array): vector of means of the Gaussian state using the ordering
              :math:`[q_1, q_2, \dots, q_n, p_1, p_2, \dots, p_n]`
          cov (array): the covariance matrix of the Gaussian state
          j (int): the j :sup:`th` mode
          hbar (float): the ``hbar`` convention used in the commutation
              relation :math:`[q, p]=i\hbar`
  
      Returns:
          float: the mean photon number in mode :math:`j`.
      """
      num_modes = len(mu) // 2
      return (
          mu[j] ** 2 + mu[j + num_modes] ** 2 + cov[j, j] + cov[j + num_modes, j + num_modes] - hbar
      ) / (2 * hbar)
  
  
  def photon_number_mean_vector(mu, cov, hbar=2):
      r"""Calculate the mean photon number of each of the modes in a Gaussian state
  
      Args:
          mu (array): vector of means of the Gaussian state using the ordering
              :math:`[q_1, q_2, \dots, q_n, p_1, p_2, \dots, p_n]`
          cov (array): the covariance matrix of the Gaussian state
          hbar (float): the ``hbar`` convention used in the commutation
              relation :math:`[q, p]=i\hbar`
  
      Returns:
          array: the vector of means of the photon number distribution
      """
  
      N = len(mu) // 2
      return np.array([photon_number_mean(mu, cov, j, hbar=hbar) for j in range(N)])
  
  
  def photon_number_covar(mu, cov, j, k, hbar=2):
      r""" Calculate the variance/covariance of the photon number distribution
      of a Gaussian state.
  
      Implements the covariance matrix of the photon number distribution of a
      Gaussian state according to the Last two eq. of Part II. in
      `'Multidimensional Hermite polynomials and photon distribution for polymode
      mixed light', Dodonov et al. <https://journals.aps.org/pra/abstract/10.1103/PhysRevA.50.813>`_
  
      .. math::
          \sigma_{n_j n_j} &= \frac{1}{2}\left(T_j^2 - 2d_j - \frac{1}{2}\right)
          + \left<\mathbf{Q}_j\right>\mathcal{M}_j\left<\mathbf{Q}_j\right>, \\
          \sigma_{n_j n_k} &= \frac{1}{2}\mathrm{Tr}\left(\Lambda_j \mathbf{M} \Lambda_k \mathbf{M}\right)
          + \left<\mathbf{Q}\right>\Lambda_j \mathbf{M} \Lambda_k\left<\mathbf{Q}\right>,
  
      where :math:`T_j` and :math:`d_j` are the trace and the determinant of
      :math:`2 \times 2` matrix :math:`\mathcal{M}_j` whose elements coincide
      with the nonzero elements of matrix :math:`\mathbf{M}_j = \Lambda_j \mathbf{M} \Lambda_k`
      while the two-vector :math:`\mathbf{Q}_j` has the components :math:`(q_j, p_j)`.
      :math:`2N \times 2N` projector matrix :math:`\Lambda_j` has only two nonzero
      elements: :math:`\left(\Lambda_j\right)_{jj} = \left(\Lambda_j\right)_{j+N,j+N} = 1`.
      Note that the convention for ``mu`` used here differs from the one used in Dodonov et al.,
      They both provide the same results in this particular case.
      Also note that the original reference of Dodonov et al. has an incorrect prefactor of 1/2
      in the last terms of the last equation above.
  
      Args:
          mu (array): vector of means of the Gaussian state using the ordering
              :math:`[q_1, q_2, \dots, q_n, p_1, p_2, \dots, p_n]`
          cov (array): the covariance matrix of the Gaussian state
          j (int): the j :sup:`th` mode
          k (int): the k :sup:`th` mode
          hbar (float): the ``hbar`` convention used in the commutation
              relation :math:`[q, p]=i\hbar`
  
      Returns:
          float: the covariance for the photon numbers at modes :math:`j` and  :math:`k`.
      """
  
      if j == k:
          mu, cov = reduced_gaussian(mu, cov, [j])
          term_1 = 0.5 * np.trace(cov) ** 2 - np.linalg.det(cov)
          term_2 = mu @ cov @ mu
          return ((term_1 + term_2) / hbar ** 2) - 0.25
  
      mu, cov = reduced_gaussian(mu, cov, [j, k])
      term_1 = cov[0, 1] ** 2 + cov[0, 3] ** 2 + cov[2, 1] ** 2 + cov[2, 3] ** 2
      term_2 = (
          cov[0, 1] * mu[0] * mu[1]
          + cov[2, 1] * mu[1] * mu[2]
          + cov[0, 3] * mu[0] * mu[3]
          + cov[2, 3] * mu[2] * mu[3]
      )
  
      return (term_1 + 2 * term_2) / (2 * hbar ** 2)
  
  
  def photon_number_covmat(mu, cov, hbar=2):
      r"""Calculate the covariance matrix of the photon number distribution of a
      Gaussian state.
  
      Args:
          mu (array): vector of means of the Gaussian state using the ordering
              :math:`[q_1, q_2, \dots, q_n, p_1, p_2, \dots, p_n]`
          cov (array): the covariance matrix of the Gaussian state
          hbar (float): the ``hbar`` convention used in the commutation
              relation :math:`[q, p]=i\hbar`
  
      Returns:
          array: the covariance matrix of the photon number distribution
      """
      N = len(mu) // 2
      pnd_cov = np.zeros((N, N))
      for i in range(N):
          for j in range(i + 1):
              pnd_cov[i][j] = photon_number_covar(mu, cov, i, j, hbar=hbar)
              pnd_cov[j][i] = pnd_cov[i][j]
      return pnd_cov
  
  
  def photon_number_expectation(mu, cov, modes, hbar=2):
      r"""Calculates the expectation value of the product of the number operator of the modes in a Gaussian state.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          modes (list): list of modes
          hbar (float): value of hbar in the uncertainty relation.
  
      Returns:
          (float): expectation value of the product of the number operators of the modes.
      """
      n, _ = cov.shape
      n_modes = n // 2
      rpt = np.zeros([n], dtype=int)
      for i in modes:
          rpt[i] = 1
          rpt[i + n_modes] = 1
  
      return normal_ordered_expectation(mu, cov, rpt, hbar=hbar)
  
  
  def photon_number_squared_expectation(mu, cov, modes, hbar=2):
      r"""Calculates the expectation value of the square of the product of the number operator of the modes in
      a Gaussian state.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          modes (list): list of modes
          hbar (float): value of hbar in the uncertainty relation.
  
      Returns:
          (float): expectation value of the square of the product of the number operator of the modes.
      """
      n_modes = len(modes)
  
      mu_red, cov_red = reduced_gaussian(mu, cov, modes)
      result = 0
      for item in product([1, 2], repeat=n_modes):
          rpt = item + item
          term = normal_ordered_expectation(mu_red, cov_red, rpt, hbar=hbar)
          result += term
      return result
  
  
  def normal_ordered_expectation(mu, cov, rpt, hbar=2):
      r"""Calculates the expectation value of the normal ordered product
      :math:`\prod_{i=0}^{N-1} a_i^{\dagger n_i} \prod_{j=0}^{N-1} a_j^{m_j}` with respect to an N-mode Gaussian state,
      where :math:`\text{rpt}=(n_0, n_1, \ldots, n_{N-1}, m_0, m_1, \ldots, m_{N-1})`.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          rpt (list): integers specifying the terms to calculate.
          hbar (float): value of hbar in the uncertainty relation.
  
      Returns:
          (float): expectation value of the normal ordered product of operators
      """
      return s_ordered_expectation(mu, cov, rpt, hbar, s=1)
  
  
  def s_ordered_expectation(mu, cov, rpt, hbar=2, s=0):
      r"""Calculates the expectation value of the s-ordered product
      obtained by taking deirvatives of the characteristic function of a Gaussian states,
      Here, :math:`\text{rpt}=(n_0, n_1, \ldots, n_{N-1}, m_0, m_1, \ldots, m_{N-1})`.
      indicates how many derivatives are taken with respect to the complex argument and its
      conjugate.
      The values :math:`s=\{1,0,-1\}` correspond respectively to normal, symmetric and antinormal order.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          rpt (list): integers specifying the terms to calculate.
          hbar (float): value of hbar in the uncertainty relation.
          s (float): value setting the ordering it must be between -1 and 1.
  
      Returns:
          (float): expectation value of the normal ordered product of operators
      """
  
      # The following seven lines are written so that we remove from the calculation the
      # modes k that we don't care about. These modes have rpt[k] = rpt[k+M] = 0
      if np.allclose(rpt, 0):
          return 1.0
      M = len(cov) // 2
      modes = np.where(np.array(rpt[0:M]) + np.array(rpt[M : 2 * M]) != 0)[0]
      mu, cov = reduced_gaussian(mu, cov, list(modes))
      ind = list(modes) + list(modes + M)
      rpt = list(np.array(rpt)[np.array(ind)])
  
      alpha = complex_to_real_displacements(mu, hbar=hbar)
      n = len(cov)
      V = (Qmat(cov, hbar=hbar) - 0.5 * (s + 1) * np.identity(n)) @ Xmat(n // 2)
      A = reduction(V, rpt)
      if np.allclose(mu, 0):
          return hafnian(A)
  
      np.fill_diagonal(A, reduction(np.conj(alpha), rpt))
      return hafnian(A, loop=True)
  
  
  def mean_clicks(cov, hbar=2):
      r"""Calculates the total mean number of clicks when a zero-mean gaussian state
      is measured using threshold detectors.
  
      Args
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering
          hbar (float): the value of :math:`\hbar` in the commutation relation :math:`[\x,\p]=i\hbar`
  
      Returns
          float: mean number of clicks
      """
      n, _ = cov.shape
      nmodes = n // 2
      Q = Qmat(cov, hbar=hbar)
      meanc = 1.0 * nmodes
  
      for i in range(nmodes):
          det_val = np.real(Q[i, i] * Q[i + nmodes, i + nmodes] - Q[i + nmodes, i] * Q[i, i + nmodes])
          meanc -= 1.0 / np.sqrt(det_val)
      return meanc
  
  
  def variance_clicks(cov, hbar=2):
      r"""Calculates the variance of the total number of clicks when a zero-mean gaussian state
      is measured using threshold detectors.
  
      Args
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering
          hbar (float): the value of :math:`\hbar` in the commutation relation :math:`[\x,\p]=i\hbar`
  
      Returns
          float: variance in the total number of clicks
      """
      n, _ = cov.shape
      means = np.zeros([n])
      nmodes = n // 2
      Q = Qmat(cov, hbar=hbar)
      vac_probs = np.array(
          [
              np.real(Q[i, i] * Q[i + nmodes, i + nmodes] - Q[i + nmodes, i] * Q[i, i + nmodes])
              for i in range(nmodes)
          ]
      )
      vac_probs = np.sqrt(vac_probs)
      vac_probs = 1 / vac_probs
      term1 = np.sum(vac_probs * (1 - vac_probs))
      term2 = 0
      for i in range(nmodes):
          for j in range(i):
              _, Qij = reduced_gaussian(means, Q, [i, j])
              prob_vac_ij = np.linalg.det(Qij).real
              prob_vac_ij = 1.0 / np.sqrt(prob_vac_ij)
              term2 += prob_vac_ij - vac_probs[i] * vac_probs[j]
  
      return term1 + 2 * term2
  
  
  def _coeff_normal_ordered(m, k):
      r"""Returns the coefficients giving the expansion of a photon number power in terms of normal ordered power of creation
      and annihilation operators. The coefficient is given by :math:`\sum_{\mu=0}^k \frac{(-1)^{k-\mu} \mu^m}{\mu!(k-\mu)!}`.
  
      Args:
          m (int): power of the photon number operator, :math:`(a^\dagger a)^m `.
          k (int): power of the normal ordered term, :math:`a^{\dagger i} a^i`.
  
      Returns:
          (float): expansion coefficient
      """
  
      return sum(
          [
              (1 / (factorial(mu) * factorial(k - mu))) * ((-1) ** (k - mu) * (mu ** m))
              for mu in range(0, k + 1)
          ]
      )
  
  
  def photon_number_moment(mu, cov, indices, hbar=2):
      r"""Calculates the expectation value of product of powers of photon number operators of a Gaussian state.
      The powers are specified by a dictionary with modes as keys and powers as values.
  
      The calculation is performed by first writing any power of the photon number as
  
      :math:`(a^\dagger a)^m = \sum_{k=1}^m c_k a^{\dagger k} a^k`
  
      where the coefficients :math:`c_i` are provided by the function `_coeff_normal_ordered`.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          indices (dictionary): specification of the different modes and their power of their photon number
          hbar (float): value of hbar in the uncertainty relation.
  
      Returns:
          float: the expectation value of the photon number powers.
  
      """
      N = len(cov) // 2
      list_indices = [indices[key] for key in indices]
      modes = list(indices)
      # Find the expansion coefficients of all the different powers
      expansion_coeff = [
          [_coeff_normal_ordered(indices[key], i) for i in range(1, 1 + indices[key])]
          for key in indices
      ]
      values = [list(range(i)) for i in list_indices]
      net_sum = 0.0
      # Construct the product of each possible term appearing in the normal ordered expansion
      for item in product(*values):
          rpt = [0] * N
          for i, key in enumerate(modes):
              rpt[key] = item[i] + 1
          rpt = rpt + rpt
          prod_coeff = np.prod([expansion_coeff[i][coeff] for i, coeff in enumerate(item)])
          net_sum += prod_coeff * s_ordered_expectation(mu, cov, rpt, s=1, hbar=hbar)
      return np.real_if_close(net_sum)
  
  
  def partition(collection):
      """Generate all set partitions of a collection.
  
      Taken from: https://stackoverflow.com/a/30134039
  
      Args:
          collection (sequence): set to find partitions of
  
      Yields:
          list[list]: set partition of collection
      """
      if len(collection) == 1:
          yield [collection]
          return
  
      first = collection[0]
      for smaller in partition(collection[1:]):
          for n, subset in enumerate(smaller):
              yield smaller[:n] + [[first] + subset] + smaller[n + 1 :]
          yield [[first]] + smaller
  
  
  def _list_to_freq_dict(words):
      """Convert between a list which of "words" and a dictionary
      which shows how many times each word appears in word
  
      Args:
          words (list): list of words
      Returns:
          dict : how many times a word appears. key is word, value is multiplicity
      """
      return {i: words.count(i) for i in set(words)}
  
  
  def photon_number_cumulant(mu, cov, modes, hbar=2):
      r"""Calculates the photon-number cumulant of the modes in the Gaussian state.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          modes (list or array): list of modes. Note that it can have repetitions.
          hbar (float): value of hbar in the uncertainty relation.
  
      Returns:
          (float): the cumulant
      """
  
      modes = list(modes)  # turns modes from array to list if passed in as array
  
      kappa = 0
      for pi in partition(modes):
          size = len(pi)
          term = factorial(size - 1) * (-1) ** (size - 1)
          for B in pi:
              indices = _list_to_freq_dict(B)
              term *= photon_number_moment(mu, cov, indices, hbar=hbar)
          kappa += term
  
      return kappa
  
  
  def click_cumulant(mu, cov, modes, hbar=2):
      r"""Calculates the click cumulant of the modes in the Gaussian state.
  
      Args:
          mu (array): length-:math:`2N` means vector in xp-ordering.
          cov (array): :math:`2N\times 2N` covariance matrix in xp-ordering.
          modes (list or array): list of modes.
          hbar (float): value of hbar in the uncertainty relation.
  
      Returns:
          (float): the cumulant
      """
  
      modes = list(modes)  # turns modes from array to list if passed in as array
      kappa = 0
      for pi in partition(modes):
          size = len(pi)
          term = factorial(size - 1) * (-1) ** (size - 1)
          for B in pi:
              B = list(set(B))  # remove repetitions
              pattern = np.ones_like(B)
              mu_red, cov_red = reduced_gaussian(mu, cov, B)
              summand = threshold_detection_prob(mu_red, cov_red, pattern, hbar=hbar)
              term *= summand
          kappa += term
  
      return kappa
  