content: """
  This module contains functions for calculating current and current noise using
  the counting statistics formalism.
  """
  __all__ = ['countstat_current', 'countstat_current_noise']
  
  import numpy as np
  
  from itertools import product
  from ..core import (
      sprepost, spre, qeye_like, Qobj, operator_to_vector, CoreOptions
  )
  from ..core import data as _data
  from .steadystate import pseudo_inverse, steadystate
  
  
  def countstat_current(L, c_ops=None, rhoss=None, J_ops=None):
      """
      Calculate the current corresponding to a system Liouvillian ``L`` and a
      list of current collapse operators ``c_ops`` or current superoperators
      ``J_ops``.
  
      Parameters
      ----------
  
      L : :class:`.Qobj`
          Qobj representing the system Liouvillian.
  
      c_ops : array / list (optional)
          List of current collapse operators. Required if either ``rhoss`` or
          ``J_ops`` is not given.
  
      rhoss : :class:`.Qobj` (optional)
          The steadystate density matrix for the given system Liouvillian ``L``
          and collapse operators. If not given, it defaults to
          ``steadystate(L, c_ops)``.
  
      J_ops : array / list (optional)
          List of current superoperators. If not given, they default to
          ``sprepost(c, c.dag())`` for each ``c`` from ``c_ops``.
  
      Returns
      --------
      I : array
          The currents ``I`` corresponding to each current collapse operator
          ``J_ops`` (or to each ``c_ops`` if ``J_ops`` was not given).
      """
  
      if J_ops is None:
          if c_ops is None:
              raise ValueError("c_ops must be given if J_ops is not")
          J_ops = [sprepost(c, c.dag()) for c in c_ops]
  
      if rhoss is None:
          if c_ops is None:
              raise ValueError("c_ops must be given if rhoss is not")
          rhoss = steadystate(L, c_ops)
  
      rhoss_vec = _data.column_stack(rhoss.data.copy())
  
      N = len(J_ops)
      current = np.zeros(N)
  
      for i, Ji in enumerate(J_ops):
          current[i] = _data.expect_super(Ji.data, rhoss_vec).real
  
      return current
  
  
  def _solve(A, V):
      try:
          return _data.solve(A, V)
      except ValueError:
          return _data.solve(A, V, "lstsq")
  
  
  def _noise_direct(L, wlist, rhoss, J_ops, I_ops):
      J_ops = [op.data for op in J_ops]
      I_ops = [op.data for op in I_ops]
      rhoss_vec = operator_to_vector(rhoss).data
  
      N_j_ops = len(J_ops)
      current = np.zeros(N_j_ops)
      noise = np.zeros((N_j_ops, N_j_ops, len(wlist)))
      skewness = np.zeros((N_j_ops, len(wlist)))
  
      tr_op = qeye_like(rhoss)
      tr_op_vec = operator_to_vector(tr_op)
  
      Pop = _data.kron(rhoss_vec, tr_op_vec.data.transpose())
      Iop = _data.identity_like(L.data)
      Q = _data.sub(Iop, Pop)
  
      QJ_ops = [
          _data.matmul(Q, _data.matmul(op, rhoss_vec))
          for op in J_ops
      ]
      QI_ops = [
          _data.matmul(Q, _data.matmul(op, rhoss_vec))
          for op in I_ops
      ]
  
      QIPI_ops = [
          _data.matmul(
              Q,
              _data.matmul(
                  op, _data.matmul(Pop, _data.matmul(op, rhoss_vec))
              )
          )
          for op in I_ops
      ]
  
      for k, w in enumerate(wlist):
          if w != 0.0:
              L_temp = 1.0j * w * spre(tr_op) + L
          else:
              # At zero frequency some solvers fail for small systems.
              # Adding a small finite frequency of order 1e-15
              # helps prevent the solvers from throwing an exception.
              with CoreOptions(auto_tidyup=False):
                  L_temp = 1e-15j * spre(tr_op) + L
  
          XI_rho = [
              _solve(L_temp.data, op) for op in QI_ops
          ]
          XJ_rho = [
              _solve(L_temp.data, op) for op in QJ_ops
          ]
  
          RIRI = [
              _solve(
                  L_temp.data,
                  _data.matmul(Q, _data.matmul(opI, _data.matmul(Q, op)))
              )
              for op, opI in zip(XI_rho, I_ops)
          ]
  
          RRIPI = [
              _solve(
                  L_temp.data,
                  _data.matmul(
                      Q, _data.matmul(Q, _solve(L_temp.data, op))
                  )
              )
              for op in QIPI_ops
          ]
  
          for i, j in product(range(N_j_ops), repeat=2):
              if i == j:
                  current[i] = _data.expect_super(I_ops[i], rhoss_vec).real
                  noise[j, i, k] = _data.expect_super(J_ops[i], rhoss_vec).real
                  skewness[i, k] = _data.expect_super(I_ops[i], rhoss_vec).real
  
                  skewness[i, k] -= (
                      3 * (
                          _data.expect_super(
                              _data.matmul(J_ops[i], Q), XI_rho[i]
                          )
                          + _data.expect_super(
                              _data.matmul(I_ops[i], Q), XJ_rho[i]
                          )
                      ).real
                  )
  
                  skewness[i, k] -= (
                      6 * (
                          _data.expect_super(
                              _data.matmul(I_ops[i], Q), RRIPI[i]
                          )
                          - _data.expect_super(
                              _data.matmul(I_ops[i], Q), RIRI[i]
                          )
                      ).real
                  )
  
              noise[j, i, k] -= (
                  _data.expect_super(
                      _data.matmul(I_ops[j], Q), XI_rho[i]
                  )
                  + _data.expect_super(
                      _data.matmul(I_ops[i], Q), XI_rho[j]
                  )
              ).real
  
      return current, noise, skewness
  
  
  def _noise_pseudoinv(L, wlist, rhoss, J_ops, I_ops, sparse, method):
      N_j_ops = len(J_ops)
      current = np.zeros(N_j_ops)
      noise = np.zeros((N_j_ops, N_j_ops, len(wlist)))
      skewness = np.zeros((N_j_ops, len(wlist)))
  
      rhoss_vec = operator_to_vector(rhoss)
  
      dtype = type(L.data)
  
      tr_op = qeye_like(rhoss)
      tr_op_vec = operator_to_vector(tr_op)
  
      P = _data.kron(rhoss_vec.data, tr_op_vec.data.transpose(), dtype=dtype)
  
      Pop_new = Qobj(P, dims=L._dims)
      for k, w in enumerate(wlist):
          R = pseudo_inverse(L, rhoss=rhoss, w=w, sparse=sparse, method=method)
          for i, j in product(range(N_j_ops), repeat=2):
              if i == j:
                  current[i] = I_ops[i](rhoss).tr().real
                  noise[i, j, k] = J_ops[i](rhoss).tr().real
                  skewness[i, k] = I_ops[i](rhoss).tr().real
                  skewness[i, k] -= (
                      3 * (
                          (I_ops[i] @ R @ J_ops[i])(rhoss).tr().real +
                          (J_ops[i] @ R @ I_ops[i])(rhoss).tr().real
                      )
                  )
  
                  skewness[i, k] -= (
                      6 * (
                          (I_ops[i] @ R @ (
                              R @ I_ops[i] @ Pop_new - I_ops[i] @ R
                          ) @ I_ops[i])(rhoss).tr().real
                      )
                  )
              op = I_ops[i] @ R @ I_ops[j] + I_ops[j] @ R @ I_ops[i]
              noise[i, j, k] -= op(rhoss).tr().real
  
      return current, noise, skewness
  
  
  def countstat_current_noise(L, c_ops, wlist=None, rhoss=None, J_ops=None,
                              I_ops=None, sparse=True, method='direct'):
      """
      Compute the cross-current noise spectrum for a list of collapse operators
      `c_ops` corresponding to monitored currents, given the system
      Liouvillian `L`. The current collapse operators `c_ops` should be part
      of the dissipative processes in `L`, but the `c_ops` given here does not
      necessarily need to be all collapse operators contributing to dissipation
      in the Liouvillian. Optionally, the steadystate density matrix `rhoss`,
      the net current operator `I_ops` and the total current operators `J_ops`
      corresponding to the current collapse operators `c_ops` can also be
      specified. If either of `rhoss` or `J_ops` are omitted, they will be
      computed internally. 'wlist' is an optional list of frequencies
      at which to evaluate the noise spectrum.
  
      Parameters
      ----------
  
      L : :class:`.Qobj`
          Qobj representing the system Liouvillian.
  
      c_ops : array / list
          List of current collapse operators.
  
      rhoss : :class:`.Qobj` (optional)
          The steadystate density matrix corresponding the system Liouvillian
          `L`.
  
      wlist : array / list (optional)
          List of frequencies at which to evaluate (if none are given, evaluates
          at zero frequency)
  
      J_ops : array / list (optional)
          List of total current superoperators. Where J_ops is defined as the
          sum of currents across a junction J_ops ≡ I+ + I-. If not provided,
          it will be calculated using the collapse operators.
  
      I_ops : array / list (optional)
          List of net current superoperators. Where I_ops is defined as the
          difference of currents across a junction I_ops ≡ I+ - I-.
          If not provided, it will be assumed that I- = 0 and
          set I_ops = J_ops.
  
      sparse : bool [True]
          Flag that indicates whether to use sparse or dense matrix methods when
          computing the pseudo inverse. Default is false, as sparse solvers
          can fail for small systems. For larger systems the sparse solvers
          are recommended.
  
      method : str, ['direct']
          Method used to compute the noise. The default, 'direct' with
          ``sparse=True``, compute the noise directly. It is the recommended
          method for larger systems. Otherwise, the pseudo inverse is computed
          using the given method. Pseudo inverse supports 'splu' and 'spilu' for
          sparse matrices and 'direct', 'scipy' and 'numpy' methods for
          ``sparse=False``.
  
      .. note::
          The algoryth is described in page 67 of "Electrons in nanostructures"
          C. Flindt, PhD Thesis, available online:
          https://orbit.dtu.dk/en/publications/electrons-in-nanostructures-coherent-manipulation-and-counting-st
  
      Returns
      --------
      current : tuple of arrays
          The first cumulant, representing the mean current associated with each
          current superoperator `I_ops`. It quantifies the average charge
          transfer rate across the system.
  
      noise : tuple of arrays
          The second cumulant, representing the current noise, which measures
          fluctuations around the mean current. This term supports
          cross-correlations between different total current
          superoperators `J_ops`.
  
      skewness : tuple of arrays
          The third cumulant, representing the asymmetry of current fluctuations.
          It describes whether deviations from the mean current are more
          likely to be positive or negative. Unlike noise, skewness only
          accounts for diagonal (self-correlated) terms and does not
          include cross-correlations between different `J_ops`.
      """
      if rhoss is None:
          rhoss = steadystate(L, c_ops)
  
      if J_ops is None:
          J_ops = [sprepost(c, c.dag()) for c in c_ops]
  
      if I_ops is None:
          I_ops = J_ops
  
      if wlist is None:
          wlist = [0.]
  
      if sparse and method == 'direct':
          # rhoss_vec = operator_to_vector(rhoss).data
          current, noise, skewness = _noise_direct(L, wlist, rhoss, J_ops, I_ops)
      else:
          current, noise, skewness = _noise_pseudoinv(
              L, wlist, rhoss, J_ops, I_ops, sparse, method)
  
      return current, noise, skewness
  