""" @brief  A module for Babai's nearest plane algorithm. """

import gmpy2;

from gmpy2 import mpz;
from gmpy2 import mpq;
from gmpy2 import mpfr;

from .gram_schmidt import gram_schmidt;
from ..matrices import dimensions;

def babai(
  B : list[list[int | mpz]],
  t : list[int | mpz | mpq],
  Bs : list[list[int | mpz | mpq | mpfr]] | None = None,
  precision : int | None = None) -> list[int | mpz]:

  """ @brief  Uses Babai's nearest plane algorithm [Babai86] to find a vector x
              that is in the lattice L, and that is close to a target vector t,
              where L is generated by the rows of a reduced basis matrix B.

      More specifically, it holds that | t - x | <= 2^(n/2) dist(t, L) if B is
      delta-LLL-reduced with delta >= 3/4.

      [Babai86] Babai, L.: "On Lovász’ lattice reduction and the nearest lattice
                            point problem". Combinatorica 6(1), pp. 1–13 (1986).

      The basis matrix B is represented as as list of lists [b_1, ..., b_n],
      where b_i = [b_i1, ..., b_id] for i = 1, ..., n represent the n row
      vectors that make up the basis.

      @param B  The n x d basis basis matrix B = [b_1, ..., b_n], where
                b_i = [b_i1, ..., b_id] for i = 1, ..., n represent the n row
                vectors that make up the basis.

      @param t  The target vector t = [t_1, ..., t_d].

      @param Bs   The Gram–Schmidt orthogonalization Bs of B, if available, or
                  None in which case the Gram–Schmidt orthogonalization B is
                  computed by this function. (If you plan on calling this
                  function several times for the same lattice L, then you may
                  wish to pre-compute Bs and pass it along in each call to this
                  function. Note also that the function lll() for computing
                  delta-LLL-reduced bases B optionally computes Bs alongside B.)

      @param precision  The precision to use when computing the Gram–Schmidt
                        projection factors in M. May be set to None, in which
                        case projection factors are represented as exact
                        quotients.

                        Note that this parameter only has an effect if gs is
                        set to None as the Gram–Schmidt orthogonalized basis Bs
                        and the associated matrix M of Gram–Schmidt projection
                        factors are otherwise pre-computed.

      @return   A vector x that is in the lattice L generated by B, and that is
                close to t. More specifically, | t - x | <= 2^(n/2) dist(t, L)
                if B is delta-LLL-reduced with delta >= 3/4. """

  # Initial setup.
  (n, d) = dimensions(B);

  if None == Bs:
    [Bs, _] = gram_schmidt(B, precision = precision);

  with gmpy2.context(gmpy2.get_context()) as context:

    # Set the precision.
    if precision != None:
      context.precision = precision;

    # Step 2.
    x = t.copy();

    # Step 3.
    for i in range(n, 0, -1):
      # Step 3.1.
      a = sum([x[j] * Bs[i-1][j] for j in range(d)]);
      b = sum([Bs[i-1][j] * Bs[i-1][j] for j in range(d)]);

      if None == precision:
        mu = mpq(a, b);
      else:
        mu = mpfr(a) / mpfr(b);

      # Step 3.2.
      mu_rounded = round(mu);

      for j in range(d):
        x[j] = x[j] - mu_rounded * B[i-1][j];

    # Step 4.
    for j in range(d):
      x[j] = t[j] - x[j];

    # Step 5.
    return x;