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

"""Buchberger's algorithm for a Groebner basis of a polynomial ideal over a field.

A polynomial in n variables is a dict that maps an exponent tuple (a1, ..., an)
to a nonzero coefficient. The zero polynomial is {}. Coefficients must be field
elements with +, -, *, / and == 0, for example Fraction or an element of Z/pZ.
The monomial order is "lex", "grlex" or "grevlex", with x1 > x2 > ... > xn.

Cost: the number of field operations has no polynomial bound; in the worst case
the degrees in the output are doubly exponential in n (Mayr and Meyer 1982).
"""

from typing import Any, Callable, Dict, List, Sequence, Tuple

Monomial = Tuple[int, ...]
Poly = Dict[Monomial, Any]


def order_key(order: str) -> Callable[[Monomial], Tuple[Any, ...]]:
    """Return a sort key: a > b in the order exactly when key(a) > key(b)."""
    if order == "lex":
        return lambda a: tuple(a)
    if order == "grlex":
        return lambda a: (sum(a), tuple(a))
    if order == "grevlex":
        return lambda a: (sum(a), tuple(-e for e in reversed(a)))
    raise ValueError("unknown monomial order: " + order)


def leading_monomial(f: Poly, key: Callable[[Monomial], Tuple[Any, ...]]) -> Monomial:
    """The largest monomial of a nonzero f."""
    return max(f, key=key)


def _add_multiple(p: Poly, g: Poly, mono: Monomial, c: Any) -> None:
    """In place: p <- p - c * mono * g."""
    for e, d in g.items():
        m = tuple(a + b for a, b in zip(e, mono))
        v = p.get(m, 0) - c * d
        if v == 0:
            p.pop(m, None)
        else:
            p[m] = v


def normal_form(f: Poly, G: Sequence[Poly], order: str) -> Poly:
    """The remainder of f on division by the list G (full reduction, iterative)."""
    key = order_key(order)
    leads = [(leading_monomial(g, key), g) for g in G if g]
    p = dict(f)
    r: Poly = {}
    while p:
        m = leading_monomial(p, key)
        c = p[m]
        for lm, g in leads:
            if all(a >= b for a, b in zip(m, lm)):
                q = tuple(a - b for a, b in zip(m, lm))
                _add_multiple(p, g, q, c / g[lm])
                break
        else:
            r[m] = c
            del p[m]
    return r


def s_polynomial(f: Poly, g: Poly, order: str) -> Poly:
    """S(f, g) = (x^gamma / LT(f)) f - (x^gamma / LT(g)) g, with x^gamma the lcm of LM(f), LM(g)."""
    key = order_key(order)
    a, b = leading_monomial(f, key), leading_monomial(g, key)
    lcm = tuple(max(x, y) for x, y in zip(a, b))
    s: Poly = {}
    _add_multiple(s, f, tuple(x - y for x, y in zip(lcm, a)), -1 / f[a])
    _add_multiple(s, g, tuple(x - y for x, y in zip(lcm, b)), 1 / g[b])
    return s


def buchberger(F: Sequence[Poly], order: str = "grevlex", product_criterion: bool = True) -> List[Poly]:
    """A Groebner basis of the ideal generated by F, for the given order.

    The result contains the nonzero elements of F, followed by the nonzero
    remainders that the algorithm adds. Pairs are taken by the normal strategy:
    the pair whose lcm is smallest in the order goes first.
    """
    key = order_key(order)
    G: List[Poly] = [dict(f) for f in F if f]  # @line L1
    lms = [leading_monomial(g, key) for g in G]
    B = [(i, j) for j in range(len(G)) for i in range(j)]  # @line L2
    while B:  # @line L3
        best = min(range(len(B)), key=lambda u: (key(tuple(max(x, y) for x, y in zip(lms[B[u][0]], lms[B[u][1]]))), B[u]))
        i, j = B.pop(best)  # @line L4
        if product_criterion and all(min(x, y) == 0 for x, y in zip(lms[i], lms[j])):  # @line L5
            continue
        h = normal_form(s_polynomial(G[i], G[j], order), G, order)  # @line L6
        if h:  # @line L7
            G.append(h)  # @line L8
            lms.append(leading_monomial(h, key))
            t = len(G) - 1
            B.extend((u, t) for u in range(t))  # @line L9
    return G  # @line L10


def reduced_basis(G: Sequence[Poly], order: str) -> List[Poly]:
    """The reduced Groebner basis from a Groebner basis G, sorted by leading monomial, largest first."""
    key = order_key(order)
    G = [g for g in G if g]
    lms = [leading_monomial(g, key) for g in G]
    keep: List[Poly] = []
    for i, g in enumerate(G):
        redundant = False
        for j in range(len(G)):
            if j != i and all(a >= b for a, b in zip(lms[i], lms[j])) and (lms[i] != lms[j] or j < i):
                redundant = True
                break
        if not redundant:
            keep.append(g)
    out: List[Poly] = []
    for i, g in enumerate(keep):
        lm = leading_monomial(g, key)
        lc = g[lm]
        tail = {m: c / lc for m, c in g.items() if m != lm}
        r = normal_form(tail, keep[:i] + keep[i + 1:], order)
        r[lm] = lc / lc
        out.append(r)
    out.sort(key=lambda p: key(leading_monomial(p, key)), reverse=True)
    return out