"""Small F_p linear algebra helpers (standard library only).

Vectors are tuples of ints mod p.  Subspaces are kept in reduced row
echelon form (RREF), a canonical key.
"""


def rref(rows, p):
    """Return the RREF basis (tuple of tuples) of the span of rows over F_p."""
    rows = [list(r) for r in rows]
    if not rows:
        return tuple()
    n = len(rows[0])
    out = []
    r = 0
    piv_cols = []
    for c in range(n):
        piv = None
        for i in range(r, len(rows)):
            if rows[i][c] % p:
                piv = i
                break
        if piv is None:
            continue
        rows[r], rows[piv] = rows[piv], rows[r]
        inv = pow(rows[r][c], p - 2, p)
        rows[r] = [(x * inv) % p for x in rows[r]]
        for i in range(len(rows)):
            if i != r and rows[i][c] % p:
                f = rows[i][c]
                rows[i] = [(a - f * b) % p for a, b in zip(rows[i], rows[r])]
        piv_cols.append(c)
        r += 1
        if r == len(rows):
            break
    return tuple(tuple(rows[i]) for i in range(r))


def in_span(basis, v, p):
    """True if v lies in the span of the RREF basis."""
    v = list(v)
    for b in basis:
        c = next(i for i, x in enumerate(b) if x)
        if v[c] % p:
            f = v[c]
            v = [(a - f * bb) % p for a, bb in zip(v, b)]
    return all(x % p == 0 for x in v)


def nullspace(mat, n, p):
    """Basis of {v in F_p^n : mat v = 0}; mat is a list of rows of length n."""
    R = rref(mat, p) if mat else tuple()
    pivots = []
    for row in R:
        pivots.append(next(i for i, x in enumerate(row) if x))
    free = [c for c in range(n) if c not in pivots]
    basis = []
    for f in free:
        v = [0] * n
        v[f] = 1
        for row, pc in zip(R, pivots):
            v[pc] = (-row[f]) % p
        basis.append(tuple(v))
    return basis


def all_vectors(basis, p):
    """All vectors in the span of basis (generator)."""
    n = len(basis)
    if n == 0:
        return
    dim = len(basis[0])
    coeffs = [0] * n
    while True:
        v = [0] * dim
        for c, b in zip(coeffs, basis):
            if c:
                v = [(a + c * bb) % p for a, bb in zip(v, b)]
        yield tuple(v)
        i = 0
        while i < n:
            coeffs[i] += 1
            if coeffs[i] < p:
                break
            coeffs[i] = 0
            i += 1
        if i == n:
            return
