# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/f4-algorithm/. The authors are listed in prov/XY5Z26.jsonl.

"""Faugere's F4 algorithm over Q, with the normal selection strategy.

A polynomial in K[x_1, ..., x_n] is a dict from exponent tuples to nonzero
Fractions. A monomial order is one of "lex", "grlex", "grevlex" with
x_1 > x_2 > ... > x_n. f4(F, order) returns the reduced Groebner basis of the
ideal generated by F, sorted by leading monomial, and the list of matrix sizes
(rows, columns) of the rounds, for inspection.

Each round takes all critical pairs whose lcm has the smallest total degree,
builds the Macaulay matrix of their two halves together with the reducers found
by symbolic preprocessing, and brings it to row echelon form. The pairs are
maintained with the Gebauer-Moeller criteria.

Cost model: operations in Q. No useful worst-case bound is proved for F4 beyond
the general one for Groebner bases, whose degrees can be doubly exponential in n
(Mayr and Meyer 1982); one round with an r x c matrix costs O(r c min(r, c)).
"""

from fractions import Fraction
from typing import Callable, Dict, List, Set, Tuple

Mono = Tuple[int, ...]
Poly = Dict[Mono, Fraction]
Key = Callable[[Mono], Tuple[object, ...]]

ORDERS: Dict[str, Key] = {
    "lex": lambda a: tuple(a),
    "grlex": lambda a: (sum(a),) + tuple(a),
    "grevlex": lambda a: (sum(a),) + tuple(-e for e in reversed(a)),
}


def lm(f: Poly, key: Key) -> Mono:
    return max(f, key=key)


def divides(a: Mono, b: Mono) -> bool:
    return all(x <= y for x, y in zip(a, b))


def lcm(a: Mono, b: Mono) -> Mono:
    return tuple(max(x, y) for x, y in zip(a, b))


def coprime(a: Mono, b: Mono) -> bool:
    return all(x == 0 or y == 0 for x, y in zip(a, b))


def shift(f: Poly, u: Mono) -> Poly:
    return {tuple(a + b for a, b in zip(e, u)): c for e, c in f.items()}


def axpy(f: Poly, c: Fraction, g: Poly) -> Poly:
    """f - c g, as a new dict."""
    h = dict(f)
    for e, v in g.items():
        w = h.get(e, Fraction(0)) - c * v
        if w:
            h[e] = w
        else:
            h.pop(e, None)
    return h


def update(G: List[Poly], pairs: List[Tuple[int, int]], h: Poly, key: Key) -> None:
    """Gebauer-Moeller UPDATE: append h to G and maintain the critical pairs in place."""
    t = len(G)
    lh = lm(h, key)
    lms = [lm(g, key) for g in G]
    C = [i for i in range(t)]
    D: List[int] = []
    while C:
        i = C.pop()
        li = lcm(lms[i], lh)
        if coprime(lms[i], lh) or not any(divides(lcm(lms[k], lh), li) for k in C + D):
            D.append(i)
    E = [(i, t) for i in D if not coprime(lms[i], lh)]
    kept = []
    for i, j in pairs:
        lij = lcm(lms[i], lms[j])
        if not divides(lh, lij) or lcm(lms[i], lh) == lij or lcm(lms[j], lh) == lij:
            kept.append((i, j))
    G.append(h)
    pairs[:] = kept + E


def symbolic_preprocessing(L: List[Poly], G: List[Poly], key: Key) -> List[Poly]:
    """Add a reducer u g for every monomial of the rows divisible by some LM(g)."""
    H = list(L)
    done: Set[Mono] = {lm(f, key) for f in L}
    todo = {e for f in L for e in f} - done
    lms = [(lm(g, key), g) for g in G]
    while todo:
        m = todo.pop()
        done.add(m)
        for gm, g in lms:
            if divides(gm, m):
                r = shift(g, tuple(a - b for a, b in zip(m, gm)))
                H.append(r)
                todo |= {e for e in r if e not in done}
                break
    return H


def row_echelon(rows: List[Poly], key: Key) -> List[Poly]:
    """Semi-echelon basis of the span of rows: distinct leading monomials, monic."""
    pivots: Dict[Mono, Poly] = {}
    for r in rows:
        r = dict(r)
        while r:
            m = lm(r, key)
            if m not in pivots:
                break
            r = axpy(r, r[m], pivots[m])
        if r:
            m = lm(r, key)
            c = r[m]
            pivots[m] = {e: v / c for e, v in r.items()}
    return list(pivots.values())


def reduce_full(f: Poly, G: List[Poly], key: Key) -> Poly:
    p, r = dict(f), {}
    lms = [(lm(g, key), g) for g in G]
    while p:
        m = lm(p, key)
        for gm, g in lms:
            if divides(gm, m):
                p = axpy(p, p[m] / g[gm], shift(g, tuple(a - b for a, b in zip(m, gm))))
                break
        else:
            r[m] = p.pop(m)
    return r


def reduced_basis(G: List[Poly], key: Key) -> List[Poly]:
    minimal: List[Poly] = []
    for g in sorted(G, key=lambda g: key(lm(g, key))):
        if not any(divides(lm(h, key), lm(g, key)) for h in minimal):
            minimal.append(g)
    out = []
    for i, g in enumerate(minimal):
        m = lm(g, key)
        tail = {e: v / g[m] for e, v in g.items() if e != m}
        r = reduce_full(tail, minimal[:i] + minimal[i + 1:], key)
        r[m] = Fraction(1)
        out.append(r)
    return sorted(out, key=lambda g: key(lm(g, key)))


def f4(F: List[Poly], order: str) -> Tuple[List[Poly], List[Tuple[int, int]]]:
    """Reduced Groebner basis of <F> for the order, and the matrix sizes per round.

    Precondition: F is a list of polynomials over Q in the same number of variables.
    Cost: one round with an r x c Macaulay matrix costs O(r c min(r, c)) operations in Q.
    """
    key = ORDERS[order]
    G: List[Poly] = []
    pairs: List[Tuple[int, int]] = []
    for f in F:  # @line L1
        if f:
            update(G, pairs, {e: Fraction(c) for e, c in f.items() if c}, key)
    sizes: List[Tuple[int, int]] = []
    while pairs:  # @line L2
        degs = [sum(lcm(lm(G[i], key), lm(G[j], key))) for i, j in pairs]
        d = min(degs)  # @line L3
        chosen = [p for p, e in zip(pairs, degs) if e == d]
        pairs[:] = [p for p, e in zip(pairs, degs) if e != d]
        halves: Dict[Tuple[int, Mono], Poly] = {}  # @line L4
        for i, j in chosen:
            l = lcm(lm(G[i], key), lm(G[j], key))
            for k in (i, j):
                u = tuple(a - b for a, b in zip(l, lm(G[k], key)))
                halves[(k, u)] = shift(G[k], u)
        H = symbolic_preprocessing(list(halves.values()), G, key)  # @line L5
        sizes.append((len(H), len({e for r in H for e in r})))
        lead_H = {lm(r, key) for r in H}
        N = row_echelon(H, key)  # @line L6
        new = [r for r in N if lm(r, key) not in lead_H]  # @line L7
        for h in sorted(new, key=lambda r: key(lm(r, key))):  # @line L8
            update(G, pairs, h, key)
    return reduced_basis(G, key), sizes  # @line L9