<gh_stars>0
import parasail
from ..util._multiprocessing import EnhancedPool as Pool
import itertools
from anndata import AnnData
from typing import Union, Sequence, List, Tuple, Dict, Optional
from .._compat import Literal
import numpy as np
from scanpy import logging
from ..util import _is_na, deprecated
import abc
from Levenshtein import distance as levenshtein_dist
from Levenshtein import hamming as hamming_dist
import scipy.spatial
import scipy.sparse
from scipy.sparse import coo_matrix, csr_matrix
from ..util import _doc_params
from tqdm import tqdm

_doc_metrics = """\
metric
    You can choose one of the following metrics:
      * `identity` -- 1 for identical sequences, 0 otherwise.
        See :class:`~scirpy.ir_dist.IdentityDistanceCalculator`.
        This metric implies a cutoff of 0.
      * `levenshtein` -- Levenshtein edit distance.
        See :class:`~scirpy.ir_dist.LevenshteinDistanceCalculator`.
      * `hamming` -- Hamming distance for CDR3 sequences of equal length.
        See :class:`~scirpy.ir_dist.HammingDistanceCalculator`.
      * `alignment` -- Distance based on pairwise sequence alignments using the
        BLOSUM62 matrix. This option is incompatible with nucleotide sequences.
        See :class:`~scirpy.ir_dist.AlignmentDistanceCalculator`.
      * any instance of :class:`~scirpy.ir_dist.DistanceCalculator`.
"""

_doc_cutoff = """\
cutoff
    All distances `> cutoff` will be replaced by `0` and eliminated from the sparse
    matrix. A sensible cutoff depends on the distance metric, you can find
    information in the corresponding docs.
"""

_doc_params_distance_calculator = """\
cutoff
    Will eleminate distances > cutoff to make efficient
    use of sparse matrices.
"""

_doc_params_parallel_distance_calculator = (
    _doc_params_distance_calculator
    + """\
n_jobs
    Number of jobs to use for the pairwise distance calculation.
    If None, use all jobs (only for ParallelDistanceCalculators).
block_size
    The width of a block of the matrix that will be delegated to a worker
    process. The block contains `block_size ** 2` elements.
"""
)

_doc_dist_mat = """\
Calculates the upper triangle, including the diagonal.

.. important::
    * Distances are offset by 1 to allow efficient use of sparse matrices
      (:math:`d' = d+1`).
    * That means, a `distance > cutoff` is represented as `0`, a `distance == 0`
      is represented as `1`, a `distance == 1` is represented as `2` and so on.
    * Only returns distances `<= cutoff`. Larger distances are eliminated
      from the sparse matrix.
    * Distances are non-negative.
"""


@_doc_params(params=_doc_params_distance_calculator)
class DistanceCalculator(abc.ABC):
    """\
    Abstract base class for a :term:`CDR3`-sequence distance calculator.

    Parameters
    ----------
    {params}

    """

    #: The sparse matrix dtype. Defaults to uint8, constraining the max distance to 255.
    DTYPE = "uint8"

    def __init__(self, cutoff: float):
        if cutoff > 255:
            raise ValueError(
                "Using a cutoff > 255 is not possible due to the `uint8` dtype used"
            )
        self.cutoff = cutoff

    @_doc_params(dist_mat=_doc_dist_mat)
    @abc.abstractmethod
    def calc_dist_mat(
        self, seqs: Sequence[str], seqs2: Optional[Sequence[str]] = None
    ) -> coo_matrix:
        """\
        Calculate pairwise distance matrix of all sequences in `seqs` and `seqs2`.

        When `seqs2` is omitted, computes the pairwise distance of `seqs` against
        itself.

        {dist_mat}

        Parameters
        ----------
        seqs
            array containing CDR3 sequences. Must not contain duplicates.
        seqs2
            second array containing CDR3 sequences. Must not contain
            duplicates either.

        Returns
        -------
        Sparse pairwise distance matrix. If only `seqs` is provided, only
        the upper triangular distance matrix is returned.
        """
        pass


@_doc_params(params=_doc_params_parallel_distance_calculator)
class ParallelDistanceCalculator(DistanceCalculator):
    """
    Abstract base class for a DistanceCalculator that computes distances in parallel.

    It does so in a blockwise fashion. The function computing distances
    for a single block needs to be overriden.

    Parameters
    ----------
    {params}
    """

    def __init__(
        self,
        cutoff: float,
        n_jobs: Optional[int] = None,
        block_size: Optional[int] = 50,
    ):
        super().__init__(cutoff)
        self.n_jobs = n_jobs
        self.block_size = block_size

    @abc.abstractmethod
    def _compute_block(
        self,
        seqs1: Sequence[str],
        seqs2: Union[Sequence[str], None],
        origin: Tuple[int, int],
    ) -> Tuple[int, int, int]:
        """Compute the distances for a block of the matrix

        Parameters
        ----------
        seqs1
            array containing sequences
        seqs2
            other array containing sequences. If `None` compute the square matrix
            of `seqs1` and iteratoe over the upper triangle including the diagonal only.
        origin
            row, col coordinates of the origin of the block.

        Returns
        ------
        List of (distance, row, col) tuples for all elements with distance != 0.
        row, col must be the coordinates in the final matrix (they can be derived using
        `origin`). Can't be a generator because this needs to be picklable.
        """
        pass

    @staticmethod
    def _block_iter(
        seqs1: Sequence[str],
        seqs2: Optional[Sequence[str]] = None,
        block_size: Optional[int] = 50,
    ) -> Tuple[Sequence[str], Union[Sequence[str], None], Tuple[int, int]]:
        """Iterate over sequences in blocks.

        Parameters
        ----------
        seqs1
            array containing (unique) sequences
        seqs2
            array containing other sequences. If `None` compute
            the square matrix of `seqs1` and iterate over the upper triangle (including
            the diagonal) only.
        block_size
            side length of a block (will have `block_size ** 2` elements.)

        Yields
        ------
        seqs1
            subset of length `block_size` of seqs1
        seqs2
            subset of length `block_size` of seqs2. If seqs2 is None, this will
            be `None` if the block is on the diagonal, or a subset of seqs1 otherwise.
        origin
            (row, col) coordinates of the origin of the block.
        """
        square_mat = seqs2 is None
        if square_mat:
            seqs2 = seqs1
        for row in range(0, len(seqs1), block_size):
            start_col = row if square_mat else 0
            for col in range(start_col, len(seqs2), block_size):
                if row == col and square_mat:
                    # block on the diagonal.
                    # yield None for seqs2 to indicate that we only want the upper
                    # diagonal.
                    yield seqs1[row : row + block_size], None, (row, row)
                else:
                    yield seqs1[row : row + block_size], seqs2[
                        col : col + block_size
                    ], (row, col)

    def calc_dist_mat(
        self, seqs: Sequence[str], seqs2: Optional[Sequence[str]] = None
    ) -> coo_matrix:
        """Calculate the distance matrix.

        See :meth:`DistanceCalculator.calc_dist_mat`."""
        # precompute blocks as list to have total number of blocks for progressbar
        blocks = list(self._block_iter(seqs, seqs2, self.block_size))

        with Pool(self.n_jobs) as p:
            block_results = p.starmap_progress(
                self._compute_block, blocks, total=len(blocks)
            )

        try:
            dists, rows, cols = zip(*itertools.chain(*block_results))
        except ValueError:
            # happens when there is no match at all
            dists, rows, cols = (), (), ()

        shape = (len(seqs), len(seqs2)) if seqs2 is not None else (len(seqs), len(seqs))
        score_mat = scipy.sparse.coo_matrix(
            (dists, (rows, cols)), dtype=self.DTYPE, shape=shape
        )
        score_mat.eliminate_zeros()

        return score_mat


@_doc_params(params=_doc_params_distance_calculator)
class IdentityDistanceCalculator(DistanceCalculator):
    """\
    Calculates the Identity-distance between :term:`CDR3` sequences.

    The identity distance is defined as
        * `0`, if sequences are identical
        * `1`, if sequences are not identical.

    Choosing a cutoff:
        For this DistanceCalculator, per definition, the cutoff = 0.
        The `cutoff` argument is ignored.

    Parameters
    ----------
    {params}
    """

    def __init__(self, cutoff: float = 0):
        super().__init__(cutoff)

    def calc_dist_mat(self, seqs: np.ndarray, seqs2: np.ndarray = None) -> coo_matrix:
        """In this case, the offseted distance matrix is the identity matrix.

        More details: :meth:`DistanceCalculator.calc_dist_mat`"""
        if seqs2 is None:
            # In this case, the offsetted distance matrix is the identity matrix
            return scipy.sparse.identity(len(seqs), dtype=self.DTYPE, format="coo")
        else:
            # actually compare the values
            def coord_generator():
                for (i1, s1), (i2, s2) in itertools.product(
                    enumerate(seqs), enumerate(seqs2)
                ):
                    if s1 == s2:
                        yield 1, i1, i2

            try:
                d, row, col = zip(*coord_generator())
            except ValueError:
                # happens when there is no match at all
                d, row, col = (), (), ()

            return coo_matrix(
                (d, (row, col)), dtype=self.DTYPE, shape=(len(seqs), len(seqs2))
            )


@_doc_params(params=_doc_params_parallel_distance_calculator)
class LevenshteinDistanceCalculator(ParallelDistanceCalculator):
    """\
    Calculates the Levenshtein edit-distance between sequences.

    The edit distance is the total number of deletion, addition and modification
    events.

    This class relies on `Python-levenshtein <https://github.com/ztane/python-Levenshtein>`_
    to calculate the distances.

    Choosing a cutoff:
        Each modification stands for a deletion, addition or modification event.
        While lacking empirical data, it seems unlikely that CDR3 sequences with more
        than two modifications still recognize the same antigen.

    Parameters
    ----------
    {params}
    """

    def _compute_block(self, seqs1, seqs2, origin):
        origin_row, origin_col = origin
        if seqs2 is not None:
            # compute the full matrix
            coord_iterator = itertools.product(enumerate(seqs1), enumerate(seqs2))
        else:
            # compute only upper triangle in this case
            coord_iterator = itertools.combinations_with_replacement(
                enumerate(seqs1), r=2
            )

        result = []
        for (row, s1), (col, s2) in coord_iterator:
            d = levenshtein_dist(s1, s2)
            if d <= self.cutoff:
                result.append((d + 1, origin_row + row, origin_col + col))

        return result


