# sage.doctest: needs sage.geometry.polyhedron
r"""
Elements of bounded height in number fields

This module provides functions to list all elements of a given number field
with height less than a specified bound.

REFERENCES:

- [DK2013]_

AUTHORS:

- John Doyle, David Krumm (2013): initial version
- TJ Combs, Raghukul Raman (2018): added Doyle-Krumm algorithm-4
"""
# ****************************************************************************
#       Copyright (C) 2013 John Doyle and David Krumm
#
#  Distributed under the terms of the GNU General Public License (GPL)
#  as published by the Free Software Foundation; either version 2 of
#  the License, or (at your option) any later version.
#                  https://www.gnu.org/licenses/
# ****************************************************************************

from copy import copy
from itertools import product
from sage.rings.real_mpfr import RealField
from sage.rings.number_field.unit_group import UnitGroup
from sage.modules.free_module_element import vector
from sage.matrix.constructor import column_matrix
from sage.rings.rational_field import QQ
from sage.rings.infinity import Infinity
from sage.rings.real_mpfr import RR
from sage.geometry.polyhedron.constructor import Polyhedron


def bdd_norm_pr_gens_iq(K, norm_list):
    r"""
    Compute generators for all principal ideals in an imaginary quadratic field
    `K` whose norms are in ``norm_list``.

    The only keys for the output dictionary are integers n appearing in
    ``norm_list``.

    The function will only be called with `K` an imaginary quadratic field.

    The function will return a dictionary for other number fields, but it may be
    incorrect.

    INPUT:

    - ``K`` -- an imaginary quadratic number field

    - ``norm_list`` -- list of positive integers

    OUTPUT: dictionary of number field elements, keyed by norm

    EXAMPLES:

    In `QQ(i)`, there is one principal ideal of norm 4, two principal ideals of
    norm 5, but no principal ideals of norm 7::

        sage: from sage.rings.number_field.bdd_height import bdd_norm_pr_gens_iq
        sage: x = polygen(ZZ, 'x')
        sage: K.<g> = NumberField(x^2 + 1)
        sage: L = range(10)
        sage: bdd_pr_ideals = bdd_norm_pr_gens_iq(K, L)
        sage: bdd_pr_ideals[4]
        [2]
        sage: bdd_pr_ideals[5]
        [-g - 2, -g + 2]
        sage: bdd_pr_ideals[7]
        []

    There are no ideals in the ring of integers with negative norm::

        sage: from sage.rings.number_field.bdd_height import bdd_norm_pr_gens_iq
        sage: K.<g> = NumberField(x^2 + 10)
        sage: L = range(-5,-1)
        sage: bdd_pr_ideals = bdd_norm_pr_gens_iq(K,L)
        sage: bdd_pr_ideals
        {-5: [], -4: [], -3: [], -2: []}

    Calling a key that is not in the input ``norm_list`` raises a KeyError::

        sage: from sage.rings.number_field.bdd_height import bdd_norm_pr_gens_iq
        sage: K.<g> = NumberField(x^2 + 20)
        sage: L = range(100)
        sage: bdd_pr_ideals = bdd_norm_pr_gens_iq(K, L)
        sage: bdd_pr_ideals[100]
        Traceback (most recent call last):
        ...
        KeyError: 100
    """
    return {n: K.elements_of_norm(n) for n in norm_list}


