"""
Neighbors
"""

from sage.modules.free_module_element import vector
from sage.rings.integer_ring import ZZ
from sage.rings.all import GF, QQ
from copy import deepcopy
from sage.matrix.constructor import matrix


# ############################################################################
# Routines used for understanding p-neighbors and computing classes in a genus
# ############################################################################


def find_primitive_p_divisible_vector__random(self, p):
    """
    Find a random `p`-primitive vector in `L/pL` whose value is `p`-divisible.

    .. note::

        Since there are about `p^{(n-2)}` of these lines, we have a `1/p`
        chance of randomly finding an appropriate vector.

    EXAMPLES::

        sage: Q = QuadraticForm(ZZ, 2, [10,1,4])
        sage: v = Q.find_primitive_p_divisible_vector__random(5)    # random
        sage: v
        (3, 3)
        sage: 5.divides(Q(v))
        True
        sage: Q = QuadraticForm(QQ,matrix.diagonal([1,1,1,1]))
        sage: v = Q.find_primitive_p_divisible_vector__random(2)
        sage: Q(v)
        2
    """
    n = self.dim()
    v = vector([ZZ.random_element(p) for _ in range(n)])

    # Repeatedly choose random vectors, and evaluate until the value
    # is p-divisible.
    k = 0
    while k < 1000:
        k = k + 1
        a = self(v)
        if a in ZZ and (a % p == 0) and (v != 0):
            return v
        else:
            v[ZZ.random_element(n)] = ZZ.random_element(p)
            # Replace a random entry and try again.
    raise RuntimeError("unable to find a p divisible vector")


def find_primitive_p_divisible_vector__next(self, p, v=None):
    """
    Find the next `p`-primitive vector (up to scaling) in `L/pL` whose
    value is `p`-divisible, where the last vector returned was `v`.  For
    an initial call, no `v` needs to be passed.

    Returns vectors whose last non-zero entry is normalized to 0 or 1 (so no
    lines are counted repeatedly).  The ordering is by increasing the
    first non-normalized entry.  If we have tested all (lines of)
    vectors, then return None.

    OUTPUT:

    vector or None

    EXAMPLES::

        sage: Q = QuadraticForm(ZZ, 2, [10,1,4])
        sage: v = Q.find_primitive_p_divisible_vector__next(5); v
        (1, 1)
        sage: v = Q.find_primitive_p_divisible_vector__next(5, v); v
        (1, 0)
        sage: v = Q.find_primitive_p_divisible_vector__next(5, v); v
        sage: Q = QuadraticForm(QQ,matrix.diagonal([1,1,1,1]))
        sage: v = Q.find_primitive_p_divisible_vector__next(2)
        sage: Q(v)
        2
    """
    # Initialize
    n = self.dim()
    if v is None:
        w = vector(ZZ, [0] * (n - 1) + [1])
    else:
        w = deepcopy(v)

    # Handle n = 1 separately.
    if n <= 1:
        raise NotImplementedError("Sorry -- Not implemented yet!")

    # Look for the last non-zero entry (which must be 1)
    nz = n - 1
    while w[nz] == 0:
        nz += -1

    # Test that the last non-zero entry is 1 (to detect tampering).
    if w[nz] != 1:
        print("Warning: The input vector to QuadraticForm.find_primitive_p_divisible_vector__next() is not normalized properly.")

    # Look for the next vector, until w == 0
    while True:

        # Look for the first non-maximal (non-normalized) entry
        ind = 0
        while (ind < nz) and (w[ind] == p - 1):
            ind += 1

        # Increment
        if ind < nz:
            w[ind] += 1
            for j in range(ind):
                w[j] = 0
        else:
            for j in range(ind + 1):    # Clear all entries
                w[j] = 0

            if nz != 0:
                # Move the non-zero normalized index over by one, or
                # return the zero vector
                w[nz - 1] = 1
                nz += -1

        # Test for zero vector
        if w == 0:
            return None

        # Test for p-divisibility
        a = self(w)
        if a in ZZ and (a % p == 0):
            return w