@_doc_params(params=_doc_params_parallel_distance_calculator)
class HammingDistanceCalculator(ParallelDistanceCalculator):
    """\
    Calculates the Hamming distance between sequences of identical length.

    The edit distance is the total number of substitution events. Sequences
    with different lengths will be treated as though they exceeded the
    distance-cutoff, i.e. they receive a distance of `0` in the sparse distance
    matrix and will not be connected by an edge in the graph.

    This class relies on `Python-levenshtein <https://github.com/ztane/python-Levenshtein>`_
    to calculate the distances.

    Choosing a cutoff:
        Each modification stands for a substitution event.
        While lacking empirical data, it seems unlikely that CDR3 sequences with more
        than two modifications still recognize the same antigen.

    Parameters
    ----------
    {params}
    """

    def _compute_block(self, seqs1, seqs2, origin):
        origin_row, origin_col = origin
        if seqs2 is not None:
            # compute the full matrix
            coord_iterator = itertools.product(enumerate(seqs1), enumerate(seqs2))
        else:
            # compute only upper triangle in this case
            coord_iterator = itertools.combinations_with_replacement(
                enumerate(seqs1), r=2
            )

        result = []
        for (row, s1), (col, s2) in coord_iterator:

            # require identical length of sequences
            if len(s1) != len(s2):
                continue
            d = hamming_dist(s1, s2)
            if d <= self.cutoff:
                result.append((d + 1, origin_row + row, origin_col + col))

        return result


@_doc_params(params=_doc_params_parallel_distance_calculator)
class AlignmentDistanceCalculator(ParallelDistanceCalculator):
    """\
    Calculates distance between sequences based on pairwise sequence alignment.

    The distance between two sequences is defined as :math:`S_{{1,2}}^{{max}} - S_{{1,2}}`,
    where :math:`S_{{1,2}}` is the alignment score of sequences 1 and 2 and
    :math:`S_{{1,2}}^{{max}}` is the max. achievable alignment score of sequences 1 and 2.
    :math:`S_{{1,2}}^{{max}}` is defined as :math:`\\min(S_{{1,1}}, S_{{2,2}})`.

    The use of alignment-based distances is heavily inspired by :cite:`TCRdist`.

    High-performance sequence alignments are calculated leveraging
    the `parasail library <https://github.com/jeffdaily/parasail-python>`_ (:cite:`Daily2016`).

    Choosing a cutoff:
        Alignment distances need to be viewed in the light of the substitution matrix.
        The alignment distance is the difference between the actual alignment
        score and the max. achievable alignment score. For instance, a mutation
        from *Leucine* (`L`) to *Isoleucine* (`I`) results in a BLOSUM62 score of `2`.
        An `L` aligned with `L` achieves a score of `4`. The distance is, therefore, `2`.

        On the other hand, a single *Tryptophane* (`W`) mutating into, e.g.
        *Proline* (`P`) already results in a distance of `15`.

        We are still lacking empirical data up to which distance a CDR3 sequence still
        is likely to recognize the same antigen, but reasonable cutoffs are `<15`.

    Parameters
    ----------
    {params}
    subst_mat
        Name of parasail substitution matrix
    gap_open
        Gap open penalty
    gap_extend
        Gap extend penatly
    """

    def __init__(
        self,
        cutoff: float,
        n_jobs: Union[int, None] = None,
        block_size: int = 50,
        *,
        subst_mat: str = "blosum62",
        gap_open: int = 11,
        gap_extend: int = 11,
    ):
        super().__init__(cutoff, n_jobs, block_size)
        self.subst_mat = subst_mat
        self.gap_open = gap_open
        self.gap_extend = gap_extend

    def _compute_block(self, seqs1, seqs2, origin):
        subst_mat = parasail.Matrix(self.subst_mat)
        origin_row, origin_col = origin

        square_matrix = seqs2 is None
        if square_matrix:
            seqs2 = seqs1

        self_alignment_scores1 = self._self_alignment_scores(seqs1)
        if square_matrix:
            self_alignment_scores2 = self_alignment_scores1
        else:
            self_alignment_scores2 = self._self_alignment_scores(seqs2)

        result = []
        for row, s1 in enumerate(seqs1):
            col_start = row if square_matrix else 0
            for col, s2 in enumerate(seqs2[col_start:], start=col_start):
                profile = parasail.profile_create_16(s1, subst_mat)
                r = parasail.nw_scan_profile_16(
                    profile, s2, self.gap_open, self.gap_extend
                )
                max_score = np.min(
                    [self_alignment_scores1[row], self_alignment_scores2[col]]
                )
                d = max_score - r.score
                if d <= self.cutoff:
                    result.append((d + 1, origin_row + row, origin_col + col))

        return result

    def _self_alignment_scores(self, seqs: Sequence) -> dict:
        """Calculate self-alignments. We need them as reference values
        to turn scores into dists"""
        return np.fromiter(
            (
                parasail.nw_scan_16(
                    s,
                    s,
                    self.gap_open,
                    self.gap_extend,
                    parasail.Matrix(self.subst_mat),
                ).score
                for s in seqs
            ),
            dtype=int,
            count=len(seqs),
        )


@deprecated(
    "Due to added BCR support, this function has been renamed "
    "to `sequence_dist`. The old version will be removed in a future release. "
)
def tcr_dist(*args, **kwargs):
    return sequence_dist(*args, **kwargs)


@_doc_params(metric=_doc_metrics, cutoff=_doc_cutoff, dist_mat=_doc_dist_mat)
def sequence_dist(
    unique_seqs: np.ndarray,
    unique_seqs2: Union[None, np.ndarray] = None,
    *,
    metric: Union[
        Literal["alignment", "identity", "levenshtein", "hamming"],
        DistanceCalculator,
    ] = "identity",
    cutoff: float = 10,
    n_jobs: Union[int, None] = None,
) -> coo_matrix:
    """\
    Calculate a sequence x sequence distance matrix.

    {dist_mat}

    Parameters
    ----------
    unique_seqs
        Numpy array of nucleotide or amino acid sequences.
        Must not contain duplicates.
        Note that not all distance metrics support nucleotide sequences.
    unique_seqs2
        Second array sequences. When omitted, `sequence_dist` computes
        the square matrix of `unique_seqs`.
    {metric}
    {cutoff}

        A cutoff of 0 implies the `identity` metric.

    Returns
    -------
    Upper triangular distance matrix when only `unique_seqs` is provided.
    A full rectangular distance matrix when both `unique_seqs` and
    `unique_seqs2` are provided.
    """
    if cutoff == 0 or metric == "identity":
        metric = "identity"
        cutoff = 0
    if isinstance(metric, DistanceCalculator):
        dist_calc = metric
    elif metric == "alignment":
        dist_calc = AlignmentDistanceCalculator(cutoff=cutoff, n_jobs=n_jobs)
    elif metric == "identity":
        dist_calc = IdentityDistanceCalculator(cutoff=cutoff)
    elif metric == "levenshtein":
        dist_calc = LevenshteinDistanceCalculator(cutoff=cutoff, n_jobs=n_jobs)
    elif metric == "hamming":
        dist_calc = HammingDistanceCalculator(cutoff=cutoff, n_jobs=n_jobs)
    else:
        raise ValueError("Invalid distance metric.")

    dist_mat = dist_calc.calc_dist_mat(unique_seqs, unique_seqs2)
    return dist_mat


@deprecated(
    "Due to added BCR support, this function has been renamed "
    "to `IrNeighbors`. The old version will be removed in a future release. "
)
def TcrNeighbors(*args, **kwargs):
    return IrNeighbors(*args, **kwargs)