def bdd_height_iq(K, height_bound):
    r"""
    Compute all elements in the imaginary quadratic field `K` which have
    relative multiplicative height at most ``height_bound``.

    The function will only be called with `K` an imaginary quadratic field.

    If called with `K` not an imaginary quadratic, the function will likely
    yield incorrect output.

    ALGORITHM:

    This is an implementation of Algorithm 5 in [DK2013]_.

    INPUT:

    - ``K`` -- an imaginary quadratic number field

    - ``height_bound`` -- a real number

    OUTPUT: an iterator of number field elements

    EXAMPLES::

        sage: from sage.rings.number_field.bdd_height import bdd_height_iq
        sage: x = polygen(ZZ, 'x')
        sage: K.<a> = NumberField(x^2 + 191)
        sage: for t in bdd_height_iq(K,8):
        ....:     print(exp(2*t.global_height()))
        1.00000000000000
        1.00000000000000
        1.00000000000000
        4.00000000000000
        4.00000000000000
        4.00000000000000
        4.00000000000000
        8.00000000000000
        8.00000000000000
        8.00000000000000
        8.00000000000000
        8.00000000000000
        8.00000000000000
        8.00000000000000
        8.00000000000000

    There are 175 elements of height at most 10 in `QQ(\sqrt(-3))`::

        sage: from sage.rings.number_field.bdd_height import bdd_height_iq
        sage: K.<a> = NumberField(x^2 + 3)
        sage: len(list(bdd_height_iq(K,10)))
        175

    The only elements of multiplicative height 1 in a number field are 0 and
    the roots of unity::

        sage: from sage.rings.number_field.bdd_height import bdd_height_iq
        sage: K.<a> = NumberField(x^2 + x + 1)
        sage: list(bdd_height_iq(K,1))
        [0, a + 1, a, -1, -a - 1, -a, 1]

    A number field has no elements of multiplicative height less than 1::

        sage: from sage.rings.number_field.bdd_height import bdd_height_iq
        sage: K.<a> = NumberField(x^2 + 5)
        sage: list(bdd_height_iq(K,0.9))
        []
    """
    if height_bound < 1:
        return
    yield K(0)
    roots_of_unity = K.roots_of_unity()
    for zeta in roots_of_unity:
        yield zeta

    # Get a complete set of ideal class representatives
    class_group_reps = []
    class_group_rep_norms = []
    for c in K.class_group():
        a = c.ideal()
        class_group_reps.append(a)
        class_group_rep_norms.append(a.norm())
    class_number = len(class_group_reps)

    # Find principal ideals of bounded norm
    possible_norm_set = set()
    for n in range(class_number):
        for m in range(1, int(height_bound + 1)):
            possible_norm_set.add(m * class_group_rep_norms[n])
    bdd_ideals = bdd_norm_pr_gens_iq(K, possible_norm_set)

    # Distribute the principal ideals
    generator_lists = []
    for n in range(class_number):
        this_ideal = class_group_reps[n]
        this_ideal_norm = class_group_rep_norms[n]
        gens = [g for i in range(1, int(height_bound + 1))
                for g in bdd_ideals[i * this_ideal_norm] if g in this_ideal]
        generator_lists.append(gens)

    # Build all the output numbers
    for n in range(class_number):
        gens = generator_lists[n]
        s = len(gens)
        for i in range(s):
            for j in range(i + 1, s):
                if K.ideal(gens[i], gens[j]) == class_group_reps[n]:
                    new_number = gens[i] / gens[j]
                    for zeta in roots_of_unity:
                        yield zeta * new_number
                        yield zeta / new_number


