"""
Counting Congruence Solutions

This file provides more user-friendly Python front-ends to the Cython code in
:mod:`sage.quadratic_forms.count_local`.
"""
##################################################################
## Methods for counting/computing the number of representations ##
## of a number by a quadratic form in Z/(p^k)Z of various types ##
##################################################################


from sage.quadratic_forms.count_local_2 import CountAllLocalTypesNaive



def count_congruence_solutions_as_vector(self, p, k, m, zvec, nzvec):
    """
    Gives the number of integer solution vectors `x` satisfying the
    congruence Q(`x`) `= m (mod p^k)` of each solution type (i.e. All,
    Good, Zero, Bad, BadI, BadII) which satisfy the additional
    congruence conditions of having certain coefficients = 0 (mod `p`)
    and certain collections of coefficients not congruent to the zero
    vector (mod `p`).

    A solution vector `x` satisfies the additional congruence conditions
    specified by zvec and nzvec (and therefore is counted) iff both of
    the following conditions hold:

        1) `x[i] == 0 (mod p)` for all `i` in zvec

        2) `x[i] != 0 (mod p) for all i` in nzvec


    REFERENCES:

    See Hanke's (????) paper "Local Densities and explicit bounds...", p??? for
    the definitions of the solution types and congruence conditions.

    INPUT:

    - `p` -- prime number > 0
    - `k` -- an integer > 0
    - `m` -- an integer (depending only on mod `p^k`)
    - zvec, nzvec -- a list of integers in range(self.dim()), or None

    OUTPUT:

    a list of six integers >= 0 representing the solution types:
    [All, Good, Zero, Bad, BadI, BadII]


    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions_as_vector(3, 1, 1, [], [])
        [0, 0, 0, 0, 0, 0]
        sage: Q.count_congruence_solutions_as_vector(3, 1, 1, None, [])
        [0, 0, 0, 0, 0, 0]
        sage: Q.count_congruence_solutions_as_vector(3, 1, 1, [], None)
        [6, 6, 0, 0, 0, 0]
        sage: Q.count_congruence_solutions_as_vector(3, 1, 1, None, None)
        [6, 6, 0, 0, 0, 0]
        sage: Q.count_congruence_solutions_as_vector(3, 1, 2, None, None)
        [6, 6, 0, 0, 0, 0]
        sage: Q.count_congruence_solutions_as_vector(3, 1, 0, None, None)
        [15, 12, 1, 2, 0, 2]

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)






##///////////////////////////////////////////
##/// Front-ends for our counting routines //
##///////////////////////////////////////////

def count_congruence_solutions(self, p, k, m, zvec, nzvec):
    """
    Counts all solutions of Q(`x`) `= m (mod p^k)` satisfying the
    additional congruence conditions described in
    QuadraticForm.count_congruence_solutions_as_vector().

    INPUT:

        `p` -- prime number > 0

        `k` -- an integer > 0

        `m` -- an integer (depending only on mod `p^k`)

        zvec, nzvec -- a list of integers in range(self.dim()), or None

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions(3, 1, 0, None, None)
        15

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)[0]



def count_congruence_solutions__good_type(self, p, k, m, zvec, nzvec):
    """
    Counts the good-type solutions of Q(x) = m (mod p^k) satisfying the
    additional congruence conditions described in
    QuadraticForm.count_congruence_solutions_as_vector().

    INPUT:

        `p` -- prime number > 0

        `k` -- an integer > 0

        `m` -- an integer (depending only on mod `p^k`)

        zvec, nzvec -- a list of integers up to dim(Q)

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions__good_type(3, 1, 0, None, None)
        12

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)[1]



def count_congruence_solutions__zero_type(self, p, k, m, zvec, nzvec):
    """
    Counts the zero-type solutions of Q(`x`) = `m (mod p^k)` satisfying the
    additional congruence conditions described in
    QuadraticForm.count_congruence_solutions_as_vector().

    INPUT:

        `p` -- prime number > 0

        `k` -- an integer > 0

        `m` -- an integer (depending only on mod `p^k`)

        zvec, nzvec -- a list of integers up to dim(Q)

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions__zero_type(3, 1, 0, None, None)
        1

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)[2]


def count_congruence_solutions__bad_type(self, p, k, m, zvec, nzvec):
    """
    Counts the bad-type solutions of Q(`x`) `= m (mod p^k)` satisfying the
    additional congruence conditions described in
    QuadraticForm.count_congruence_solutions_as_vector().

    INPUT:

        `p` -- prime number > 0

        `k` -- an integer > 0

        `m` -- an integer (depending only on mod `p^k`)

        zvec, nzvec -- a list of integers up to dim(Q)

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions__bad_type(3, 1, 0, None, None)
        2

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)[3]


def count_congruence_solutions__bad_type_I(self, p, k, m, zvec, nzvec):
    """
    Counts the bad-typeI solutions of Q(`x`) = `m (mod p^k)` satisfying
    the additional congruence conditions described in
    QuadraticForm.count_congruence_solutions_as_vector().

    INPUT:

        `p` -- prime number > 0

        `k` -- an integer > 0

        `m` -- an integer (depending only on mod `p^k`)

        zvec, nzvec -- a list of integers up to dim(Q)

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions__bad_type_I(3, 1, 0, None, None)
        0

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)[4]


def count_congruence_solutions__bad_type_II(self, p, k, m, zvec, nzvec):
    """
    Counts the bad-typeII solutions of Q(`x`) `= m (mod p^k)` satisfying
    the additional congruence conditions described in
    QuadraticForm.count_congruence_solutions_as_vector().

    INPUT:

        `p` -- prime number > 0

        `k` -- an integer > 0

        `m` -- an integer (depending only on mod `p^k`)

        zvec, nzvec -- a list of integers up to dim(Q)

    EXAMPLES::

        sage: Q = DiagonalQuadraticForm(ZZ, [1,2,3])
        sage: Q.count_congruence_solutions__bad_type_II(3, 1, 0, None, None)
        2

    """
    return CountAllLocalTypesNaive(self, p, k, m, zvec, nzvec)[5]