class IrNeighbors:
    def __init__(
        self,
        adata: AnnData,
        *,
        metric: Union[
            Literal["alignment", "identity", "levenshtein", "hamming"],
            DistanceCalculator,
        ] = "identity",
        cutoff: float = 10,
        receptor_arms: Literal["VJ", "VDJ", "all", "any"] = "all",
        dual_ir: Literal["primary_only", "all", "any"] = "primary_only",
        sequence: Literal["aa", "nt"] = "aa",
    ):
        """Class to compute Neighborhood graphs of CDR3 sequences.

        For documentation of the parameters, see :func:`ir_neighbors`.
        """
        start = logging.info("Initializing IrNeighbors object...")
        if metric == "identity" and cutoff != 0:
            raise ValueError("Identity metric only works with cutoff == 0")
        if metric != "identity" and cutoff == 0:
            logging.warning(f"Running with {metric} metric, but cutoff == 0. ")
        if sequence == "nt" and metric == "alignment":
            raise ValueError(
                "Using nucleotide sequences with alignment metric is not supported. "
            )
        if receptor_arms not in ["VJ", "VDJ", "all", "any"]:
            raise ValueError(
                "Invalid value for `receptor_arms`. Note that starting with v0.5 "
                "`TRA` and `TRB` are not longer valid values."
            )
        if dual_ir not in ["primary_only", "all", "any"]:
            raise ValueError("Invalid value for `dual_ir")
        if sequence not in ["aa", "nt"]:
            raise ValueError("Invalid value for `sequence`")
        self.adata = adata
        self.metric = metric
        self.cutoff = cutoff
        self.receptor_arms = receptor_arms
        self.dual_ir = dual_ir
        self.sequence = sequence
        self._build_index_dict()
        self._dist_mat = None
        logging.info("Finished initalizing IrNeighbors object. ", time=start)

    @staticmethod
    def _seq_to_cell_idx(
        unique_seqs: np.ndarray, cdr_seqs: np.ndarray
    ) -> Dict[int, List[int]]:
        """
        Compute sequence to cell index for a single chain (e.g. `TRA_1`).

        Maps cell_idx -> [list, of, seq_idx].
        Useful to build a cell x cell matrix from a seq x seq matrix.

        Computes magic lookup indexes in linear time.

        Parameters
        ----------
        unique_seqs
            Pool of all unique cdr3 sequences (length = #unique cdr3 sequences)
        cdr_seqs
            CDR3 sequences for the current chain (length = #cells)

        Returns
        -------
        Sequence2Cell mapping
        """
        # 1) reverse mapping of amino acid sequence to index in sequence-distance matrix
        seq_to_index = {seq: i for i, seq in enumerate(unique_seqs)}

        # 2) indices of cells in adata that have a CDR3 sequence.
        cells_with_chain = np.where(~_is_na(cdr_seqs))[0]

        # 3) indices of the corresponding sequences in the distance matrix.
        seq_inds = {
            chain_id: seq_to_index[cdr_seqs[chain_id]] for chain_id in cells_with_chain
        }

        # 4) list of cell-indices in the cell distance matrix for each sequence
        seq_to_cell = {seq_id: list() for seq_id in seq_to_index.values()}
        for cell_id in cells_with_chain:
            seq_id = seq_inds[cell_id]
            seq_to_cell[seq_id].append(cell_id)

        return seq_to_cell

    def _build_index_dict(self):
        """Build nested dictionary for each receptor arm (TRA, TRB) containing all
        combinations of receptor_arms x primary/secondary_chain

        If the merge mode for either `receptor_arm` or `dual_ir` is `all`,
        includes a lookup table that contains the number of CDR3 sequences for
        each cell.
        """
        receptor_arms = (
            ["VJ", "VDJ"]
            if self.receptor_arms not in ["VJ", "VDJ"]
            else [self.receptor_arms]
        )
        chain_inds = [1] if self.dual_ir == "primary_only" else [1, 2]
        sequence = "" if self.sequence == "aa" else "_nt"

        arm_dict = {}
        for arm in receptor_arms:
            cdr_seqs = {
                k: self.adata.obs[f"IR_{arm}_{k}_cdr3{sequence}"].values
                for k in chain_inds
            }
            unique_seqs = np.hstack(list(cdr_seqs.values()))
            unique_seqs = np.unique(unique_seqs[~_is_na(unique_seqs)]).astype(str)
            seq_to_cell = {
                k: self._seq_to_cell_idx(unique_seqs, cdr_seqs[k]) for k in chain_inds
            }
            arm_dict[arm] = {
                "chain_inds": chain_inds,
                "unique_seqs": unique_seqs,
                "seq_to_cell": seq_to_cell,
            }

            # need the count of chains per cell for the `all` strategies.
            if self.receptor_arms == "all" or self.dual_ir == "all":
                arm_dict[arm]["chains_per_cell"] = np.sum(
                    ~_is_na(
                        self.adata.obs.loc[
                            :, [f"IR_{arm}_{k}_cdr3{sequence}" for k in chain_inds]
                        ]
                    ),
                    axis=1,
                )

        self.index_dict = arm_dict

    def _reduce_dual_all(self, d, chain, cell_row, cell_col):
        """Reduce dual IRs into a single value when 'all' sequences
        need to match. This requires additional checking effort for the number
        of chains in the given cell, since we can't make the distinction between
        no chain and dist > cutoff based on the distances (both would contain a
        0 in the distance matrix)."""
        chain_count = (
            self.index_dict[chain]["chains_per_cell"][cell_row],
            self.index_dict[chain]["chains_per_cell"][cell_col],
        )
        if len(d) == 1 and chain_count == (1, 1):
            # exactely one chain for both cells -> return that value
            return next(iter(d.values()))
        elif chain_count == (2, 2):
            # two options: either (1 matches 2 and 2 matches 1)
            # or (1 matches 1 and 2 matches 2).
            try:
                # minus 1, because both dists are offseted by 1.
                d1 = d[(1, 2)] + d[(2, 1)] - 1
            except KeyError:
                d1 = None
            try:
                d2 = d[(1, 1)] + d[(2, 2)] - 1
            except KeyError:
                d2 = None

            if d1 is not None and d2 is not None:
                return min(d1, d2)
            elif d1 is not None:
                return d1
            elif d2 is not None:
                return d2
            else:
                return 0

        else:
            return 0

    def _reduce_arms_all(self, values, cell_row, cell_col):
        """Reduce multiple receptor arms into a single value when 'all' sequences
        need to match. This requires additional checking effort for teh number
        of chains in the given cell, since we can't make the distinction between
        no chain and dist > cutoff based on the distances (both would contain a
        0 in the distance matrix)."""
        values = (x for x in values if x != 0)

        try:
            arm1 = next(values)
        except StopIteration:
            # no value > 0
            return 0

        try:
            arm2 = next(values)
            # two receptor arms -> easy
            # -1 because both distances are offseted by 1
            return arm1 + arm2 - 1
        except StopIteration:
            # only one arm
            tra_chains = (
                self.index_dict["VJ"]["chains_per_cell"][cell_row],
                self.index_dict["VJ"]["chains_per_cell"][cell_col],
            )
            trb_chains = (
                self.index_dict["VDJ"]["chains_per_cell"][cell_row],
                self.index_dict["VDJ"]["chains_per_cell"][cell_col],
            )
            # Either exactely one chain for VJ or
            # exactely one chain for VDJ for both cells.
            if tra_chains == (0, 0) or trb_chains == (0, 0):
                return arm1
            else:
                return 0

    @staticmethod
    def _reduce_arms_any(lst, *args):
        """Reduce arms when *any* of the sequences needs to match.
        This is the simpler case. This also works with only one entry
        (e.g. arms = "VJ")"""
        # need to exclude 0 values, since the dist mat is offseted by 1.
        try:
            return min(x for x in lst if x != 0)
        except ValueError:
            # no values in generator
            return 0

    @staticmethod
    def _reduce_dual_any(d, *args):
        """Reduce dual irs to a single value when *any* of the sequences needs
        to match (by minimum). This also works with only one entry (i.e. 'primary only')
        """
        # need to exclude 0 values, since the dist mat is offseted by 1.
        try:
            return min(x for x in d.values() if x != 0)
        except ValueError:
            # no values in generator
            return 0

    def _reduce_coord_dict(self, coord_dict):
        """Applies reduction functions to the coord dict.
        Yield (coords, value) pairs."""
        start = logging.info("Constructing cell x cell distance matrix...")
        reduce_dual = (
            self._reduce_dual_all if self.dual_ir == "all" else self._reduce_dual_any
        )
        reduce_arms = (
            self._reduce_arms_all
            if self.receptor_arms == "all"
            else self._reduce_arms_any
        )
        for (cell_row, cell_col), entry in tqdm(
            coord_dict.items(), total=len(coord_dict)
        ):
            reduced_dual = (
                reduce_dual(value_dict, chain, cell_row, cell_col)
                for chain, value_dict in entry.items()
            )
            reduced = reduce_arms(
                reduced_dual,
                cell_row,
                cell_col,
            )
            yield (cell_row, cell_col), reduced
        logging.info("Finished constructing cell x cell distance matrix. ", time=start)

    def _cell_dist_mat_reduce(self):
        """Compute the distance matrix by using custom reduction functions.
        More flexible than `_build_cell_dist_mat_min`, but requires more memory.
        Reduce dual is called before reduce arms.
        """
        coord_dict = dict()

        def _add_to_dict(d, c1, c2, cell_row, cell_col, value):
            """Add a value to the nested coord dict"""
            try:
                tmp_dict = d[(cell_row, cell_col)]
                try:
                    tmp_dict2 = tmp_dict[arm]
                    try:
                        if (c1, c2) in tmp_dict2:
                            # can be in arbitrary order apprarently
                            assert (c2, c1) not in tmp_dict2
                            tmp_dict2[(c2, c1)] = value
                        tmp_dict2[(c1, c2)] = value
                    except KeyError:
                        tmp_dict2 = {(c1, c2): value}
                except KeyError:
                    tmp_dict[arm] = {(c1, c2): value}
            except KeyError:
                d[(cell_row, cell_col)] = {arm: {(c1, c2): value}}

        for arm, arm_info in self.index_dict.items():
            dist_mat, seq_to_cell, chain_inds = (
                arm_info["dist_mat"],
                arm_info["seq_to_cell"],
                arm_info["chain_inds"],
            )
            start = logging.info(f"Started comstructing {arm} coord-dictionary...")
            for row, col, value in tqdm(
                zip(dist_mat.row, dist_mat.col, dist_mat.data), total=dist_mat.nnz
            ):
                for c1, c2 in itertools.product(chain_inds, repeat=2):
                    for cell_row, cell_col in itertools.product(
                        seq_to_cell[c1][row], seq_to_cell[c2][col]
                    ):
                        # fill upper diagonal. Important: these are dist-mat row,cols
                        # not cell-mat row cols. This is required, because the
                        # itertools.product returns all combinations for the diagonal
                        # but not for the other values.
                        _add_to_dict(coord_dict, c1, c2, cell_row, cell_col, value)
                        if row != col:
                            _add_to_dict(coord_dict, c1, c2, cell_col, cell_row, value)

            logging.info(f"Finished constructing {arm} coord-dictionary", time=start)

        yield from self._reduce_coord_dict(coord_dict)

    def compute_distances(
        self,
        n_jobs: Union[int, None] = None,
    ):
        """Computes the distances between CDR3 sequences

        Parameters
        ----------
        n_jobs
            Number of CPUs to use for alignment, levenshtein or hamming distance.
            Default: use all CPUS.
        """
        for arm, arm_dict in self.index_dict.items():
            start = logging.info(f"Computing {arm} pairwise distances...")

            arm_dict["dist_mat"] = sequence_dist(
                arm_dict["unique_seqs"],
                metric=self.metric,
                cutoff=self.cutoff,
                n_jobs=n_jobs,
            )
            logging.info(f"Finished computing {arm} pairwise distances.", time=start)

        try:
            coords, values = zip(*self._cell_dist_mat_reduce())
        except ValueError:
            # happens when there are no distances computed (e.g. all CDR3s are NAN)
            # -> construct empty CSR matrix
            self._dist_mat = csr_matrix((self.adata.n_obs, self.adata.n_obs))
        else:
            rows, cols = zip(*coords)
            dist_mat = coo_matrix(
                (values, (rows, cols)), shape=(self.adata.n_obs, self.adata.n_obs)
            )
            dist_mat.eliminate_zeros()
            self._dist_mat = dist_mat.tocsr()

    @property
    def dist(self):
        """The computed distance matrix.
        Requires to invoke `compute_distances() first."""
        return self._dist_mat

    @property
    def connectivities(self):
        """Get the weighted adjacecency matrix derived from the distance matrix.

        The cutoff will be used to normalize the distances.
        """
        if self.cutoff == 0:
            return self._dist_mat

        start = logging.debug("Started converting distances to connectivities. ")

        connectivities = self._dist_mat.copy()

        # actual distances
        d = connectivities.data - 1

        # structure of the matrix stayes the same, we can safely change the data only
        connectivities.data = (self.cutoff - d) / self.cutoff
        connectivities.eliminate_zeros()
        logging.debug("Finished converting distances to connectivities. ", time=start)
        return connectivities