def bdd_norm_pr_ideal_gens(K, norm_list):
    r"""
    Compute generators for all principal ideals in a number field `K` whose
    norms are in ``norm_list``.

    INPUT:

    - ``K`` -- a number field

    - ``norm_list`` -- list of positive integers

    OUTPUT: dictionary of number field elements, keyed by norm

    EXAMPLES:

    There is only one principal ideal of norm 1, and it is generated by the
    element 1::

        sage: from sage.rings.number_field.bdd_height import bdd_norm_pr_ideal_gens
        sage: K.<g> = QuadraticField(101)
        sage: bdd_norm_pr_ideal_gens(K, [1])
        {1: [1]}

    In this example the output differs slightly on 32- and 64-bit
    machines, but in both cases the norm of `g \pm 11` is the same::

        sage: from sage.rings.number_field.bdd_height import (
        ....:     bdd_norm_pr_ideal_gens
        ....: )
        sage: K.<g> = QuadraticField(123)
        sage: d = bdd_norm_pr_ideal_gens(K, range(5))
        sage: d  # needs 32_bit
        {0: [0], 1: [1], 2: [g - 11], 3: [], 4: [2]}
        sage: d  # needs !32_bit
        {0: [0], 1: [1], 2: [g + 11], 3: [], 4: [2]}

    ::

        sage: from sage.rings.number_field.bdd_height import bdd_norm_pr_ideal_gens
        sage: x = polygen(ZZ, 'x')
        sage: K.<g> = NumberField(x^5 - x + 19)
        sage: b = bdd_norm_pr_ideal_gens(K, range(30))
        sage: key = ZZ(28)
        sage: b[key]
        [157*g^4 - 139*g^3 - 369*g^2 + 848*g + 158, g^4 + g^3 - g - 7]
    """
    negative_norm_units = K.elements_of_norm(-1)
    gens = {}
    if not negative_norm_units:
        for n in norm_list:
            if not n:
                gens[n] = [K.zero()]
            else:
                gens[n] = K.elements_of_norm(n) + K.elements_of_norm(-n)
    else:
        for n in norm_list:
            gens[n] = K.elements_of_norm(n)
    return gens


def integer_points_in_polytope(matrix, interval_radius):
    r"""
    Return the set of integer points in the polytope obtained by acting on a
    cube by a linear transformation.

    Given an `r`-by-`r` matrix ``matrix`` and a real number ``interval_radius``,
    this function finds all integer lattice points in the polytope obtained by
    transforming the cube ``[-interval_radius, interval_radius]^r`` via the linear
    map induced by ``matrix``.

    INPUT:

    - ``matrix`` -- a square matrix of real numbers

    - ``interval_radius`` -- a real number

    OUTPUT: list of tuples of integers

    EXAMPLES:

    Stretch the interval `[-1,1]` by a factor of 2 and find the integers in the
    resulting interval::

        sage: from sage.rings.number_field.bdd_height import integer_points_in_polytope
        sage: m = matrix([2])
        sage: r = 1
        sage: integer_points_in_polytope(m, r)
        [(-2), (-1), (0), (1), (2)]

    Integer points inside a parallelogram::

        sage: from sage.rings.number_field.bdd_height import integer_points_in_polytope
        sage: m = matrix([[1, 2], [3, 4]])
        sage: r = RealField()(1.3)
        sage: integer_points_in_polytope(m, r)
        [(-3, -7), (-2, -5), (-2, -4), (-1, -3), (-1, -2), (-1, -1), (0, -1),
         (0, 0), (0, 1), (1, 1), (1, 2), (1, 3), (2, 4), (2, 5), (3, 7)]

    Integer points inside a parallelepiped::

        sage: from sage.rings.number_field.bdd_height import integer_points_in_polytope
        sage: m = matrix([[1.2,3.7,0.2], [-5.3,-.43,3], [1.2,4.7,-2.1]])
        sage: r = 2.2
        sage: L = integer_points_in_polytope(m, r)
        sage: len(L)
        4143

    If ``interval_radius`` is 0, the output should include only the zero tuple::

        sage: from sage.rings.number_field.bdd_height import integer_points_in_polytope
        sage: m = matrix([[1,2,3,7], [4,5,6,2], [7,8,9,3], [0,3,4,5]])
        sage: integer_points_in_polytope(m, 0)
        [(0, 0, 0, 0)]
    """
    T = matrix
    d = interval_radius
    r = T.nrows()

    # Find the vertices of the given box
    box_vertices = [vector(x) for x in product([-d, d], repeat=r)]

    # Transform the vertices
    T_trans = T.transpose()
    transformed_vertices = [v * T_trans for v in box_vertices]

    # Create polyhedron from transformed vertices and find integer points inside
    return list(Polyhedron(transformed_vertices, base_ring=QQ).integral_points())


