"""Suffix array construction by prefix doubling.

The algorithm is the classic prefix-doubling scheme of Manber and Myers
(SIAM J. Comput. 22(5), 1993). It runs in O(n log^2 n) time and O(n)
auxiliary space, using a comparison sort at each of the O(log n) rounds.

Linear-time alternatives exist (DC3/skew, SA-IS) but are considerably more
intricate; prefix doubling is the right trade-off when clarity matters or
when inputs stay moderate in size.
"""

from __future__ import annotations

__all__ = ["build_suffix_array", "build_inverse_suffix_array"]


def build_suffix_array(text: str) -> list[int]:
    """Build the suffix array of ``text``.

    The suffix array is the permutation of ``range(len(text))`` that lists
    the starting positions of every suffix of ``text`` in increasing
    lexicographic order.

    The construction proceeds by prefix doubling. At the start of each
    round, ``rank[i]`` holds the rank of suffix ``i`` restricted to its
    first ``k`` characters. Sorting the suffixes on the pair
    ``(rank[i], rank[i + k])`` therefore orders them over ``2 * k``
    characters, since the second component is already the rank of the
    following block of length ``k``. Ranks are then recomputed and ``k``
    is doubled.

    Args:
        text: The string to index. May be empty. A terminal sentinel such
            as ``"$"`` is not required, but adding one guarantees that no
            suffix is a prefix of another, which downstream structures
            (BWT, FM-index, LCP array) usually expect.

    Returns:
        A list of ``len(text)`` distinct indices into ``text``, ordered so
        that ``text[sa[0]:] < text[sa[1]:] < ...``.

    Examples:
        >>> build_suffix_array("TAATAC$")
        [6, 1, 4, 2, 5, 0, 3]
        >>> build_suffix_array("banana")
        [5, 3, 1, 0, 4, 2]
        >>> build_suffix_array("")
        []
        >>> build_suffix_array("a")
        [0]
    """
    n = len(text)
    if n == 0:
        return []

    # Dense initial ranks. Mapping each distinct character onto 0..m-1
    # (rather than keeping raw code points) preserves the lexicographic
    # order while guaranteeing the invariant the early exit below relies
    # on: ranks always form a contiguous range starting at 0.
    alphabet = {char: code for code, char in enumerate(sorted(set(text)))}
    rank = [alphabet[char] for char in text]

    suffix_array = list(range(n))
    next_rank = [0] * n

    k = 1
    while k < n:
        # Materialise the sort keys once per round. Recomputing them inside
        # the comparison callback would evaluate each key O(log n) times
        # during the sort, then twice more during the rank scan.
        # The -1 sentinel marks a suffix shorter than 2 * k: having no
        # second half, it must sort before any suffix sharing its first
        # half, which is exactly the lexicographic rule for prefixes.
        keys = [(rank[i], rank[i + k] if i + k < n else -1) for i in range(n)]
        suffix_array.sort(key=lambda i: keys[i])

        # Recompute ranks by scanning the sorted order. The counter only
        # ever increments by one, so the resulting ranks span 0..c-1 where
        # c is the number of distinct keys. Equal keys share a rank.
        next_rank[suffix_array[0]] = 0
        for position in range(1, n):
            previous = suffix_array[position - 1]
            current = suffix_array[position]
            next_rank[current] = next_rank[previous] + (
                keys[current] != keys[previous]
            )

        # Swap rather than assign: the outgoing rank array is recycled as
        # the scratch buffer for the next round, where it is fully
        # overwritten before being read.
        rank, next_rank = next_rank, rank

        # Ranks being dense, a maximum of n - 1 means n distinct classes,
        # hence every suffix is already distinguishable and the order is
        # total. Further rounds would reproduce the same permutation.
        # Ranks increase along suffix_array, so the maximum sits last.
        if rank[suffix_array[-1]] == n - 1:
            break

        k <<= 1

    return suffix_array


def build_inverse_suffix_array(suffix_array: list[int]) -> list[int]:
    """Invert a suffix array.

    Args:
        suffix_array: A suffix array, as returned by
            :func:`build_suffix_array`.

    Returns:
        A list ``isa`` such that ``isa[i]`` is the position of suffix ``i``
        within ``suffix_array``, i.e. its lexicographic rank.

    Examples:
        >>> build_inverse_suffix_array([6, 1, 4, 2, 5, 0, 3])
        [5, 1, 3, 6, 2, 4, 0]
    """
    inverse = [0] * len(suffix_array)
    for position, index in enumerate(suffix_array):
        inverse[index] = position
    return inverse


if __name__ == "__main__":
    import doctest

    failures, _ = doctest.testmod()
    if failures == 0:
        print("All doctests passed.")