text: #  Author: Artem Pulkin
  # flake8: noqa
  """
  This and other `_slow` modules implement the time-dependent Hartree-Fock procedure. The primary performance drawback is
  that, unlike other 'fast' routines with an implicit construction of the eigenvalue problem, these modules construct
  TDHF matrices explicitly via an AO-MO transformation, i.e. with a O(N^5) complexity scaling. As a result, regular
  `numpy.linalg.eig` can be used to retrieve TDHF roots in a reliable fashion without any issues related to the Davidson
  procedure. Several variants of TDHF are available:
  
   * `pyscf.tdscf.rhf_slow`: the molecular implementation;
   * `pyscf.pbc.tdscf.rhf_slow`: PBC (periodic boundary condition) implementation for RHF objects of `pyscf.pbc.scf`
     modules;
   * `pyscf.pbc.tdscf.krhf_slow_supercell`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
     an arbitrary number of k-points but has a overhead due to an effective construction of a supercell;
   * (this module) `pyscf.pbc.tdscf.krhf_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krhf`
     module. Despite its name, it accepts KRHF objects with an arbitrary number of k-points but finds only few TDHF roots
     corresponding to collective oscillations without momentum transfer;
   * `pyscf.pbc.tdscf.krhf_slow`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
     an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
  """
  
  from pyscf.pbc.tdscf import krhf_slow_supercell as td
  from pyscf.tdscf import rhf_slow
  
  import numpy
  
  
  # Convention for these modules:
  # * PhysERI, PhysERI4, PhysERI8 are 2-electron integral routines computed directly (for debug purposes), with a 4-fold
  #   symmetry and with an 8-fold symmetry
  # * vector_to_amplitudes reshapes and normalizes the solution
  # * TDRHF provides a container
  
  
  class PhysERI(td.PhysERI):
  
      def __init__(self, model, frozen=None):
          """
          The TDHF ERI implementation performing a full transformation of integrals to Bloch functions. No symmetries are
          employed in this class. Only a subset of transformed ERI is returned, corresponding to oscillations without a
          momentum transfer. The k-points of the returned ERIs come in two pairs :math:`(k_1 k_1 | k_2 k_2)`. As a result,
          one of the diagonal blocks of the full TDHF matrix can be reconstructed. For other blocks, look into `PhysERI`
          class of this momdule.
  
          Args:
              model (KRHF): the base model;
              frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals for all
              k-points or multiple lists of frozen orbitals for each k-point;
          """
          super(PhysERI, self).__init__(model, frozen=frozen)
  
      def tdhf_diag(self):
          """
          Retrieves the merged diagonal block with equal pairs of k-indexes (k, k) only.
  
          Returns:
              The diagonal block.
          """
          return super(PhysERI, self).tdhf_diag(pairs=((i, i) for i in range(len(self.model.kpts))))
  
      def __calc_block__(self, item, k):
          if k in self.__full_eri_k__:
              return super(PhysERI, self).__calc_block__(item, k)
          else:
              raise RuntimeError("The block k = {:d} + {:d} - {:d} - {:d} does not conserve momentum".format(*k))
  
      def eri_mknj(self, item):
          """
          Retrieves the merged ERI block using 'mknj' notation with pairs of k-indexes (k1, k1, k2, k2).
          Args:
              item (str): a 4-character string of 'mknj' letters;
  
          Returns:
              The corresponding block of ERI (phys notation).
          """
          return super(PhysERI, self).eri_mknj(
              item,
              pairs_row=((i, i) for i in range(len(self.model.kpts))),
              pairs_column=((i, i) for i in range(len(self.model.kpts))),
          )
  
  
  class PhysERI4(PhysERI):
  
      def __init__(self, model, frozen=None):
          """
          The TDHF ERI implementation performing partial transformations of integrals to Bloch functions. A 4-fold
          symmetry of complex-valued functions is employed in this class. Only a subset of transformed ERI is returned,
          corresponding to oscillations without a momentum transfer. The k-points of the returned ERIs come in two pairs
          :math:`(k_1 k_1 | k_2 k_2)`. As a result, one of the diagonal blocks of the full TDHF matrix can be
          reconstructed. For other blocks, look into `PhysERI` class of this momdule.
  
          Args:
              model (KRHF): the base model;
              frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals for all
              k-points or multiple lists of frozen orbitals for each k-point;
          """
          super(PhysERI4, self).__init__(model, frozen=frozen)
  
      symmetries = [
          ((0, 1, 2, 3), False),
          ((1, 0, 3, 2), False),
          ((2, 3, 0, 1), True),
          ((3, 2, 1, 0), True),
      ]
  
      def __calc_block__(self, item, k):
          if k in self.__full_eri_k__:
              return td.PhysERI4.__calc_block__.im_func(self, item, k)
          else:
              raise RuntimeError("The block k = {:d} + {:d} - {:d} - {:d} does not conserve momentum".format(*k))
  
  
  class PhysERI8(PhysERI4):
      def __init__(self, model, frozen=None):
          """
          The TDHF ERI implementation performing partial transformations of integrals to Bloch functions. An 8-fold
          symmetry of real-valued functions is employed in this class. Only a subset of transformed ERI is returned,
          corresponding to oscillations without a momentum transfer. The k-points of the returned ERIs come in two pairs
          :math:`(k_1 k_1 | k_2 k_2)`. As a result, one of the diagonal blocks of the full TDHF matrix can be
          reconstructed. For other blocks, look into `PhysERI` class of this momdule.
  
          Args:
              model (KRHF): the base model;
              frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals for all
              k-points or multiple lists of frozen orbitals for each k-point;
          """
          super(PhysERI8, self).__init__(model, frozen=frozen)
  
      symmetries = [
          ((0, 1, 2, 3), False),
          ((1, 0, 3, 2), False),
          ((2, 3, 0, 1), False),
          ((3, 2, 1, 0), False),
  
          ((2, 1, 0, 3), False),
          ((3, 0, 1, 2), False),
          ((0, 3, 2, 1), False),
          ((1, 2, 3, 0), False),
      ]
  
  
  def vector_to_amplitudes(vectors, nocc, nmo):
      """
      Transforms (reshapes) and normalizes vectors into amplitudes.
      Args:
          vectors (numpy.ndarray): raw eigenvectors to transform;
          nocc (tuple): numbers of occupied orbitals;
          nmo (int): the total number of orbitals per k-point;
  
      Returns:
          Amplitudes with the following shape: (# of roots, 2 (x or y), # of kpts, # of kpts, # of occupied orbitals,
          # of virtual orbitals).
      """
      if not all(i == nocc[0] for i in nocc):
          raise NotImplementedError("Varying occupation numbers are not implemented yet")
      nk = len(nocc)
      nocc = nocc[0]
      if not all(i == nmo[0] for i in nmo):
          raise NotImplementedError("Varying AO spaces are not implemented yet")
      nmo = nmo[0]
      vectors = numpy.asanyarray(vectors)
      # Compared to krhf_slow_supercell, only one k-point index is present here. The second index corresponds to
      # momentum transfer and is integrated out
      vectors = vectors.reshape(2, nk, nocc, nmo-nocc, vectors.shape[1])
      norm = (abs(vectors) ** 2).sum(axis=(1, 2, 3))
      norm = 2 * (norm[0] - norm[1])
      vectors /= norm ** .5
      return vectors.transpose(4, 0, 1, 2, 3)
  
  
  class TDRHF(rhf_slow.TDRHF):
      eri4 = PhysERI4
      eri8 = PhysERI8
      v2a = staticmethod(vector_to_amplitudes)
  