@_doc_params(metric=_doc_metrics, cutoff=_doc_cutoff, dist_mat=_doc_dist_mat)
def ir_neighbors(
    adata: AnnData,
    *,
    metric: Union[
        Literal["identity", "alignment", "levenshtein", "hamming"],
        DistanceCalculator,
    ] = "identity",
    cutoff: int = 10,
    receptor_arms: Literal["VJ", "VDJ", "all", "any"] = "all",
    dual_ir: Literal["primary_only", "any", "all"] = "primary_only",
    key_added: Union[str, None] = None,
    sequence: Literal["aa", "nt"] = "nt",
    inplace: bool = True,
    n_jobs: Union[int, None] = None,
) -> Union[Tuple[csr_matrix, csr_matrix], None]:
    """\
    Construct a neighborhood graph based on :term:`CDR3` sequence similarity.

    All cells with a CDR3 distance `< cutoff` receive an edge in the graph.
    Edges are weighted by the distance.

    Parameters
    ----------
    adata
        annotated data matrix
    {metric}
    {cutoff}

        Two cells with a distance <= the cutoff will be connected.
        A cutoff of 0 implies the use of the `identity` metric.

    receptor_arms:
         * `"TRA"` - only consider TRA sequences
         * `"TRB"` - only consider TRB sequences
         * `"all"` - both TRA and TRB need to match
         * `"any"` - either TRA or TRB need to match

    dual_ir:
         * `"primary_only"` - only consider most abundant pair of TRA/TRB chains
         * `"any"` - consider both pairs of TRA/TRB sequences. Distance must be below
           cutoff for any of the chains.
         * `"all"` - consider both pairs of TRA/TRB sequences. Distance must be below
           cutoff for all of the chains.

        See also :term:`Dual IR`

    key_added:
        dict key under which the result will be stored in `adata.uns`
        when `inplace` is True. Defaults to `ir_neighbors_{{sequence}}_{{metric}}`.

        If metric is an instance of :class:`scirpy.ir_dist.DistanceCalculator`,
        `{{metric}}` defaults to `"custom"`.
    sequence:
        Use amino acid (`aa`) or nulceotide (`nt`) sequences?
    inplace:
        If `True`, store the results in `adata.uns`. Otherwise return
        the results.
    n_jobs:
        Number of cores to use for alignment, levenshtein or hamming distance.

    Returns
    -------
    connectivities
        weighted adjacency matrix
    dist
        cell x cell distance matrix with the distances as computed according to `metric`
        offsetted by 1 to make use of sparse matrices.
    """
    if cutoff == 0 or metric == "identity":
        metric = "identity"
        cutoff = 0
    if key_added is None:
        tmp_metric = "custom" if isinstance(metric, DistanceCalculator) else metric
        key_added = f"ir_neighbors_{sequence}_{tmp_metric}"
    ad = IrNeighbors(
        adata,
        metric=metric,
        cutoff=cutoff,
        receptor_arms=receptor_arms,
        dual_ir=dual_ir,
        sequence=sequence,
    )
    ad.compute_distances(n_jobs)

    if not inplace:
        return ad.connectivities, ad.dist
    else:
        adata.uns[key_added] = dict()
        adata.uns[key_added]["params"] = {
            "metric": metric,
            "cutoff": cutoff,
            "dual_ir": dual_ir,
            "receptor_arms": receptor_arms,
        }
        adata.uns[key_added]["connectivities"] = ad.connectivities
        adata.uns[key_added]["distances"] = ad.dist

import parasail
from ..util._multiprocessing import EnhancedPool as Pool
import itertools
from anndata import AnnData
from typing import Union, Sequence, List, Tuple, Dict, Optional
from .._compat import Literal
import numpy as np
from scanpy import logging
from ..util import _is_na, deprecated
import abc
from Levenshtein import distance as levenshtein_dist
from Levenshtein import hamming as hamming_dist
import scipy.spatial
import scipy.sparse
from scipy.sparse import coo_matrix, csr_matrix
from ..util import _doc_params
from tqdm import tqdm

_doc_metrics = """\
metric
    You can choose one of the following metrics:
      * `identity` -- 1 for identical sequences, 0 otherwise.
        See :class:`~scirpy.ir_dist.IdentityDistanceCalculator`.
        This metric implies a cutoff of 0.
      * `levenshtein` -- Levenshtein edit distance.
        See :class:`~scirpy.ir_dist.LevenshteinDistanceCalculator`.
      * `hamming` -- Hamming distance for CDR3 sequences of equal length.
        See :class:`~scirpy.ir_dist.HammingDistanceCalculator`.
      * `alignment` -- Distance based on pairwise sequence alignments using the
        BLOSUM62 matrix. This option is incompatible with nucleotide sequences.
        See :class:`~scirpy.ir_dist.AlignmentDistanceCalculator`.
      * any instance of :class:`~scirpy.ir_dist.DistanceCalculator`.
"""

_doc_cutoff = """\
cutoff
    All distances `> cutoff` will be replaced by `0` and eliminated from the sparse
    matrix. A sensible cutoff depends on the distance metric, you can find
    information in the corresponding docs.
"""

_doc_params_distance_calculator = """\
cutoff
    Will eleminate distances > cutoff to make efficient
    use of sparse matrices.
"""

_doc_params_parallel_distance_calculator = (
    _doc_params_distance_calculator
    + """\
n_jobs
    Number of jobs to use for the pairwise distance calculation.
    If None, use all jobs (only for ParallelDistanceCalculators).
block_size
    The width of a block of the matrix that will be delegated to a worker
    process. The block contains `block_size ** 2` elements.
"""
)

_doc_dist_mat = """\
Calculates the upper triangle, including the diagonal.

.. important::
    * Distances are offset by 1 to allow efficient use of sparse matrices
      (:math:`d' = d+1`).
    * That means, a `distance > cutoff` is represented as `0`, a `distance == 0`
      is represented as `1`, a `distance == 1` is represented as `2` and so on.
    * Only returns distances `<= cutoff`. Larger distances are eliminated
      from the sparse matrix.
    * Distances are non-negative.
"""


@_doc_params(params=_doc_params_distance_calculator)
class DistanceCalculator(abc.ABC):
    """\
    Abstract base class for a :term:`CDR3`-sequence distance calculator.

    Parameters
    ----------
    {params}

    """

    #: The sparse matrix dtype. Defaults to uint8, constraining the max distance to 255.
    DTYPE = "uint8"

    def __init__(self, cutoff: float):
        if cutoff > 255:
            raise ValueError(
                "Using a cutoff > 255 is not possible due to the `uint8` dtype used"
            )
        self.cutoff = cutoff

    @_doc_params(dist_mat=_doc_dist_mat)
    @abc.abstractmethod
    def calc_dist_mat(
        self, seqs: Sequence[str], seqs2: Optional[Sequence[str]] = None
    ) -> coo_matrix:
        """\
        Calculate pairwise distance matrix of all sequences in `seqs` and `seqs2`.

        When `seqs2` is omitted, computes the pairwise distance of `seqs` against
        itself.

        {dist_mat}

        Parameters
        ----------
        seqs
            array containing CDR3 sequences. Must not contain duplicates.
        seqs2
            second array containing CDR3 sequences. Must not contain
            duplicates either.

        Returns
        -------
        Sparse pairwise distance matrix. If only `seqs` is provided, only
        the upper triangular distance matrix is returned.
        """
        pass


@_doc_params(params=_doc_params_parallel_distance_calculator)
class ParallelDistanceCalculator(DistanceCalculator):
    """
    Abstract base class for a DistanceCalculator that computes distances in parallel.

    It does so in a blockwise fashion. The function computing distances
    for a single block needs to be overriden.

    Parameters
    ----------
    {params}
    """

    def __init__(
        self,
        cutoff: float,
        n_jobs: Optional[int] = None,
        block_size: Optional[int] = 50,
    ):
        super().__init__(cutoff)
        self.n_jobs = n_jobs
        self.block_size = block_size

    @abc.abstractmethod
    def _compute_block(
        self,
        seqs1: Sequence[str],
        seqs2: Union[Sequence[str], None],
        origin: Tuple[int, int],
    ) -> Tuple[int, int, int]:
        """Compute the distances for a block of the matrix

        Parameters
        ----------
        seqs1
            array containing sequences
        seqs2
            other array containing sequences. If `None` compute the square matrix
            of `seqs1` and iteratoe over the upper triangle including the diagonal only.
        origin
            row, col coordinates of the origin of the block.

        Returns
        ------
        List of (distance, row, col) tuples for all elements with distance != 0.
        row, col must be the coordinates in the final matrix (they can be derived using
        `origin`). Can't be a generator because this needs to be picklable.
        """
        pass

    @staticmethod
    def _block_iter(
        seqs1: Sequence[str],
        seqs2: Optional[Sequence[str]] = None,
        block_size: Optional[int] = 50,
    ) -> Tuple[Sequence[str], Union[Sequence[str], None], Tuple[int, int]]:
        """Iterate over sequences in blocks.

        Parameters
        ----------
        seqs1
            array containing (unique) sequences
        seqs2
            array containing other sequences. If `None` compute
            the square matrix of `seqs1` and iterate over the upper triangle (including
            the diagonal) only.
        block_size
            side length of a block (will have `block_size ** 2` elements.)

        Yields
        ------
        seqs1
            subset of length `block_size` of seqs1
        seqs2
            subset of length `block_size` of seqs2. If seqs2 is None, this will
            be `None` if the block is on the diagonal, or a subset of seqs1 otherwise.
        origin
            (row, col) coordinates of the origin of the block.
        """
        square_mat = seqs2 is None
        if square_mat:
            seqs2 = seqs1
        for row in range(0, len(seqs1), block_size):
            start_col = row if square_mat else 0
            for col in range(start_col, len(seqs2), block_size):
                if row == col and square_mat:
                    # block on the diagonal.
                    # yield None for seqs2 to indicate that we only want the upper
                    # diagonal.
                    yield seqs1[row : row + block_size], None, (row, row)
                else:
                    yield seqs1[row : row + block_size], seqs2[
                        col : col + block_size
                    ], (row, col)

    def calc_dist_mat(
        self, seqs: Sequence[str], seqs2: Optional[Sequence[str]] = None
    ) -> coo_matrix:
        """Calculate the distance matrix.

        See :meth:`DistanceCalculator.calc_dist_mat`."""
        # precompute blocks as list to have total number of blocks for progressbar
        blocks = list(self._block_iter(seqs, seqs2, self.block_size))

        with Pool(self.n_jobs) as p:
            block_results = p.starmap_progress(
                self._compute_block, blocks, total=len(blocks)
            )

        try:
            dists, rows, cols = zip(*itertools.chain(*block_results))
        except ValueError:
            # happens when there is no match at all
            dists, rows, cols = (), (), ()

        shape = (len(seqs), len(seqs2)) if seqs2 is not None else (len(seqs), len(seqs))
        score_mat = scipy.sparse.coo_matrix(
            (dists, (rows, cols)), dtype=self.DTYPE, shape=shape
        )
        score_mat.eliminate_zeros()

        return score_mat


