text: import numpy as np
  from numpy import matlib
  import scipy.optimize
  from scipy import sparse
  import logging
  from typing import *
  from sklearn.neighbors import NearestNeighbors
  from .speedboosted import _colDeltaCor, _colDeltaCorLog10, _colDeltaCorSqrt
  from .speedboosted import _colDeltaCorpartial, _colDeltaCorLog10partial, _colDeltaCorSqrtpartial
  
  
  def colDeltaCor(emat: np.ndarray, dmat: np.ndarray, threads: int=None) -> np.ndarray:
      """Calculate the correlation between the displacement (d[:,i])
      and the difference between a cell and every other (e - e[:, i])
      
      Parallel cython+OpenMP implemetation
  
      Arguments
      ---------
      emat: np.ndarray (ngenes, ncells)
          gene expression matrix
      dmat: np.ndarray (ngenes, ncells)
          gene velocity/displacement matrix
      threads: int
          number of parallel threads to use
      """
      import multiprocessing
      if threads is None:
          num_threads = int(multiprocessing.cpu_count() / 2)
      else:
          num_threads = max(threads, multiprocessing.cpu_count())
      out = np.zeros((emat.shape[1], emat.shape[1]))
      _colDeltaCor(emat, dmat, out, num_threads)
      return out
  
  def colDeltaBool(emat: np.ndarray, dmat: np.ndarray) -> np.ndarray:    
      ncell = emat.shape[1]
      ngene = emat.shape[0]
      corrcoef = np.zeros((ncell,ncell))
      for i_c in range(ncell):
          v_bool = np.matlib.repmat(np.sign(dmat[:,i_c]),ncell,1)
          displ_bool= np.sign(emat.T-np.matlib.repmat(emat[:,i_c],ncell,1))
          corrcoef[i_c,:] = np.sum(v_bool==displ_bool,1)/ngene
      return corrcoef
  
  def colDeltaCorpartial(emat: np.ndarray, dmat: np.ndarray, ixs: np.ndarray, threads: int=None) -> np.ndarray:
      """Calculate the correlation between the displacement (d[:,i])
      and the difference between a cell and every other (e - e[:, i])
      
      Parallel cython+OpenMP implemetation
  
      Arguments
      ---------
      emat: np.ndarray (ngenes, ncells)
          gene expression matrix
      dmat: np.ndarray (ngenes, ncells)
          gene velocity/displacement matrix
      ixs: the neighborhood matrix (ncells, nneighbours)
          ixs[i, k] is the kth neighbour to the cell i
      threads: int
          number of parallel threads to use
      """
      import multiprocessing
      if threads is None:
          num_threads = int(multiprocessing.cpu_count() / 2)
      else:
          num_threads = max(threads, multiprocessing.cpu_count())
      out = np.zeros((emat.shape[1], emat.shape[1]))
      emat = np.require(emat, requirements="C")
      ixs = np.require(ixs, requirements="C").astype(np.intp)
      _colDeltaCorpartial(emat, dmat, out, ixs, num_threads)
      return out
  
  
  def colDeltaCorLog10(emat: np.ndarray, dmat: np.ndarray, threads: int=None, psc: float=1.0) -> np.ndarray:
      """Calculate the correlation between the displacement (d[:,i])
      and the difference between a cell and every other (e - e[:, i])
      
      Parallel cython+OpenMP implemetation
  
      Arguments
      ---------
      emat: np.ndarray (ngenes, ncells)
          gene expression matrix
      dmat: np.ndarray (ngenes, ncells)
          gene velocity/displacement matrix
      threads: int
          number of parallel threads to use
      """
      import multiprocessing
      if threads is None:
          num_threads = int(multiprocessing.cpu_count() / 2)
      else:
          num_threads = max(threads, multiprocessing.cpu_count())
      out = np.zeros((emat.shape[1], emat.shape[1]))
      _colDeltaCorLog10(emat, dmat, out, num_threads, psc)
      return out
  
  
  def colDeltaCorLog10partial(emat: np.ndarray, dmat: np.ndarray, ixs: np.ndarray, threads: int=None, psc: float=1.0) -> np.ndarray:
      """Calculate the correlation between the displacement (d[:,i])
      and the difference between a cell and every other (e - e[:, i])
      
      Parallel cython+OpenMP implemetation
  
      Arguments
      ---------
      emat: np.ndarray (ngenes, ncells)
          gene expression matrix
      dmat: np.ndarray (ngenes, ncells)
          gene velocity/displacement matrix
      ixs: the neighborhood matrix (ncells, nneighbours)
          ixs[i, k] is the kth neighbour to the cell i
      threads: int
          number of parallel threads to use
      """
      import multiprocessing
      if threads is None:
          num_threads = int(multiprocessing.cpu_count() / 2)
      else:
          num_threads = max(threads, multiprocessing.cpu_count())
      out = np.zeros((emat.shape[1], emat.shape[1]))
      emat = np.require(emat, requirements="C")
      ixs = np.require(ixs, requirements="C").astype(np.intp)
      _colDeltaCorLog10partial(emat, dmat, out, ixs, num_threads, psc)
      return out
  
  
  def colDeltaCorSqrt(emat: np.ndarray, dmat: np.ndarray, threads: int=None, psc: float=0.0) -> np.ndarray:
      """Calculate the correlation between the displacement (d[:,i])
      and the difference between a cell and every other (e - e[:, i])
      
      Parallel cython+OpenMP implemetation
  
      Arguments
      ---------
      emat: np.ndarray (ngenes, ncells)
          gene expression matrix
      dmat: np.ndarray (ngenes, ncells)
          gene velocity/displacement matrix
      threads: int
          number of parallel threads to use
      """
      import multiprocessing
      if threads is None:
          num_threads = int(multiprocessing.cpu_count() / 2)
      else:
          num_threads = max(threads, multiprocessing.cpu_count())
      out = np.zeros((emat.shape[1], emat.shape[1]))
      _colDeltaCorSqrt(emat, dmat, out, num_threads, psc)
      return out
  
  
  def colDeltaCorSqrtpartial(emat: np.ndarray, dmat: np.ndarray, ixs: np.ndarray, threads: int=None, psc: float=0.0) -> np.ndarray:
      """Calculate the correlation between the displacement (d[:,i])
      and the difference between a cell and every other (e - e[:, i])
      
      Parallel cython+OpenMP implemetation
  
      Arguments
      ---------
      emat: np.ndarray (ngenes, ncells)
          gene expression matrix
      dmat: np.ndarray (ngenes, ncells)
          gene velocity/displacement matrix
      ixs: the neighborhood matrix (ncells, nneighbours)
          ixs[i, k] is the kth neighbour to the cell i
      threads: int
          number of parallel threads to use
      """
      import multiprocessing
      if threads is None:
          num_threads = int(multiprocessing.cpu_count() / 2)
      else:
          num_threads = max(threads, multiprocessing.cpu_count())
      out = np.zeros((emat.shape[1], emat.shape[1]))
      emat = np.require(emat, requirements="C")
      ixs = np.require(ixs, requirements="C").astype(np.intp)
      _colDeltaCorSqrtpartial(emat, dmat, out, ixs, num_threads, psc)
      return out
  
  
  def _fit1_slope(y: np.ndarray, x: np.ndarray) -> float:
      """Simple function that fit a linear regression model without intercept
      """
      if not np.any(x):
          m = np.NAN  # It is definetelly not at steady state!!!
      elif not np.any(y):
          m = 0
      else:
          result, rnorm = scipy.optimize.nnls(x[:, None], y)  # Fastest but costrains result >= 0
          m = result[0]
          # Second fastest: m, _ = scipy.optimize.leastsq(lambda m: x*m - y, x0=(0,))
          # Third fastest: m = scipy.optimize.minimize_scalar(lambda m: np.sum((x*m - y)**2 )).x
          # Before I was doinf fastest: scipy.optimize.minimize_scalar(lambda m: np.sum((y - m * x)**2), bounds=(0, 3), method="bounded").x
          # Optionally one could clip m if high value make no sense
          # m = np.clip(m,0,3)
      return m
  
  
  def _fit1_slope_weighted(y: np.ndarray, x: np.ndarray, w: np.ndarray, limit_gamma: bool=False, bounds: Tuple[float, float]=(0, 20)) -> float:
      """Simple function that fit a weighted linear regression model without intercept
      """
      if not np.any(x):
          m = np.NAN  # It is definetelly not at steady state!!!
      elif not np.any(y):
          m = 0
      else:
          if limit_gamma:
              if np.median(y) > np.median(x):
                  high_x = x > np.percentile(x, 90)
                  up_gamma = np.percentile(y[high_x], 10) / np.median(x[high_x])
                  up_gamma = np.maximum(1.5, up_gamma)
              else:
                  up_gamma = 1.5  # Just a bit more than 1
              m = scipy.optimize.minimize_scalar(lambda m: np.sum(w * (x * m - y)**2), bounds=(1e-8, up_gamma), method="bounded").x
          else:
              m = scipy.optimize.minimize_scalar(lambda m: np.sum(w * (x * m - y)**2), bounds=bounds, method="bounded").x
      return m
  
  
  def _fit1_slope_weighted_offset(y: np.ndarray, x: np.ndarray, w: np.ndarray, fixperc_q: bool=False, limit_gamma: bool=False) -> Tuple[float, float]:
      """Function that fits a weighted linear regression model with intercept
      with some adhoc
      """
      if not np.any(x):
          m = (np.NAN, 0)  # It is definetelly not at steady state!!!
      elif not np.any(y):
          m = (0, 0)
      else:
          if fixperc_q:
              m1 = np.percentile(y[x <= np.percentile(x, 1)], 50)
              m0 = scipy.optimize.minimize_scalar(lambda m: np.sum(w * (x * m - y + m1)**2), bounds=(0, 20), method="bounded").x
              m = (m0, m1)
          else:
              # m, _ = scipy.optimize.leastsq(lambda m: np.sqrt(w) * (-y + x * m[0] + m[1]), x0=(0, 0))  # This is probably faster but it can have negative slope
              # NOTE: The up_gamma is to deal with cases where consistently y > x. Those should have positive velocity everywhere
              if limit_gamma:
                  if np.median(y) > np.median(x):
                      high_x = x > np.percentile(x, 90)
                      up_gamma = np.percentile(y[high_x], 10) / np.median(x[high_x])
                      up_gamma = np.maximum(1.5, up_gamma)
                  else:
                      up_gamma = 1.5  # Just a bit more than 1
              else:
                  up_gamma = 20
              up_q = 2 * np.sum(y * w) / np.sum(w)
              m = scipy.optimize.minimize(lambda m: np.sum(w * (-y + x * m[0] + m[1])**2),
                                          x0=(0.1, 1e-16), method="L-BFGS-B",
                                          bounds=[(1e-8, up_gamma), (0, up_q)]).x  # If speedup is needed either the gradient or numexpr could be used
      return m[0], m[1]
  
  
  def _fit1_slope_offset(y: np.ndarray, x: np.ndarray, fixperc_q: bool=False) -> Tuple[float, float]:
      """Simple function that fit a linear regression model with intercept
      """
      if not np.any(x):
          m = (np.NAN, 0)  # It is definetelly not at steady state!!!
      elif not np.any(y):
          m = (0, 0)
      else:
          # result, rnorm = scipy.optimize.nnls(x[:, None], y)  # Fastest but costrains result >= 0
          # m = result[0]
          if fixperc_q:
              m1 = np.percentile(y[x <= np.percentile(x, 1)], 50)
              m0 = scipy.optimize.minimize_scalar(lambda m: np.sum((x * m - y + m1)**2), bounds=(0, 20), method="bounded").x
              m = (m0, m1)
          else:
              m, _ = scipy.optimize.leastsq(lambda m: -y + x * m[0] + m[1], x0=(0, 0))
          # Third fastest: m = scipy.optimize.minimize_scalar(lambda m: np.sum((x*m - y)**2 )).x
          # Before I was doinf fastest: scipy.optimize.minimize_scalar(lambda m: np.sum((y - m * x)**2), bounds=(0, 3), method="bounded").x
          # Optionally one could clip m if high value make no sense
          # m = np.clip(m,0,3)
      return m[0], m[1]
  
  
  def fit_slope(Y: np.ndarray, X: np.ndarray) -> np.ndarray:
      """Loop through the genes and fits the slope
  
      Y: np.ndarray, shape=(genes, cells)
          the dependent variable (unspliced)
      X: np.ndarray, shape=(genes, cells)
          the independent variable (spliced)
      """
      # NOTE this could be easily parallelized
      slopes = np.fromiter((_fit1_slope(Y[i, :], X[i, :]) for i in range(Y.shape[0])),
                           dtype="float32",
                           count=Y.shape[0])
      return slopes
  
  
  def fit_slope_offset(Y: np.ndarray, X: np.ndarray, fixperc_q: bool=False) -> Tuple[np.ndarray, np.ndarray]:
      """Loop through the genes and fits the slope
  
      Y: np.ndarray, shape=(genes, cells)
          the dependent variable (unspliced)
      X: np.ndarray, shape=(genes, cells)
          the independent variable (spliced)
      """
      # NOTE this could be easily parallelized
      slopes = np.zeros(Y.shape[0], dtype="float32")
      offsets = np.zeros(Y.shape[0], dtype="float32")
      for i in range(Y.shape[0]):
          m, q = _fit1_slope_offset(Y[i, :], X[i, :], fixperc_q)
          slopes[i] = m
          offsets[i] = q
      return slopes, offsets
  
  
  def fit_slope_weighted(Y: np.ndarray, X: np.ndarray, W: np.ndarray, return_R2: bool=False, limit_gamma: bool=False, bounds: Tuple[float, float]=(0, 20)) -> np.ndarray:
      """Loop through the genes and fits the slope
  
      Y: np.ndarray, shape=(genes, cells)
          the dependent variable (unspliced)
      X: np.ndarray, shape=(genes, cells)
          the independent variable (spliced)
      W: np.ndarray, shape=(genes, cells)
          the weights that will scale the square residuals
      """
      # NOTE this could be easily parallelized
      # slopes = np.fromiter((_fit1_slope_weighted(Y[i, :], X[i, :], W[i, :], bounds=bounds) for i in range(Y.shape[0])),
      #                     dtype="float32",
      #                     count=Y.shape[0])
  
      slopes = np.zeros(Y.shape[0], dtype="float32")
      offsets = np.zeros(Y.shape[0], dtype="float32")
      if return_R2:
          R2 = np.zeros(Y.shape[0], dtype="float32")
      for i in range(Y.shape[0]):
          m = _fit1_slope_weighted(Y[i, :], X[i, :], W[i, :], limit_gamma)
          slopes[i] = m
          if return_R2:
              # NOTE: the coefficient of determination is not weighted but the fit is
              with np.errstate(divide='ignore', invalid='ignore'):
                  SSres = np.sum((m * X[i, :] - Y[i, :])**2)
                  SStot = np.sum((Y[i, :].mean() - Y[i, :])**2)
                  _R2 = 1 - (SSres / SStot)
              if np.isfinite(_R2):
                  R2[i] = _R2
              else:
                  R2[i] = -1e16
      if return_R2:
          return slopes, R2
      return slopes
  
  
  def fit_slope_weighted_offset(Y: np.ndarray, X: np.ndarray, W: np.ndarray, fixperc_q: bool=False, return_R2: bool=True, limit_gamma: bool=False) -> Any:
      """Loop through the genes and fits the slope
  
      Y: np.ndarray, shape=(genes, cells)
          the dependent variable (unspliced)
      X: np.ndarray, shape=(genes, cells)
          the independent variable (spliced)
      """
      # NOTE this could be easily parallelized
      slopes = np.zeros(Y.shape[0], dtype="float32")
      offsets = np.zeros(Y.shape[0], dtype="float32")
      if return_R2:
          R2 = np.zeros(Y.shape[0], dtype="float32")
      for i in range(Y.shape[0]):
          m, q = _fit1_slope_weighted_offset(Y[i, :], X[i, :], W[i, :], fixperc_q, limit_gamma)
          slopes[i] = m
          offsets[i] = q
          if return_R2:
              # NOTE: the coefficient of determination is not weighted but the fit is
              with np.errstate(divide='ignore', invalid='ignore'):
                  SSres = np.sum((m * X[i, :] + q - Y[i, :])**2)
                  SStot = np.sum((Y[i, :].mean() - Y[i, :])**2)
                  _R2 = 1 - (SSres / SStot)
              if np.isfinite(_R2):
                  R2[i] = _R2
              else:
                  R2[i] = -1e16
      if return_R2:
          return slopes, offsets, R2
      return slopes, offsets
  
  
  def clusters_stats(U: np.ndarray, S: np.ndarray,
                     clusters_uid: np.ndarray, cluster_ix: np.ndarray,
                     size_limit: int=40) -> Tuple[np.ndarray, np.ndarray]:
      """Calculate the averages per cluster
      
      If the cluster is too small (size<size_limit) the average of the toal is reported instead
      """
      U_avgs = np.zeros((S.shape[0], len(clusters_uid)))
      S_avgs = np.zeros((S.shape[0], len(clusters_uid)))
      avgU_div_avgS = np.zeros((S.shape[0], len(clusters_uid)))
      slopes_by_clust = np.zeros((S.shape[0], len(clusters_uid)))
      for i, uid in enumerate(clusters_uid):
          cluster_filter = cluster_ix == i
          n_cells = np.sum(cluster_filter)
          logging.info(f"Cluster: {uid} ({n_cells} cells)")
          if n_cells > size_limit:
              U_avgs[:, i], S_avgs[:, i] = U[:, cluster_filter].mean(1), S[:, cluster_filter].mean(1)
          else:
              U_avgs[:, i], S_avgs[:, i] = U.mean(1), S.mean(1)
              
      return U_avgs, S_avgs
  