Лаборатория математических методов и моделей в биоинформатике
Института проблем передачи информации им. А.А. Харкевича
Российской академии наук

from typing import List, Tuple

import numpy as np


def dist(x: str, y: str) -> int:
    return palindrome_alignment(x, y)[0]


def imp(x: str) -> float:
    return palindrome_self_alignment(x)[0] / len(x)


def palindrome_alignment(x: str, y: str, loop: bool = False) -> Tuple[int, List[int]]:
    """For a given pair of nucleotide sequences `x`, `y`, find the optimal
    alignment between x and ww' (or yw' if `loop` is `True`) where w is
    some prefix of y and w' is the reverse complement of w.
    Return the edit distance and the list of optimal w lengths.
    >>> palindrome_alignment("ACGT", "ACC")
    (0, [2])
    >>> palindrome_alignment("ACTG", "ACC", loop=True)
    (1, [1, 2])
    """
    x = preprocess_nucleotide_sequence(x)
    y = preprocess_nucleotide_sequence(y)
    forward = prefix_edit_distances(x, y)
    reverse = prefix_edit_distances(reverse_complement(x), y)
    if loop:
        forward = forward[:, -1:]
    total = forward + reverse[::-1]
    minima = total.min(axis=0)
    edit_distance = minima.min()
    prefix_lengths = np.where(minima == edit_distance)[0].tolist()
    return int(edit_distance), prefix_lengths


def palindrome_self_alignment(x: str) -> Tuple[int, List[int]]:
    """A special case of `palindrome_alignment` for self-alignment.
    For a given nucleotide sequence `x`, find the optimal alignment
    between x and ww' where w is some prefix of x and w' is the
    reverse complement of w.
    Return the edit distance and the list of optimal w lengths.
    >>> palindrome_self_alignment("ACGC")
    (1, [2])
    """
    x = preprocess_nucleotide_sequence(x)
    forward = np.abs(np.arange(len(x) + 1) - np.arange(len(x) + 1)[:, None])
    reverse = prefix_edit_distances(reverse_complement(x), x)
    total = forward + reverse[::-1]
    minima = total.min(axis=0)
    edit_distance = minima.min()
    prefix_lengths = np.where(minima == edit_distance)[0].tolist()
    return int(edit_distance), prefix_lengths


def preprocess_nucleotide_sequence(s: str) -> str:
    """Preprocess a nucleotide sequence for the Needleman-Wunsch algorithm."""
    s = s.upper().replace("U", "T").replace("-", "")
    if set(s) - set("ATCG"):
        raise ValueError(f"invalid nucleotide sequence: {s}")
    return s


def reverse_complement(s: str) -> str:
    """Return the reverse complement of a nucleotide sequence."""
    return s[::-1].translate(str.maketrans("ATCGatcg", "TAGCtagc"))


def prefix_edit_distances(s0: str, s1: str) -> np.ndarray:
    """Return the matrix of edit distances between all prefixes of `s0`
    and all prefixes of `s1` indexed by the lengths of the prefixes."""
    n0, n1 = len(s0), len(s1)
    distance = np.zeros((n0 + 1, n1 + 1), np.int32)
    distance[:, 0] = np.arange(n0 + 1)
    distance[0, :] = np.arange(n1 + 1)
    for i in range(1, n0 + 1):
        for j in range(1, n1 + 1):
            distance[i, j] = min(
                distance[i, j - 1] + 1,
                distance[i - 1, j] + 1,
                distance[i - 1, j - 1] + (s0[i - 1] != s1[j - 1]),
            )
    return distance