def bdd_height(K, height_bound, tolerance=1e-2, precision=53):
    r"""
    Compute all elements in the number field `K` which have relative
    multiplicative height at most ``height_bound``.

    The function can only be called for number fields `K` with positive unit
    rank. An error will occur if `K` is `\QQ` or an imaginary quadratic field.

    This algorithm computes 2 lists: `L`, containing elements `x` in `K` such that
    `H_k(x) \leq B`, and a list `L'` containing elements `x` in `K` that, due to
    floating point issues,
    may be slightly larger then the bound. This can be controlled
    by lowering the tolerance.

    In current implementation both lists `(L,L')` are merged and returned in
    form of iterator.

    ALGORITHM:

    This is an implementation of the revised algorithm (Algorithm 4) in
    [DK2013]_.

    INPUT:

    - ``height_bound`` -- real number

    - ``tolerance`` -- (default: 0.01) a rational number in (0,1]

    - ``precision`` -- (default: 53) positive integer

    OUTPUT: an iterator of number field elements

    EXAMPLES:

    There are no elements of negative height::

        sage: from sage.rings.number_field.bdd_height import bdd_height
        sage: x = polygen(ZZ, 'x')
        sage: K.<g> = NumberField(x^5 - x + 7)
        sage: list(bdd_height(K, -3))
        []

    The only nonzero elements of height 1 are the roots of unity::

        sage: from sage.rings.number_field.bdd_height import bdd_height
        sage: K.<g> = QuadraticField(3)
        sage: list(bdd_height(K, 1))
        [0, -1, 1]

    ::

        sage: from sage.rings.number_field.bdd_height import bdd_height
        sage: K.<g> = QuadraticField(36865)
        sage: len(list(bdd_height(K, 101))) # long time (4 s)
        131

    ::

        sage: from sage.rings.number_field.bdd_height import bdd_height
        sage: K.<g> = NumberField(x^6 + 2)
        sage: len(list(bdd_height(K, 60))) # long time (5 s)
        1899

    ::

        sage: from sage.rings.number_field.bdd_height import bdd_height
        sage: K.<g> = NumberField(x^4 - x^3 - 3*x^2 + x + 1)
        sage: len(list(bdd_height(K, 10)))
        99

    TESTS:

    Check that :issue:`22771` is fixed::

        sage: from sage.rings.number_field.bdd_height import bdd_height
        sage: K.<v> = NumberField(x^3 + x + 1)
        sage: len(list(bdd_height(K,3)))
        23
    """
    # global values, used in internal function
    B = height_bound
    theta = tolerance
    if B < 1:
        return
    embeddings = K.places(prec=precision)
    r1, r2 = K.signature()
    r = r1 + r2 - 1
    RF = RealField(precision)
    lambda_gens_approx = {}
    class_group_rep_norm_log_approx = []
    unit_log_dict = {}

    def rational_in(x, y):
        r"""
        Compute a rational number q, such that x<q<y using Archimedes' axiom
        """
        z = y - x
        if z == 0:
            n = 1
        else:
            n = RR(1/z).ceil() + 1
        if RR(n*y).ceil() is n*y:  # WHAT !?
            m = n*y - 1
        else:
            m = RR(n*y).floor()
        return m / n

    def delta_approximation(x, delta):
        r"""
        Compute a rational number in range `(x-\delta, x+\delta)`
        """
        return rational_in(x - delta, x + delta)

    def vector_delta_approximation(v, delta):
        r"""
        Compute a rational vector `w=(w_1, ..., w_n)`
        such that `|v_i-w_i|<\delta` for all `i` in `[1, n]`
        """
        return [delta_approximation(vi, delta) for vi in v]

    def log_map(number):
        r"""
        Compute the image of an element of `K` under the logarithmic map.
        """
        x = number
        x_logs = []
        for i in range(r1):
            sigma = embeddings[i]  # real embeddings
            x_logs.append(sigma(x).abs().log())
        for i in range(r1, r + 1):
            tau = embeddings[i]  # Complex embeddings
            x_logs.append(2 * tau(x).abs().log())
        return vector(x_logs)

    def log_height_for_generators_approx(alpha, beta, Lambda):
        r"""
        Compute the rational approximation of logarithmic height function.
        Return a lambda approximation h_K(alpha/beta)
        """
        delta = Lambda / (r + 2)
        norm_log = delta_approximation(RR(K.ideal(alpha, beta).norm()).log(), delta)
        log_ga = vector_delta_approximation(log_map(alpha), delta)
        log_gb = vector_delta_approximation(log_map(beta), delta)
        arch_sum = sum([max(log_ga[k], log_gb[k]) for k in range(r + 1)])
        return (arch_sum - norm_log)

    def packet_height(n, pair, u):
        r"""
        Compute the height of the element of `K` encoded by a given packet.
        """
        gens = generator_lists[n]
        i = pair[0]
        j = pair[1]
        Log_gi = lambda_gens_approx[gens[i]]
        Log_gj = lambda_gens_approx[gens[j]]
        Log_u_gi = vector(Log_gi) + unit_log_dict[u]
        arch_sum = sum([max(Log_u_gi[k], Log_gj[k]) for k in range(r + 1)])
        return (arch_sum - class_group_rep_norm_log_approx[n])

    # Step 1
    # Computes ideal class representative and their rational approx norm
    t = theta / (3*B)
    delta_1 = t / (6*r+12)

    class_group_reps = []
    class_group_rep_norms = []

    for c in K.class_group():
        a = c.ideal()
        a_norm = a.norm()
        log_norm = RF(a_norm).log()
        log_norm_approx = delta_approximation(log_norm, delta_1)
        class_group_reps.append(a)
        class_group_rep_norms.append(a_norm)
        class_group_rep_norm_log_approx.append(log_norm_approx)
    class_number = len(class_group_reps)

    # Step 2
    # Find generators for principal ideals of bounded norm
    possible_norm_set = set()
    for n in range(class_number):
        for m in range(1, (B + 1).ceil()):
            possible_norm_set.add(m * class_group_rep_norms[n])
    bdd_ideals = bdd_norm_pr_ideal_gens(K, possible_norm_set)

    # Stores it in form of a dictionary and gives lambda(g)_approx for key g
    for norm in possible_norm_set:
        gens = bdd_ideals[norm]
        for g in gens:
            lambda_g_approx = vector_delta_approximation(log_map(g), delta_1)
            lambda_gens_approx[g] = lambda_g_approx

    # Step 3
    # Find a list of all generators corresponding to each ideal a_l
    generator_lists = []
    for l in range(class_number):
        this_ideal = class_group_reps[l]
        this_ideal_norm = class_group_rep_norms[l]
        gens = []
        for i in range(1, (B + 1).ceil()):
            for g in bdd_ideals[i * this_ideal_norm]:
                if g in this_ideal:
                    gens.append(g)
        generator_lists.append(gens)

    # Step 4
    # Finds all relevant pair and their height
    gen_height_approx_dictionary = {}
    relevant_pair_lists = []

    for n in range(class_number):
        relevant_pairs = []
        gens = generator_lists[n]
        l = len(gens)
        for i in range(l):
            for j in range(i+1, l):
                if K.ideal(gens[i], gens[j]) == class_group_reps[n]:
                    relevant_pairs.append([i, j])
                    gen_height_approx_dictionary[(n, i, j)] = log_height_for_generators_approx(gens[i], gens[j], t/6)
        relevant_pair_lists.append(relevant_pairs)

    # Step 5
    b = rational_in(t/12 + RR(B).log(), t/4 + RR(B).log())
    maximum = 0
    for n in range(class_number):
        for p in relevant_pair_lists[n]:
            maximum = max(maximum, gen_height_approx_dictionary[(n, p[0], p[1])])
    d_tilde = b + t/6 + maximum

    # Step 6
    # computes fundamental units and their value under log map
    fund_units = UnitGroup(K).fundamental_units()
    fund_unit_logs = [log_map(fund_units[i]) for i in range(r)]
    S = column_matrix(fund_unit_logs).delete_rows([r])
    S_inverse = S.inverse()
    S_norm = S.norm(Infinity)
    S_inverse_norm = S_inverse.norm(Infinity)

    upper_bound = (r**2) * max(S_norm, S_inverse_norm)
    m = RR(upper_bound).ceil() + 1

    # Step 7
    # Variables needed for rational approximation
    lambda_tilde = (t/12) / (d_tilde*r*(1+m))
    delta_tilde = min(lambda_tilde/((r**2)*((m**2)+m*lambda_tilde)), 1/(r**2))
    M = d_tilde * (upper_bound+lambda_tilde*RR(r).sqrt())
    M = RR(M).ceil()
    d_tilde = RR(d_tilde)
    delta_2 = min(delta_tilde, (t/6)/(r*(r+1)*M))

    # Step 8, 9
    # Computes relevant points in polytope
    fund_unit_log_approx = [vector_delta_approximation(fund_unit_logs[i], delta_2) for i in range(r)]
    S_tilde = column_matrix(fund_unit_log_approx).delete_rows([r])
    S_tilde_inverse = S_tilde.inverse()
    U = integer_points_in_polytope(S_tilde_inverse, d_tilde)

    # Step 10
    # tilde suffixed list are used for computing second list (L_primed)
    yield K(0)
    U0 = []
    U0_tilde = []
    L0 = []
    L0_tilde = []

    # Step 11
    # Computes unit height
    unit_height_dict = {}
    U_copy = copy(U)
    inter_bound = b - (5*t)/12

    for u in U:
        u_log = sum([u[j]*vector(fund_unit_log_approx[j]) for j in range(r)])
        unit_log_dict[u] = u_log
        u_height = sum([max(u_log[k], 0) for k in range(r + 1)])
        unit_height_dict[u] = u_height
        if u_height < inter_bound:
            U0.append(u)
        if inter_bound <= u_height < b - (t/12):
            U0_tilde.append(u)
        if u_height > t/12 + d_tilde:
            U_copy.remove(u)
    U = U_copy

    relevant_tuples = set(U0 + U0_tilde)

    # Step 12
    # check for relevant packets
    for n in range(class_number):
        for pair in relevant_pair_lists[n]:
            i = pair[0]
            j = pair[1]
            u_height_bound = b + gen_height_approx_dictionary[(n, i, j)] + t/4
            for u in U:
                if unit_height_dict[u] < u_height_bound:
                    candidate_height = packet_height(n, pair, u)
                    if candidate_height <= b - 7*t/12:
                        L0.append([n, pair, u])
                        relevant_tuples.add(u)
                    elif candidate_height < b + t/4:
                        L0_tilde.append([n, pair, u])
                        relevant_tuples.add(u)

    # Step 13
    # forms a dictionary of all_unit_tuples and their value
    tuple_to_unit_dict = {}
    for u in relevant_tuples:
        unit = K.one()
        for k in range(r):
            unit *= fund_units[k]**u[k]
        tuple_to_unit_dict[u] = unit

    # Step 14
    # Build all output numbers
    roots_of_unity = K.roots_of_unity()
    for u in U0 + U0_tilde:
        for zeta in roots_of_unity:
            yield zeta * tuple_to_unit_dict[u]

    # Step 15
    for p in L0 + L0_tilde:
        gens = generator_lists[p[0]]
        i = p[1][0]
        j = p[1][1]
        u = p[2]
        c_p = tuple_to_unit_dict[u] * (gens[i] / gens[j])
        for zeta in roots_of_unity:
            yield zeta * c_p
            yield zeta / c_p
