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

from fractions import Fraction
from math import floor
from typing import Sequence, Union

Number = Union[int, Fraction]
Basis = list[list[Number]]


def _dot(u: Sequence[Number], v: Sequence[Number]) -> Number:
    return sum((a * b for a, b in zip(u, v)), 0)


def _inverse(a: list[list[Fraction]]) -> list[list[Fraction]]:
    """Exact inverse by Gauss-Jordan elimination; ValueError if a is singular."""
    n = len(a)
    m = [list(row) + [Fraction(int(i == j)) for j in range(n)] for i, row in enumerate(a)]
    for c in range(n):
        p = next((r for r in range(c, n) if m[r][c] != 0), None)
        if p is None:
            raise ValueError("the rows are linearly dependent")
        m[c], m[p] = m[p], m[c]
        piv = m[c][c]
        m[c] = [x / piv for x in m[c]]
        for r in range(n):
            if r != c and m[r][c] != 0:
                f = m[r][c]
                m[r] = [x - f * y for x, y in zip(m[r], m[c])]
    return [row[n:] for row in m]


def seysen_measure(basis: Sequence[Sequence[Number]]) -> Fraction:
    """S(B) = sum_i ||b_i||^2 ||d_i||^2, where d_i is the dual basis."""
    b = [list(r) for r in basis]
    g = [[Fraction(_dot(u, v)) for v in b] for u in b]
    h = _inverse(g)
    return sum((g[i][i] * h[i][i] for i in range(len(b))), Fraction(0))


def seysen_reduce_counted(basis: Sequence[Sequence[Number]]) -> tuple[Basis, int]:
    """Seysen reduction with exact arithmetic; returns (reduced basis, number of updates).

    Input: n linearly independent rows of equal length, int or Fraction entries.
    Output: a basis of the same lattice such that for every ordered pair i != j,
    |lambda*_ij| <= 1/2, so no update b_j <- b_j + l b_i with integer l lowers S.
    Each update lowers S(B) det(G); for an integer basis this is a positive integer,
    so there are at most (S_0 - n) det(G) updates.
    Raises ValueError for ragged or linearly dependent rows.
    """
    b: Basis = [list(r) for r in basis]
    n = len(b)
    if n == 0:
        return [], 0
    if any(len(r) != len(b[0]) for r in b):
        raise ValueError("all rows must have the same length")
    g = [[Fraction(_dot(u, v)) for v in b] for u in b]  # @line L1
    h = _inverse(g)  # h[i][j] = <d_i, d_j>
    steps = 0
    while True:  # @line L2
        changed = False  # @line L3
        for i in range(n):  # @line L4
            for j in range(n):
                if i == j:
                    continue
                star = (h[i][j] / h[j][j] - g[i][j] / g[i][i]) / 2  # @line L5
                lam = floor(star + Fraction(1, 2))  # @line L6
                if lam * lam - 2 * lam * star < 0:  # @line L7
                    b[j] = [x + lam * y for x, y in zip(b[j], b[i])]  # @line L8
                    # @line L9: b_j += lam b_i and d_i -= lam d_j, on both Gram matrices
                    g[j] = [x + lam * y for x, y in zip(g[j], g[i])]
                    for r in range(n):
                        g[r][j] += lam * g[r][i]
                    h[i] = [x - lam * y for x, y in zip(h[i], h[j])]
                    for r in range(n):
                        h[r][i] -= lam * h[r][j]
                    steps += 1
                    changed = True  # @line L10
        if not changed:  # @line L11
            return b, steps  # @line L12


def seysen_reduce(basis: Sequence[Sequence[Number]]) -> Basis:
    """Seysen-reduced basis of the lattice generated by the rows of `basis`."""
    return seysen_reduce_counted(basis)[0]