@_doc_params(params=_doc_params_distance_calculator)
class IdentityDistanceCalculator(DistanceCalculator):
    """\
    Calculates the Identity-distance between :term:`CDR3` sequences.

    The identity distance is defined as
        * `0`, if sequences are identical
        * `1`, if sequences are not identical.

    Choosing a cutoff:
        For this DistanceCalculator, per definition, the cutoff = 0.
        The `cutoff` argument is ignored.

    Parameters
    ----------
    {params}
    """

    def __init__(self, cutoff: float = 0):
        super().__init__(cutoff)

    def calc_dist_mat(self, seqs: np.ndarray, seqs2: np.ndarray = None) -> coo_matrix:
        """In this case, the offseted distance matrix is the identity matrix.

        More details: :meth:`DistanceCalculator.calc_dist_mat`"""
        if seqs2 is None:
            # In this case, the offsetted distance matrix is the identity matrix
            return scipy.sparse.identity(len(seqs), dtype=self.DTYPE, format="coo")
        else:
            # actually compare the values
            def coord_generator():
                for (i1, s1), (i2, s2) in itertools.product(
                    enumerate(seqs), enumerate(seqs2)
                ):
                    if s1 == s2:
                        yield 1, i1, i2

            try:
                d, row, col = zip(*coord_generator())
            except ValueError:
                # happens when there is no match at all
                d, row, col = (), (), ()

            return coo_matrix(
                (d, (row, col)), dtype=self.DTYPE, shape=(len(seqs), len(seqs2))
            )


@_doc_params(params=_doc_params_parallel_distance_calculator)
class LevenshteinDistanceCalculator(ParallelDistanceCalculator):
    """\
    Calculates the Levenshtein edit-distance between sequences.

    The edit distance is the total number of deletion, addition and modification
    events.

    This class relies on `Python-levenshtein <https://github.com/ztane/python-Levenshtein>`_
    to calculate the distances.

    Choosing a cutoff:
        Each modification stands for a deletion, addition or modification event.
        While lacking empirical data, it seems unlikely that CDR3 sequences with more
        than two modifications still recognize the same antigen.

    Parameters
    ----------
    {params}
    """

    def _compute_block(self, seqs1, seqs2, origin):
        origin_row, origin_col = origin
        if seqs2 is not None:
            # compute the full matrix
            coord_iterator = itertools.product(enumerate(seqs1), enumerate(seqs2))
        else:
            # compute only upper triangle in this case
            coord_iterator = itertools.combinations_with_replacement(
                enumerate(seqs1), r=2
            )

        result = []
        for (row, s1), (col, s2) in coord_iterator:
            d = levenshtein_dist(s1, s2)
            if d <= self.cutoff:
                result.append((d + 1, origin_row + row, origin_col + col))

        return result


@_doc_params(params=_doc_params_parallel_distance_calculator)
class HammingDistanceCalculator(ParallelDistanceCalculator):
    """\
    Calculates the Hamming distance between sequences of identical length.

    The edit distance is the total number of substitution events. Sequences
    with different lengths will be treated as though they exceeded the
    distance-cutoff, i.e. they receive a distance of `0` in the sparse distance
    matrix and will not be connected by an edge in the graph.

    This class relies on `Python-levenshtein <https://github.com/ztane/python-Levenshtein>`_
    to calculate the distances.

    Choosing a cutoff:
        Each modification stands for a substitution event.
        While lacking empirical data, it seems unlikely that CDR3 sequences with more
        than two modifications still recognize the same antigen.

    Parameters
    ----------
    {params}
    """

    def _compute_block(self, seqs1, seqs2, origin):
        origin_row, origin_col = origin
        if seqs2 is not None:
            # compute the full matrix
            coord_iterator = itertools.product(enumerate(seqs1), enumerate(seqs2))
        else:
            # compute only upper triangle in this case
            coord_iterator = itertools.combinations_with_replacement(
                enumerate(seqs1), r=2
            )

        result = []
        for (row, s1), (col, s2) in coord_iterator:

            # require identical length of sequences
            if len(s1) != len(s2):
                continue
            d = hamming_dist(s1, s2)
            if d <= self.cutoff:
                result.append((d + 1, origin_row + row, origin_col + col))

        return result


@_doc_params(params=_doc_params_parallel_distance_calculator)
class AlignmentDistanceCalculator(ParallelDistanceCalculator):
    """\
    Calculates distance between sequences based on pairwise sequence alignment.

    The distance between two sequences is defined as :math:`S_{{1,2}}^{{max}} - S_{{1,2}}`,
    where :math:`S_{{1,2}}` is the alignment score of sequences 1 and 2 and
    :math:`S_{{1,2}}^{{max}}` is the max. achievable alignment score of sequences 1 and 2.
    :math:`S_{{1,2}}^{{max}}` is defined as :math:`\\min(S_{{1,1}}, S_{{2,2}})`.

    The use of alignment-based distances is heavily inspired by :cite:`TCRdist`.

    High-performance sequence alignments are calculated leveraging
    the `parasail library <https://github.com/jeffdaily/parasail-python>`_ (:cite:`Daily2016`).

    Choosing a cutoff:
        Alignment distances need to be viewed in the light of the substitution matrix.
        The alignment distance is the difference between the actual alignment
        score and the max. achievable alignment score. For instance, a mutation
        from *Leucine* (`L`) to *Isoleucine* (`I`) results in a BLOSUM62 score of `2`.
        An `L` aligned with `L` achieves a score of `4`. The distance is, therefore, `2`.

        On the other hand, a single *Tryptophane* (`W`) mutating into, e.g.
        *Proline* (`P`) already results in a distance of `15`.

        We are still lacking empirical data up to which distance a CDR3 sequence still
        is likely to recognize the same antigen, but reasonable cutoffs are `<15`.

    Parameters
    ----------
    {params}
    subst_mat
        Name of parasail substitution matrix
    gap_open
        Gap open penalty
    gap_extend
        Gap extend penatly
    """

    def __init__(
        self,
        cutoff: float,
        n_jobs: Union[int, None] = None,
        block_size: int = 50,
        *,
        subst_mat: str = "blosum62",
        gap_open: int = 11,
        gap_extend: int = 11,
    ):
        super().__init__(cutoff, n_jobs, block_size)
        self.subst_mat = subst_mat
        self.gap_open = gap_open
        self.gap_extend = gap_extend

    def _compute_block(self, seqs1, seqs2, origin):
        subst_mat = parasail.Matrix(self.subst_mat)
        origin_row, origin_col = origin

        square_matrix = seqs2 is None
        if square_matrix:
            seqs2 = seqs1

        self_alignment_scores1 = self._self_alignment_scores(seqs1)
        if square_matrix:
            self_alignment_scores2 = self_alignment_scores1
        else:
            self_alignment_scores2 = self._self_alignment_scores(seqs2)

        result = []
        for row, s1 in enumerate(seqs1):
            col_start = row if square_matrix else 0
            for col, s2 in enumerate(seqs2[col_start:], start=col_start):
                profile = parasail.profile_create_16(s1, subst_mat)
                r = parasail.nw_scan_profile_16(
                    profile, s2, self.gap_open, self.gap_extend
                )
                max_score = np.min(
                    [self_alignment_scores1[row], self_alignment_scores2[col]]
                )
                d = max_score - r.score
                if d <= self.cutoff:
                    result.append((d + 1, origin_row + row, origin_col + col))

        return result

    def _self_alignment_scores(self, seqs: Sequence) -> dict:
        """Calculate self-alignments. We need them as reference values
        to turn scores into dists"""
        return np.fromiter(
            (
                parasail.nw_scan_16(
                    s,
                    s,
                    self.gap_open,
                    self.gap_extend,
                    parasail.Matrix(self.subst_mat),
                ).score
                for s in seqs
            ),
            dtype=int,
            count=len(seqs),
        )


@deprecated(
    "Due to added BCR support, this function has been renamed "
    "to `sequence_dist`. The old version will be removed in a future release. "
)
def tcr_dist(*args, **kwargs):
    return sequence_dist(*args, **kwargs)