def find_p_neighbor_from_vec(self, p, y):
    r"""
    Return the `p`-neighbor of ``self`` defined by ``y``.

    Let `(L,q)` be a lattice with `b(L,L) \subseteq \ZZ` which is maximal at `p`.
    Let `y \in L` with `b(y,y) \in p^2\ZZ` then the `p`-neighbor of
    `L` at `y` is given by
    `\ZZ y/p + L_y` where `L_y = \{x \in L | b(x,y) \in p \ZZ \}`
    and `b(x,y) = q(x+y)-q(x)-q(y)` is the bilinear form associated to `q`.

    INPUT:

    - ``p`` -- a prime number
    - ``y`` -- a vector with `q(y) \in p \ZZ`.
    - ``odd`` -- (default=``False``) if `p=2` return also odd neighbors

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ,[1,1,1,1])
        sage: v = vector([0,2,1,1])
        sage: X = Q.find_p_neighbor_from_vec(3,v); X
        Quadratic form in 4 variables over Integer Ring with coefficients:
        [ 1 0 0 0 ]
        [ * 1 4 4 ]
        [ * * 5 12 ]
        [ * * * 9 ]

    Since the base ring and the domain are not yet separate,
    for rational, half integral forms we just pretend
    the base ring is `ZZ`::

        sage: Q = QuadraticForm(QQ,matrix.diagonal([1,1,1,1]))
        sage: v = vector([1,1,1,1])
        sage: Q.find_p_neighbor_from_vec(2,v)
        Quadratic form in 4 variables over Rational Field with coefficients:
        [ 1/2 1 1 1 ]
        [ * 1 1 2 ]
        [ * * 1 2 ]
        [ * * * 2 ]
    """
    p = ZZ(p)
    if not p.divides(self(y)):
        raise ValueError("y=%s must be of square divisible by p=%s"%(y,p))
    if self.base_ring() not in [ZZ, QQ]:
        raise NotImplementedError("the base ring of this form must be the integers or the rationals")
    n = self.dim()
    G = self.Hessian_matrix()
    R = self.base_ring()
    odd = False
    if R is QQ:
      odd = True
      if G.denominator() != 1:
        raise ValueError("the associated bilinear form q(x+y)-q(x)-q(y) must be integral.")
    b = y*G*y
    if not b % p == 0:
        raise ValueError("y^2 must be divisible by p=%s"%p)
    y_dual = y*G
    if p != 2 and b % p**2 != 0:
        for k in range(n):
            if y_dual[k] % p != 0:
                z = (ZZ**n).gen(k)
                break
        else:
            raise ValueError("either y is not primitive or self is not maximal at %s"%p)
        z *= (2*y*G*z).inverse_mod(p)
        y = y - b*z
        # assert y*G*y % p^2 == 0
    if p == 2:
        val = b.valuation(p)
        if val <= 1:
            raise ValueError("y=%s must be of square divisible by 2"%y)
        if val == 2 and not odd:
            # modify it to have square 4
            for k in range(n):
                if y_dual[k] % p != 0:
                    z = (ZZ**n).gen(k)
                    break
            else:
                raise ValueError("either y is not primitive or self is not even, maximal at 2")
            y += 2*z
            # assert y*G*y % 8 == 0

    y_dual = G*y
    Ly = y_dual.change_ring(GF(p)).column().kernel().matrix().lift()
    B = Ly.stack(p * matrix.identity(n))
    # the rows of B now generate L_y = { x in L | (x,y)=0 mod p}
    B = y.row().stack(p*B)
    B = B.hermite_form()[:n, :] / p
    # the rows of B generate ZZ * y/p + L_y
    # by definition this is the p-neighbor of L at y
    # assert B.det().abs() == 1

    QF = self.parent()
    Gnew = (B*G*B.T).change_ring(R)
    return QF(Gnew)


