content: """
  Photon scattering in quantum optical systems
  
  This module includes a collection of functions for numerically computing photon
  scattering in driven arbitrary systems coupled to some configuration of output
  waveguides. The implementation of these functions closely follows the
  mathematical treatment given in K.A. Fischer, et. al., Scattering of Coherent
  Pulses from Quantum Optical Systems (2017, arXiv:1710.02875).
  """
  # Author:  Ben Bartlett
  # Contact: benbartlett@stanford.edu
  
  import numpy as np
  from scipy.integrate import trapezoid
  from itertools import product, combinations_with_replacement
  from ..core import basis, tensor, zero_ket, Qobj, QobjEvo
  from .propagator import propagator, Propagator
  
  __all__ = ['temporal_basis_vector',
             'temporal_scattered_state',
             'scattering_probability']
  
  
  def set_partition(collection, num_sets):
      """
      Enumerate all ways of partitioning collection into num_sets different
      lists, e.g. :
      list(set_partition([1,2], 2))
      >>> [[[1, 2], []], [[1], [2]], [[2], [1]], [[], [1, 2]]]
  
      Parameters
      ----------
      collection : iterable
          Collection to generate a set partition of.
      num_sets : int
          Number of sets to partition collection into.
  
      Returns
      -------
      partition : iterable
          The partitioning of collection into num_sets sets.
      """
      for partitioning in product(range(num_sets), repeat=len(collection)):
          partition = [[] for _ in range(num_sets)]
          for i, set_index in enumerate(partitioning):
              partition[set_index].append(collection[i])
          yield tuple(tuple(indices) for indices in partition)
  
  
  def photon_scattering_amplitude(propagator, c_ops, tlist, taus, psi, psit):
      """
      Compute the scattering amplitude for a system emitting into multiple
      waveguides.
  
      Parameters
      ----------
      propagator : :class:`.Propagator`
          Propagator
      c_ops : list
          list of collapse operators for each waveguide; these are assumed to
          include spontaneous decay rates, e.g.
          :math:`\\sigma = \\sqrt \\gamma \\cdot a`
      tlist : array_like
          List of times starting at the beginning and ending at the end of the
          evolution.
      taus : list-like
          List of (list of emission times) for each waveguide.
      psi : Qobj
          State at the start of the evolution
      psit : Qobj
          State at the end of the evolution.
      """
      # Extract the full list of taus
      tau_collapse = []
      for i, tau_wg in enumerate(taus):
          for t in tau_wg:
              tau_collapse.append((t, i))
      tau_collapse.sort(key=lambda tup: tup[0])  # sort tau_collapse by time
  
      tq = tlist[0]
      # Compute Prod Ueff(tq, tq-1)
      for tau in tau_collapse:
          tprev = tq
          tq, q = tau
          psi = c_ops[q] * propagator(tq, tprev) * psi
  
      psi = propagator(tlist[-1], tq) * psi
      return psit.overlap(psi)
  
  
  def _temporal_basis_idx(waveguide_emission_indices, n_time_bins):
      """
      Generate a the global index for the excitation.
      """
      idx = []
      for i, wg_indices in enumerate(waveguide_emission_indices):
          for index in wg_indices:
              idx += [i * n_time_bins + index]
      idx = idx or [0]
      return tuple(idx)
  
  
  def _temporal_basis_dims(waveguide_emission_indices, n_time_bins,
                           n_emissions=None):
      """
      Return the dims of the ``temporal_basis_vector``.
      """
      # TODO: Review n_emissions: change the number of dims but the equivalent
      # does not exist in _temporal_basis_idx
      num_col = len(waveguide_emission_indices)
      if n_emissions is None:
          n_emissions = sum(
              [len(waveguide_indices)
               for waveguide_indices in waveguide_emission_indices])
      n_emissions = n_emissions or 1
      return [num_col * n_time_bins] * n_emissions
  
  
  def temporal_basis_vector(waveguide_emission_indices, n_time_bins):
      """
      Generate a temporal basis vector for emissions at specified time bins into
      specified waveguides.
  
      Parameters
      ----------
      waveguide_emission_indices : list or tuple
          List of indices where photon emission occurs for each waveguide,
          e.g. [[t1_wg1], [t1_wg2, t2_wg2], [], [t1_wg4, t2_wg4, t3_wg4]].
      n_time_bins : int
          Number of time bins; the range over which each index can vary.
  
      Returns
      -------
      temporal_basis_vector : :class:`.Qobj`
          A basis vector representing photon scattering at the specified indices.
          If there are W waveguides, T times, and N photon emissions, then the
          basis vector has dimensionality (W*T)^N.
      """
      idx = _temporal_basis_idx(waveguide_emission_indices, n_time_bins)
      dims = _temporal_basis_dims(waveguide_emission_indices, n_time_bins, None)
      return basis(dims, list(idx))
  
  
  def _temporal_scattered_matrix(H, psi0, n_emissions, c_ops, tlist,
                                 system_zero_state=None,
                                 construct_effective_hamiltonian=True):
      """
      Compute the scattered n-photon state as an ndarray.
      """
      T = len(tlist)
      W = len(c_ops)
      em_dims = max(n_emissions, 1)
      phi_n = np.zeros([W * T] * em_dims, dtype=complex)
  
      if construct_effective_hamiltonian:
          # Construct an effective Hamiltonian from system hamiltonian and c_ops
          Heff = QobjEvo(H) - 1j / 2 * sum([op.dag() * op for op in c_ops])
      else:
          Heff = H
  
      evolver = Propagator(Heff, memoize=len(tlist))
  
      all_emission_indices = combinations_with_replacement(range(T), n_emissions)
  
      if system_zero_state is None:
          system_zero_state = psi0
  
      # Compute <omega_tau> for all combinations of tau
      for emission_indices in all_emission_indices:
          # Consider unique partitionings of emission times into waveguides
          partition = tuple(set(set_partition(emission_indices, W)))
          # Consider all possible partitionings of time bins by waveguide
          for indices in partition:
              taus = [[tlist[i] for i in wg_indices] for wg_indices in indices]
              phi_n_amp = photon_scattering_amplitude(
                  evolver, c_ops, tlist,
                  taus, psi0, system_zero_state
              )
              # Add scatter amplitude times temporal basis to overall state
              idx = _temporal_basis_idx(indices, T)
              phi_n[idx] = phi_n_amp
      return phi_n
  
  
  def temporal_scattered_state(H, psi0, n_emissions, c_ops, tlist,
                               system_zero_state=None,
                               construct_effective_hamiltonian=True):
      """
      Compute the scattered n-photon state projected onto the temporal basis.
  
      Parameters
      ----------
      H : :class:`.Qobj` or list
          System-waveguide(s) Hamiltonian or effective Hamiltonian in Qobj or
          list-callback format. If construct_effective_hamiltonian is not
          specified, an effective Hamiltonian is constructed from `H` and
          `c_ops`.
      psi0 : :class:`.Qobj`
          Initial state density matrix :math:`\\rho(t_0)` or state vector
          :math:`\\psi(t_0)`.
      n_emissions : int
          Number of photon emissions to calculate.
      c_ops : list
          List of collapse operators for each waveguide; these are assumed to
          include spontaneous decay rates, e.g.
          :math:`\\sigma = \\sqrt \\gamma \\cdot a`
      tlist : array_like
          List of times for :math:`\\tau_i`. tlist should contain 0 and exceed
          the pulse duration / temporal region of interest.
      system_zero_state : :class:`.Qobj`, optional
          State representing zero excitations in the system. Defaults to
          :math:`\\psi(t_0)`
      construct_effective_hamiltonian : bool, default: True
          Whether an effective Hamiltonian should be constructed from H and
          c_ops:
          :math:`H_{eff} = H - \\frac{i}{2} \\sum_n \\sigma_n^\\dagger \\sigma_n`
          Default: True.
  
      Returns
      -------
      phi_n : :class:`.Qobj`
          The scattered bath state projected onto the temporal basis given by
          tlist. If there are W waveguides, T times, and N photon emissions, then
          the state is a tensor product state with dimensionality T^(W*N).
      """
      T = len(tlist)
      W = len(c_ops)
      em_dims = max(n_emissions, 1)
      phi_n = _temporal_scattered_matrix(
          H, psi0, n_emissions, c_ops, tlist,
          system_zero_state, construct_effective_hamiltonian
      )
      return Qobj(phi_n.ravel(), dims=[[W * T] * em_dims, [1] * em_dims])
  
  
  def scattering_probability(H, psi0, n_emissions, c_ops, tlist,
                             system_zero_state=None,
                             construct_effective_hamiltonian=True):
      """
      Compute the integrated probability of scattering n photons in an arbitrary
      system. This function accepts a nonlinearly spaced array of times.
  
      Parameters
      ----------
      H : :class:`.Qobj` or list
          System-waveguide(s) Hamiltonian or effective Hamiltonian in Qobj or
          list-callback format. If construct_effective_hamiltonian is not
          specified, an effective Hamiltonian is constructed from H and
          `c_ops`.
      psi0 : :class:`.Qobj`
          Initial state density matrix :math:`\\rho(t_0)` or state vector
          :math:`\\psi(t_0)`.
      n_emissions : int
          Number of photons emitted by the system (into any combination of
          waveguides).
      c_ops : list
          List of collapse operators for each waveguide; these are assumed to
          include spontaneous decay rates, e.g.
          :math:`\\sigma = \\sqrt \\gamma \\cdot a`.
      tlist : array_like
          List of times for :math:`\\tau_i`. tlist should contain 0 and exceed
          the pulse duration / temporal region of interest; tlist need not be
          linearly spaced.
      system_zero_state : :class:`.Qobj`, optional
          State representing zero excitations in the system. Defaults to
          `basis(systemDims, 0)`.
      construct_effective_hamiltonian : bool, default: True
          Whether an effective Hamiltonian should be constructed from H and
          c_ops:
          :math:`H_{eff} = H - \\frac{i}{2} \\sum_n \\sigma_n^\\dagger \\sigma_n`
          Default: True.
  
      Returns
      -------
      scattering_prob : float
          The probability of scattering n photons from the system over the time
          range specified.
      """
      phi_n = _temporal_scattered_matrix(H, psi0, n_emissions, c_ops, tlist,
                                         system_zero_state,
                                         construct_effective_hamiltonian)
      T = len(tlist)
      W = len(c_ops)
  
      # Compute <omega_tau> for all combinations of tau
      all_emission_indices = combinations_with_replacement(range(T), n_emissions)
      probs = np.zeros([T] * n_emissions)
  
      # Project scattered state onto temporal basis
      for emit_indices in all_emission_indices:
          # Consider unique emission time partitionings
          partition = tuple(set(set_partition(emit_indices, W)))
          # wg_indices_list = list(set_partition(indices, W))
          for wg_indices in partition:
              idx = _temporal_basis_idx(wg_indices, T)
              amplitude = phi_n[idx]
              probs[emit_indices] += np.real(amplitude.conjugate() * amplitude)
  
      # Iteratively integrate to obtain single value
      while probs.shape != ():
          probs = trapezoid(probs, x=tlist)
      return np.abs(probs)
  