@_doc_params(metric=_doc_metrics, cutoff=_doc_cutoff, dist_mat=_doc_dist_mat)
def sequence_dist(
    unique_seqs: np.ndarray,
    unique_seqs2: Union[None, np.ndarray] = None,
    *,
    metric: Union[
        Literal["alignment", "identity", "levenshtein", "hamming"],
        DistanceCalculator,
    ] = "identity",
    cutoff: float = 10,
    n_jobs: Union[int, None] = None,
) -> coo_matrix:
    """\
    Calculate a sequence x sequence distance matrix.

    {dist_mat}

    Parameters
    ----------
    unique_seqs
        Numpy array of nucleotide or amino acid sequences.
        Must not contain duplicates.
        Note that not all distance metrics support nucleotide sequences.
    unique_seqs2
        Second array sequences. When omitted, `sequence_dist` computes
        the square matrix of `unique_seqs`.
    {metric}
    {cutoff}

        A cutoff of 0 implies the `identity` metric.

    Returns
    -------
    Upper triangular distance matrix when only `unique_seqs` is provided.
    A full rectangular distance matrix when both `unique_seqs` and
    `unique_seqs2` are provided.
    """
    if cutoff == 0 or metric == "identity":
        metric = "identity"
        cutoff = 0
    if isinstance(metric, DistanceCalculator):
        dist_calc = metric
    elif metric == "alignment":
        dist_calc = AlignmentDistanceCalculator(cutoff=cutoff, n_jobs=n_jobs)
    elif metric == "identity":
        dist_calc = IdentityDistanceCalculator(cutoff=cutoff)
    elif metric == "levenshtein":
        dist_calc = LevenshteinDistanceCalculator(cutoff=cutoff, n_jobs=n_jobs)
    elif metric == "hamming":
        dist_calc = HammingDistanceCalculator(cutoff=cutoff, n_jobs=n_jobs)
    else:
        raise ValueError("Invalid distance metric.")

    dist_mat = dist_calc.calc_dist_mat(unique_seqs, unique_seqs2)
    return dist_mat


@deprecated(
    "Due to added BCR support, this function has been renamed "
    "to `IrNeighbors`. The old version will be removed in a future release. "
)
def TcrNeighbors(*args, **kwargs):
    return IrNeighbors(*args, **kwargs)


class IrNeighbors:
    def __init__(
        self,
        adata: AnnData,
        *,
        metric: Union[
            Literal["alignment", "identity", "levenshtein", "hamming"],
            DistanceCalculator,
        ] = "identity",
        cutoff: float = 10,
        receptor_arms: Literal["VJ", "VDJ", "all", "any"] = "all",
        dual_ir: Literal["primary_only", "all", "any"] = "primary_only",
        sequence: Literal["aa", "nt"] = "aa",
    ):
        """Class to compute Neighborhood graphs of CDR3 sequences.

        For documentation of the parameters, see :func:`ir_neighbors`.
        """
        start = logging.info("Initializing IrNeighbors object...")
        if metric == "identity" and cutoff != 0:
            raise ValueError("Identity metric only works with cutoff == 0")
        if metric != "identity" and cutoff == 0:
            logging.warning(f"Running with {metric} metric, but cutoff == 0. ")
        if sequence == "nt" and metric == "alignment":
            raise ValueError(
                "Using nucleotide sequences with alignment metric is not supported. "
            )
        if receptor_arms not in ["VJ", "VDJ", "all", "any"]:
            raise ValueError(
                "Invalid value for `receptor_arms`. Note that starting with v0.5 "
                "`TRA` and `TRB` are not longer valid values."
            )
        if dual_ir not in ["primary_only", "all", "any"]:
            raise ValueError("Invalid value for `dual_ir")
        if sequence not in ["aa", "nt"]:
            raise ValueError("Invalid value for `sequence`")
        self.adata = adata
        self.metric = metric
        self.cutoff = cutoff
        self.receptor_arms = receptor_arms
        self.dual_ir = dual_ir
        self.sequence = sequence
        self._build_index_dict()
        self._dist_mat = None
        logging.info("Finished initalizing IrNeighbors object. ", time=start)

    @staticmethod
    def _seq_to_cell_idx(
        unique_seqs: np.ndarray, cdr_seqs: np.ndarray
    ) -> Dict[int, List[int]]:
        """
        Compute sequence to cell index for a single chain (e.g. `TRA_1`).

        Maps cell_idx -> [list, of, seq_idx].
        Useful to build a cell x cell matrix from a seq x seq matrix.

        Computes magic lookup indexes in linear time.

        Parameters
        ----------
        unique_seqs
            Pool of all unique cdr3 sequences (length = #unique cdr3 sequences)
        cdr_seqs
            CDR3 sequences for the current chain (length = #cells)

        Returns
        -------
        Sequence2Cell mapping
        """
        # 1) reverse mapping of amino acid sequence to index in sequence-distance matrix
        seq_to_index = {seq: i for i, seq in enumerate(unique_seqs)}

        # 2) indices of cells in adata that have a CDR3 sequence.
        cells_with_chain = np.where(~_is_na(cdr_seqs))[0]

        # 3) indices of the corresponding sequences in the distance matrix.
        seq_inds = {
            chain_id: seq_to_index[cdr_seqs[chain_id]] for chain_id in cells_with_chain
        }

        # 4) list of cell-indices in the cell distance matrix for each sequence
        seq_to_cell = {seq_id: list() for seq_id in seq_to_index.values()}
        for cell_id in cells_with_chain:
            seq_id = seq_inds[cell_id]
            seq_to_cell[seq_id].append(cell_id)

        return seq_to_cell

    def _build_index_dict(self):
        """Build nested dictionary for each receptor arm (TRA, TRB) containing all
        combinations of receptor_arms x primary/secondary_chain

        If the merge mode for either `receptor_arm` or `dual_ir` is `all`,
        includes a lookup table that contains the number of CDR3 sequences for
        each cell.
        """
        receptor_arms = (
            ["VJ", "VDJ"]
            if self.receptor_arms not in ["VJ", "VDJ"]
            else [self.receptor_arms]
        )
        chain_inds = [1] if self.dual_ir == "primary_only" else [1, 2]
        sequence = "" if self.sequence == "aa" else "_nt"

        arm_dict = {}
        for arm in receptor_arms:
            cdr_seqs = {
                k: self.adata.obs[f"IR_{arm}_{k}_cdr3{sequence}"].values
                for k in chain_inds
            }
            unique_seqs = np.hstack(list(cdr_seqs.values()))
            unique_seqs = np.unique(unique_seqs[~_is_na(unique_seqs)]).astype(str)
            seq_to_cell = {
                k: self._seq_to_cell_idx(unique_seqs, cdr_seqs[k]) for k in chain_inds
            }
            arm_dict[arm] = {
                "chain_inds": chain_inds,
                "unique_seqs": unique_seqs,
                "seq_to_cell": seq_to_cell,
            }

            # need the count of chains per cell for the `all` strategies.
            if self.receptor_arms == "all" or self.dual_ir == "all":
                arm_dict[arm]["chains_per_cell"] = np.sum(
                    ~_is_na(
                        self.adata.obs.loc[
                            :, [f"IR_{arm}_{k}_cdr3{sequence}" for k in chain_inds]
                        ]
                    ),
                    axis=1,
                )

        self.index_dict = arm_dict

    def _reduce_dual_all(self, d, chain, cell_row, cell_col):
        """Reduce dual IRs into a single value when 'all' sequences
        need to match. This requires additional checking effort for the number
        of chains in the given cell, since we can't make the distinction between
        no chain and dist > cutoff based on the distances (both would contain a
        0 in the distance matrix)."""
        chain_count = (
            self.index_dict[chain]["chains_per_cell"][cell_row],
            self.index_dict[chain]["chains_per_cell"][cell_col],
        )
        if len(d) == 1 and chain_count == (1, 1):
            # exactely one chain for both cells -> return that value
            return next(iter(d.values()))
        elif chain_count == (2, 2):
            # two options: either (1 matches 2 and 2 matches 1)
            # or (1 matches 1 and 2 matches 2).
            try:
                # minus 1, because both dists are offseted by 1.
                d1 = d[(1, 2)] + d[(2, 1)] - 1
            except KeyError:
                d1 = None
            try:
                d2 = d[(1, 1)] + d[(2, 2)] - 1
            except KeyError:
                d2 = None

            if d1 is not None and d2 is not None:
                return min(d1, d2)
            elif d1 is not None:
                return d1
            elif d2 is not None:
                return d2
            else:
                return 0

        else:
            return 0

    def _reduce_arms_all(self, values, cell_row, cell_col):
        """Reduce multiple receptor arms into a single value when 'all' sequences
        need to match. This requires additional checking effort for teh number
        of chains in the given cell, since we can't make the distinction between
        no chain and dist > cutoff based on the distances (both would contain a
        0 in the distance matrix)."""
        values = (x for x in values if x != 0)

        try:
            arm1 = next(values)
        except StopIteration:
            # no value > 0
            return 0

        try:
            arm2 = next(values)
            # two receptor arms -> easy
            # -1 because both distances are offseted by 1
            return arm1 + arm2 - 1
        except StopIteration:
            # only one arm
            tra_chains = (
                self.index_dict["VJ"]["chains_per_cell"][cell_row],
                self.index_dict["VJ"]["chains_per_cell"][cell_col],
            )
            trb_chains = (
                self.index_dict["VDJ"]["chains_per_cell"][cell_row],
                self.index_dict["VDJ"]["chains_per_cell"][cell_col],
            )
            # Either exactely one chain for VJ or
            # exactely one chain for VDJ for both cells.
            if tra_chains == (0, 0) or trb_chains == (0, 0):
                return arm1
            else:
                return 0

    @staticmethod
    def _reduce_arms_any(lst, *args):
        """Reduce arms when *any* of the sequences needs to match.
        This is the simpler case. This also works with only one entry
        (e.g. arms = "VJ")"""
        # need to exclude 0 values, since the dist mat is offseted by 1.
        try:
            return min(x for x in lst if x != 0)
        except ValueError:
            # no values in generator
            return 0

    @staticmethod
    def _reduce_dual_any(d, *args):
        """Reduce dual irs to a single value when *any* of the sequences needs
        to match (by minimum). This also works with only one entry (i.e. 'primary only')
        """
        # need to exclude 0 values, since the dist mat is offseted by 1.
        try:
            return min(x for x in d.values() if x != 0)
        except ValueError:
            # no values in generator
            return 0

    def _reduce_coord_dict(self, coord_dict):
        """Applies reduction functions to the coord dict.
        Yield (coords, value) pairs."""
        start = logging.info("Constructing cell x cell distance matrix...")
        reduce_dual = (
            self._reduce_dual_all if self.dual_ir == "all" else self._reduce_dual_any
        )
        reduce_arms = (
            self._reduce_arms_all
            if self.receptor_arms == "all"
            else self._reduce_arms_any
        )
        for (cell_row, cell_col), entry in tqdm(
            coord_dict.items(), total=len(coord_dict)
        ):
            reduced_dual = (
                reduce_dual(value_dict, chain, cell_row, cell_col)
                for chain, value_dict in entry.items()
            )
            reduced = reduce_arms(
                reduced_dual,
                cell_row,
                cell_col,
            )
            yield (cell_row, cell_col), reduced
        logging.info("Finished constructing cell x cell distance matrix. ", time=start)

    def _cell_dist_mat_reduce(self):
        """Compute the distance matrix by using custom reduction functions.
        More flexible than `_build_cell_dist_mat_min`, but requires more memory.
        Reduce dual is called before reduce arms.
        """
        coord_dict = dict()

        def _add_to_dict(d, c1, c2, cell_row, cell_col, value):
            """Add a value to the nested coord dict"""
            try:
                tmp_dict = d[(cell_row, cell_col)]
                try:
                    tmp_dict2 = tmp_dict[arm]
                    try:
                        if (c1, c2) in tmp_dict2:
                            # can be in arbitrary order apprarently
                            assert (c2, c1) not in tmp_dict2
                            tmp_dict2[(c2, c1)] = value
                        tmp_dict2[(c1, c2)] = value
                    except KeyError:
                        tmp_dict2 = {(c1, c2): value}
                except KeyError:
                    tmp_dict[arm] = {(c1, c2): value}
            except KeyError:
                d[(cell_row, cell_col)] = {arm: {(c1, c2): value}}

        for arm, arm_info in self.index_dict.items():
            dist_mat, seq_to_cell, chain_inds = (
                arm_info["dist_mat"],
                arm_info["seq_to_cell"],
                arm_info["chain_inds"],
            )
            start = logging.info(f"Started comstructing {arm} coord-dictionary...")
            for row, col, value in tqdm(
                zip(dist_mat.row, dist_mat.col, dist_mat.data), total=dist_mat.nnz
            ):
                for c1, c2 in itertools.product(chain_inds, repeat=2):
                    for cell_row, cell_col in itertools.product(
                        seq_to_cell[c1][row], seq_to_cell[c2][col]
                    ):
                        # fill upper diagonal. Important: these are dist-mat row,cols
                        # not cell-mat row cols. This is required, because the
                        # itertools.product returns all combinations for the diagonal
                        # but not for the other values.
                        _add_to_dict(coord_dict, c1, c2, cell_row, cell_col, value)
                        if row != col:
                            _add_to_dict(coord_dict, c1, c2, cell_col, cell_row, value)

            logging.info(f"Finished constructing {arm} coord-dictionary", time=start)

        yield from self._reduce_coord_dict(coord_dict)

    def compute_distances(
        self,
        n_jobs: Union[int, None] = None,
    ):
        """Computes the distances between CDR3 sequences

        Parameters
        ----------
        n_jobs
            Number of CPUs to use for alignment, levenshtein or hamming distance.
            Default: use all CPUS.
        """
        for arm, arm_dict in self.index_dict.items():
            start = logging.info(f"Computing {arm} pairwise distances...")

            arm_dict["dist_mat"] = sequence_dist(
                arm_dict["unique_seqs"],
                metric=self.metric,
                cutoff=self.cutoff,
                n_jobs=n_jobs,
            )
            logging.info(f"Finished computing {arm} pairwise distances.", time=start)

        try:
            coords, values = zip(*self._cell_dist_mat_reduce())
        except ValueError:
            # happens when there are no distances computed (e.g. all CDR3s are NAN)
            # -> construct empty CSR matrix
            self._dist_mat = csr_matrix((self.adata.n_obs, self.adata.n_obs))
        else:
            rows, cols = zip(*coords)
            dist_mat = coo_matrix(
                (values, (rows, cols)), shape=(self.adata.n_obs, self.adata.n_obs)
            )
            dist_mat.eliminate_zeros()
            self._dist_mat = dist_mat.tocsr()

    @property
    def dist(self):
        """The computed distance matrix.
        Requires to invoke `compute_distances() first."""
        return self._dist_mat

    @property
    def connectivities(self):
        """Get the weighted adjacecency matrix derived from the distance matrix.

        The cutoff will be used to normalize the distances.
        """
        if self.cutoff == 0:
            return self._dist_mat

        start = logging.debug("Started converting distances to connectivities. ")

        connectivities = self._dist_mat.copy()

        # actual distances
        d = connectivities.data - 1

        # structure of the matrix stayes the same, we can safely change the data only
        connectivities.data = (self.cutoff - d) / self.cutoff
        connectivities.eliminate_zeros()
        logging.debug("Finished converting distances to connectivities. ", time=start)
        return connectivities