def neighbor_iteration(seeds, p, mass=None, max_classes=ZZ(10)**3,
                       algorithm=None, max_neighbors=1000, verbose=False):
    r"""
    Return all classes in the `p`-neighbor graph of ``self``.

    Starting from the given seeds, this function successively
    finds p-neighbors until no new quadratic form (class) is obtained.

    INPUT:

    - ``seeds`` -- a list of quadratic forms in the same genus

    - ``p`` -- a prime number

    - ``mass`` -- (optional) a rational number; the mass of this genus

    - ``max_classes`` -- (default: ``1000``) break the computation when ``max_classes`` are found

    - ``algorithm`` -- (optional) one of 'orbits', 'random', 'exaustion'

    - ``max_random_trys`` -- (default: ``1000``) the maximum number of neighbors
                             computed for a single lattice

    OUTPUT:

    - a list of quadratic forms

    EXAMPLES::

        sage: from sage.quadratic_forms.quadratic_form__neighbors import neighbor_iteration
        sage: Q = QuadraticForm(ZZ, 3, [1, 0, 0, 2, 1, 3])
        sage: Q.det()
        46
        sage: mass = Q.conway_mass()
        sage: g1 = neighbor_iteration([Q],3, mass=mass, algorithm = 'random') # long time
        sage: g2 = neighbor_iteration([Q],3, algorithm = 'exaustion') # long time
        sage: g3 = neighbor_iteration([Q],3, algorithm = 'orbits')
        sage: mass == sum(1/q.number_of_automorphisms() for q in g1) # long time
        True
        sage: mass == sum(1/q.number_of_automorphisms() for q in g2) # long time
        True
        sage: mass == sum(1/q.number_of_automorphisms() for q in g3)
        True

    TESTS::

        sage: from sage.quadratic_forms.quadratic_form__neighbors import neighbor_iteration
        sage: Q = QuadraticForm(ZZ, 3, [1, 0, 0, 2, 1, 3])
        sage: g = neighbor_iteration([Q],3,mass=Q.conway_mass(),max_classes=2)
        ...
        UserWarning: reached the maximum number of isometry classes=2. Increase the optional argument max_classes to obtain more.
        Warning: not all classes in the genus were found
        sage: neighbor_iteration([Q], 3, mass=Q.conway_mass(), max_neighbors=0, algorithm='random')
        Warning: not all classes in the genus were found
        []
    """
    p = ZZ(p)
    from sage.quadratic_forms.quadratic_form import QuadraticForm
    from warnings import warn
    if not all(isinstance(s, QuadraticForm) for s in seeds):
        raise ValueError("seeds must be a list of quadratic forms")
    if algorithm is None:
        n = seeds[0].dim()
        if p**n > ZZ(2)**18:
            # too many lines to compute the orbits fast
            algorithm = 'random'
        else:
            algorithm = 'orbits'

    if algorithm == 'orbits':
        def p_divisible_vectors(Q, max_neighbors):
            yield from iter(v.lift() for v in Q.orbits_lines_mod_p(p)
                            if v != 0 and Q(v.lift()).valuation(p) > 0)
            return
    elif algorithm == 'exaustion':
        def p_divisible_vectors(Q, max_neighbors):
            k = 0
            v = Q.find_primitive_p_divisible_vector__next(p)
            while k < max_neighbors:
                k = k + 1
                v = Q.find_primitive_p_divisible_vector__next(p, v)
                if v is not None:
                  yield v
    elif algorithm == 'random':
        def p_divisible_vectors(Q, max_neighbors):
            k = 0
            while k < max_neighbors:
                k = k +1
                v = Q.find_primitive_p_divisible_vector__random(p)
                yield v
    else:
        raise ValueError("unknown algorithm")
    waiting_list = list(seeds)
    isom_classes = []
    mass_count = QQ(0)
    n_isom_classes = ZZ(0)
    while len(waiting_list) > 0 and mass != mass_count and n_isom_classes < max_classes:
        # find all p-neighbors of Q
        Q = waiting_list.pop()
        for v in p_divisible_vectors(Q, max_neighbors):
            Q_neighbor = Q.find_p_neighbor_from_vec(p, v)
            if not any(Q_neighbor.is_globally_equivalent_to(S) for S in isom_classes):
                Q_neighbor = Q_neighbor.lll()
                isom_classes.append(Q_neighbor)
                waiting_list.append(Q_neighbor)
                n_isom_classes += 1
                mass_count += Q_neighbor.number_of_automorphisms()**(-1)
                if verbose:
                    print(max_neighbors)
                    print(len(waiting_list))
                if mass_count == mass or n_isom_classes >= max_classes:
                    break

    if len(isom_classes) >= max_classes:
        warn("reached the maximum number of isometry classes=%s. Increase the optional argument max_classes to obtain more." %max_classes)

    if mass is not None:
        assert mass_count <= mass
        if mass_count < mass:
            print("Warning: not all classes in the genus were found")
    return isom_classes

def orbits_lines_mod_p(self, p):
    r"""
    Let `(L, q)` be a lattice. This returns representatives of the
    orbits of lines in `L/pL` under the orthogonal group of `q`.

    INPUT:

    - ``p`` -- a prime number

    OUTPUT:

    - a list of vectors over ``GF(p)``

    EXAMPLES::

        sage: from sage.quadratic_forms.quadratic_form__neighbors import orbits_lines_mod_p
        sage: Q = QuadraticForm(ZZ, 3, [1, 0, 0, 2, 1, 3])
        sage: Q.orbits_lines_mod_p(2)
        [(0, 0, 1),
        (0, 1, 0),
        (0, 1, 1),
        (1, 0, 0),
        (1, 0, 1),
        (1, 1, 0),
        (1, 1, 1)]
    """
    from sage.libs.gap.libgap import libgap
    # careful the self.automorphism_group() acts from the left
    # but in gap we act from the right!! --> transpose
    gens = self.automorphism_group().gens()
    gens = [g.matrix().transpose().change_ring(GF(p)) for g in gens]
    orbs = libgap.function_factory(
    """function(gens, p)
        local one, G, reps, V, n, orb;
        one:= One(GF(p));
        G:=Group(List(gens, g -> g*one));
        n:= Size(gens[1]);
        V:= GF(p)^n;
        orb:= OrbitsDomain(G, V, OnLines);
        reps:= List(orb, g->g[1]);
        return reps;
        end;""")
    # run this at startup if you need more memory...
    #from sage.interfaces.gap import get_gap_memory_pool_size, set_gap_memory_pool_size
    #memory_gap = get_gap_memory_pool_size()
    #set_gap_memory_pool_size(1028*memory_gap)
    orbs_reps = orbs(gens, p)
    #set_gap_memory_pool_size(memory_gap)
    M = GF(p)**self.dim()
    return [M(m.sage()) for m in orbs_reps if not m.IsZero()]