@_doc_params(metric=_doc_metrics, cutoff=_doc_cutoff, dist_mat=_doc_dist_mat)
def ir_neighbors(
    adata: AnnData,
    *,
    metric: Union[
        Literal["identity", "alignment", "levenshtein", "hamming"],
        DistanceCalculator,
    ] = "identity",
    cutoff: int = 10,
    receptor_arms: Literal["VJ", "VDJ", "all", "any"] = "all",
    dual_ir: Literal["primary_only", "any", "all"] = "primary_only",
    key_added: Union[str, None] = None,
    sequence: Literal["aa", "nt"] = "nt",
    inplace: bool = True,
    n_jobs: Union[int, None] = None,
) -> Union[Tuple[csr_matrix, csr_matrix], None]:
    """\
    Construct a neighborhood graph based on :term:`CDR3` sequence similarity.

    All cells with a CDR3 distance `< cutoff` receive an edge in the graph.
    Edges are weighted by the distance.

    Parameters
    ----------
    adata
        annotated data matrix
    {metric}
    {cutoff}

        Two cells with a distance <= the cutoff will be connected.
        A cutoff of 0 implies the use of the `identity` metric.

    receptor_arms:
         * `"TRA"` - only consider TRA sequences
         * `"TRB"` - only consider TRB sequences
         * `"all"` - both TRA and TRB need to match
         * `"any"` - either TRA or TRB need to match

    dual_ir:
         * `"primary_only"` - only consider most abundant pair of TRA/TRB chains
         * `"any"` - consider both pairs of TRA/TRB sequences. Distance must be below
           cutoff for any of the chains.
         * `"all"` - consider both pairs of TRA/TRB sequences. Distance must be below
           cutoff for all of the chains.

        See also :term:`Dual IR`

    key_added:
        dict key under which the result will be stored in `adata.uns`
        when `inplace` is True. Defaults to `ir_neighbors_{{sequence}}_{{metric}}`.

        If metric is an instance of :class:`scirpy.ir_dist.DistanceCalculator`,
        `{{metric}}` defaults to `"custom"`.
    sequence:
        Use amino acid (`aa`) or nulceotide (`nt`) sequences?
    inplace:
        If `True`, store the results in `adata.uns`. Otherwise return
        the results.
    n_jobs:
        Number of cores to use for alignment, levenshtein or hamming distance.

    Returns
    -------
    connectivities
        weighted adjacency matrix
    dist
        cell x cell distance matrix with the distances as computed according to `metric`
        offsetted by 1 to make use of sparse matrices.
    """
    if cutoff == 0 or metric == "identity":
        metric = "identity"
        cutoff = 0
    if key_added is None:
        tmp_metric = "custom" if isinstance(metric, DistanceCalculator) else metric
        key_added = f"ir_neighbors_{sequence}_{tmp_metric}"
    ad = IrNeighbors(
        adata,
        metric=metric,
        cutoff=cutoff,
        receptor_arms=receptor_arms,
        dual_ir=dual_ir,
        sequence=sequence,
    )
    ad.compute_distances(n_jobs)

    if not inplace:
        return ad.connectivities, ad.dist
    else:
        adata.uns[key_added] = dict()
        adata.uns[key_added]["params"] = {
            "metric": metric,
            "cutoff": cutoff,
            "dual_ir": dual_ir,
            "receptor_arms": receptor_arms,
        }
        adata.uns[key_added]["connectivities"] = ad.connectivities
        adata.uns[key_added]["distances"] = ad.dist
\ metric     You can choose one of the following metrics:       * `identity` -- 1 for identical sequences, 0 otherwise.         See :class:`~scirpy.ir_dist.IdentityDistanceCalculator`.         This metric implies a cutoff of 0.       * `levenshtein` -- Levenshtein edit distance.         See :class:`~scirpy.ir_dist.LevenshteinDistanceCalculator`.       * `hamming` -- Hamming distance for CDR3 sequences of equal length.         See :class:`~scirpy.ir_dist.HammingDistanceCalculator`.       * `alignment` -- Distance based on pairwise sequence alignments using the         BLOSUM62 matrix. This option is incompatible with nucleotide sequences.         See :class:`~scirpy.ir_dist.AlignmentDistanceCalculator`.       * any instance of :class:`~scirpy.ir_dist.DistanceCalculator`. \ cutoff     All distances `> cutoff` will be replaced by `0` and eliminated from the sparse     matrix. A sensible cutoff depends on the distance metric, you can find     information in the corresponding docs. \ cutoff     Will eleminate distances > cutoff to make efficient     use of sparse matrices. \ n_jobs     Number of jobs to use for the pairwise distance calculation.     If None, use all jobs (only for ParallelDistanceCalculators). block_size     The width of a block of the matrix that will be delegated to a worker     process. The block contains `block_size ** 2` elements. \ Calculates the upper triangle, including the diagonal.  .. important::     * Distances are offset by 1 to allow efficient use of sparse matrices       (:math:`d' = d+1`).     * That means, a `distance > cutoff` is represented as `0`, a `distance == 0`       is represented as `1`, a `distance == 1` is represented as `2` and so on.     * Only returns distances `<= cutoff`. Larger distances are eliminated       from the sparse matrix.     * Distances are non-negative. \     Abstract base class for a :term:`CDR3`-sequence distance calculator.      Parameters     ----------     {params} #: The sparse matrix dtype. Defaults to uint8, constraining the max distance to 255. \         Calculate pairwise distance matrix of all sequences in `seqs` and `seqs2`.          When `seqs2` is omitted, computes the pairwise distance of `seqs` against         itself.          {dist_mat}          Parameters         ----------         seqs             array containing CDR3 sequences. Must not contain duplicates.         seqs2             second array containing CDR3 sequences. Must not contain             duplicates either.          Returns         -------         Sparse pairwise distance matrix. If only `seqs` is provided, only         the upper triangular distance matrix is returned. Abstract base class for a DistanceCalculator that computes distances in parallel.      It does so in a blockwise fashion. The function computing distances     for a single block needs to be overriden.      Parameters     ----------     {params} Compute the distances for a block of the matrix          Parameters         ----------         seqs1             array containing sequences         seqs2             other array containing sequences. If `None` compute the square matrix             of `seqs1` and iteratoe over the upper triangle including the diagonal only.         origin             row, col coordinates of the origin of the block.          Returns         ------         List of (distance, row, col) tuples for all elements with distance != 0.         row, col must be the coordinates in the final matrix (they can be derived using         `origin`). Can't be a generator because this needs to be picklable. Iterate over sequences in blocks.          Parameters         ----------         seqs1             array containing (unique) sequences         seqs2             array containing other sequences. If `None` compute             the square matrix of `seqs1` and iterate over the upper triangle (including             the diagonal) only.         block_size             side length of a block (will have `block_size ** 2` elements.)          Yields         ------         seqs1             subset of length `block_size` of seqs1         seqs2             subset of length `block_size` of seqs2. If seqs2 is None, this will             be `None` if the block is on the diagonal, or a subset of seqs1 otherwise.         origin             (row, col) coordinates of the origin of the block. # block on the diagonal. # yield None for seqs2 to indicate that we only want the upper # diagonal. Calculate the distance matrix.          See :meth:`DistanceCalculator.calc_dist_mat`. # precompute blocks as list to have total number of blocks for progressbar # happens when there is no match at all \     Calculates the Identity-distance between :term:`CDR3` sequences.      The identity distance is defined as         * `0`, if sequences are identical         * `1`, if sequences are not identical.      Choosing a cutoff:         For this DistanceCalculator, per definition, the cutoff = 0.         The `cutoff` argument is ignored.      Parameters     ----------     {params} In this case, the offseted distance matrix is the identity matrix.          More details: :meth:`DistanceCalculator.calc_dist_mat` # In this case, the offsetted distance matrix is the identity matrix # actually compare the values # happens when there is no match at all \     Calculates the Levenshtein edit-distance between sequences.      The edit distance is the total number of deletion, addition and modification     events.      This class relies on `Python-levenshtein <https://github.com/ztane/python-Levenshtein>`_     to calculate the distances.      Choosing a cutoff:         Each modification stands for a deletion, addition or modification event.         While lacking empirical data, it seems unlikely that CDR3 sequences with more         than two modifications still recognize the same antigen.      Parameters     ----------     {params} # compute the full matrix # compute only upper triangle in this case \     Calculates the Hamming distance between sequences of identical length.      The edit distance is the total number of substitution events. Sequences     with different lengths will be treated as though they exceeded the     distance-cutoff, i.e. they receive a distance of `0` in the sparse distance     matrix and will not be connected by an edge in the graph.      This class relies on `Python-levenshtein <https://github.com/ztane/python-Levenshtein>`_     to calculate the distances.      Choosing a cutoff:         Each modification stands for a substitution event.         While lacking empirical data, it seems unlikely that CDR3 sequences with more         than two modifications still recognize the same antigen.      Parameters     ----------     {params} # compute the full matrix # compute only upper triangle in this case # require identical length of sequences \     Calculates distance between sequences based on pairwise sequence alignment.      The distance between two sequences is defined as :math:`S_{{1,2}}^{{max}} - S_{{1,2}}`,     where :math:`S_{{1,2}}` is the alignment score of sequences 1 and 2 and     :math:`S_{{1,2}}^{{max}}` is the max. achievable alignment score of sequences 1 and 2.     :math:`S_{{1,2}}^{{max}}` is defined as :math:`\\min(S_{{1,1}}, S_{{2,2}})`.      The use of alignment-based distances is heavily inspired by :cite:`TCRdist`.      High-performance sequence alignments are calculated leveraging     the `parasail library <https://github.com/jeffdaily/parasail-python>`_ (:cite:`Daily2016`).      Choosing a cutoff:         Alignment distances need to be viewed in the light of the substitution matrix.         The alignment distance is the difference between the actual alignment         score and the max. achievable alignment score. For instance, a mutation         from *Leucine* (`L`) to *Isoleucine* (`I`) results in a BLOSUM62 score of `2`.         An `L` aligned with `L` achieves a score of `4`. The distance is, therefore, `2`.          On the other hand, a single *Tryptophane* (`W`) mutating into, e.g.         *Proline* (`P`) already results in a distance of `15`.          We are still lacking empirical data up to which distance a CDR3 sequence still         is likely to recognize the same antigen, but reasonable cutoffs are `<15`.      Parameters     ----------     {params}     subst_mat         Name of parasail substitution matrix     gap_open         Gap open penalty     gap_extend         Gap extend penatly Calculate self-alignments. We need them as reference values         to turn scores into dists \     Calculate a sequence x sequence distance matrix.      {dist_mat}      Parameters     ----------     unique_seqs         Numpy array of nucleotide or amino acid sequences.         Must not contain duplicates.         Note that not all distance metrics support nucleotide sequences.     unique_seqs2         Second array sequences. When omitted, `sequence_dist` computes         the square matrix of `unique_seqs`.     {metric}     {cutoff}          A cutoff of 0 implies the `identity` metric.      Returns     -------     Upper triangular distance matrix when only `unique_seqs` is provided.     A full rectangular distance matrix when both `unique_seqs` and     `unique_seqs2` are provided. Class to compute Neighborhood graphs of CDR3 sequences.          For documentation of the parameters, see :func:`ir_neighbors`. Compute sequence to cell index for a single chain (e.g. `TRA_1`).          Maps cell_idx -> [list, of, seq_idx].         Useful to build a cell x cell matrix from a seq x seq matrix.          Computes magic lookup indexes in linear time.          Parameters         ----------         unique_seqs             Pool of all unique cdr3 sequences (length = #unique cdr3 sequences)         cdr_seqs             CDR3 sequences for the current chain (length = #cells)          Returns         -------         Sequence2Cell mapping # 1) reverse mapping of amino acid sequence to index in sequence-distance matrix # 2) indices of cells in adata that have a CDR3 sequence. # 3) indices of the corresponding sequences in the distance matrix. # 4) list of cell-indices in the cell distance matrix for each sequence Build nested dictionary for each receptor arm (TRA, TRB) containing all         combinations of receptor_arms x primary/secondary_chain          If the merge mode for either `receptor_arm` or `dual_ir` is `all`,         includes a lookup table that contains the number of CDR3 sequences for         each cell. # need the count of chains per cell for the `all` strategies. Reduce dual IRs into a single value when 'all' sequences         need to match. This requires additional checking effort for the number         of chains in the given cell, since we can't make the distinction between         no chain and dist > cutoff based on the distances (both would contain a         0 in the distance matrix). # exactely one chain for both cells -> return that value # two options: either (1 matches 2 and 2 matches 1) # or (1 matches 1 and 2 matches 2). # minus 1, because both dists are offseted by 1. Reduce multiple receptor arms into a single value when 'all' sequences         need to match. This requires additional checking effort for teh number         of chains in the given cell, since we can't make the distinction between         no chain and dist > cutoff based on the distances (both would contain a         0 in the distance matrix). # no value > 0 # two receptor arms -> easy # -1 because both distances are offseted by 1 # only one arm # Either exactely one chain for VJ or # exactely one chain for VDJ for both cells. Reduce arms when *any* of the sequences needs to match.         This is the simpler case. This also works with only one entry         (e.g. arms = "VJ") # need to exclude 0 values, since the dist mat is offseted by 1. # no values in generator Reduce dual irs to a single value when *any* of the sequences needs         to match (by minimum). This also works with only one entry (i.e. 'primary only') # need to exclude 0 values, since the dist mat is offseted by 1. # no values in generator Applies reduction functions to the coord dict.         Yield (coords, value) pairs. Compute the distance matrix by using custom reduction functions.         More flexible than `_build_cell_dist_mat_min`, but requires more memory.         Reduce dual is called before reduce arms. Add a value to the nested coord dict # can be in arbitrary order apprarently # fill upper diagonal. Important: these are dist-mat row,cols # not cell-mat row cols. This is required, because the # itertools.product returns all combinations for the diagonal # but not for the other values. Computes the distances between CDR3 sequences          Parameters         ----------         n_jobs             Number of CPUs to use for alignment, levenshtein or hamming distance.             Default: use all CPUS. # happens when there are no distances computed (e.g. all CDR3s are NAN) # -> construct empty CSR matrix The computed distance matrix.         Requires to invoke `compute_distances() first. Get the weighted adjacecency matrix derived from the distance matrix.          The cutoff will be used to normalize the distances. # actual distances # structure of the matrix stayes the same, we can safely change the data only \     Construct a neighborhood graph based on :term:`CDR3` sequence similarity.      All cells with a CDR3 distance `< cutoff` receive an edge in the graph.     Edges are weighted by the distance.      Parameters     ----------     adata         annotated data matrix     {metric}     {cutoff}          Two cells with a distance <= the cutoff will be connected.         A cutoff of 0 implies the use of the `identity` metric.      receptor_arms:          * `"TRA"` - only consider TRA sequences          * `"TRB"` - only consider TRB sequences          * `"all"` - both TRA and TRB need to match          * `"any"` - either TRA or TRB need to match      dual_ir:          * `"primary_only"` - only consider most abundant pair of TRA/TRB chains          * `"any"` - consider both pairs of TRA/TRB sequences. Distance must be below            cutoff for any of the chains.          * `"all"` - consider both pairs of TRA/TRB sequences. Distance must be below            cutoff for all of the chains.          See also :term:`Dual IR`      key_added:         dict key under which the result will be stored in `adata.uns`         when `inplace` is True. Defaults to `ir_neighbors_{{sequence}}_{{metric}}`.          If metric is an instance of :class:`scirpy.ir_dist.DistanceCalculator`,         `{{metric}}` defaults to `"custom"`.     sequence:         Use amino acid (`aa`) or nulceotide (`nt`) sequences?     inplace:         If `True`, store the results in `adata.uns`. Otherwise return         the results.     n_jobs:         Number of cores to use for alignment, levenshtein or hamming distance.      Returns     -------     connectivities         weighted adjacency matrix     dist         cell x cell distance matrix with the distances as computed according to `metric`         offsetted by 1 to make use of sparse matrices.