import functools
import math

import mpmath

from ..core.compatibility import as_int
from ..core.numbers import igcdex
from .primetest import isprime


def symmetric_residue(a, m):
    """Return the residual mod m such that it is within half of the modulus.

    >>> symmetric_residue(1, 6)
    1
    >>> symmetric_residue(4, 6)
    -2

    """
    if a <= m // 2:
        return a
    return a - m


def crt(M, U, symmetric=False, check=True):
    r"""Chinese Remainder Theorem.

    The moduli in M are assumed to be pairwise coprime.  The output
    is then an integer f, such that f = u_i mod m_i for each pair out
    of U and M. If ``symmetric`` is False a positive integer will be
    returned, else \|f\| will be less than or equal to the LCM of the
    moduli, and thus f may be negative.

    If the moduli are not co-prime the correct result will be returned
    if/when the test of the result is found to be incorrect. This result
    will be None if there is no solution.

    The keyword ``check`` can be set to False if it is known that the moduli
    are coprime.

    Examples
    ========

    >>> crt([99, 97, 95], [49, 76, 65])
    (639985, 912285)

    This is the correct result because:

    >>> [639985 % m for m in [99, 97, 95]]
    [49, 76, 65]

    If the moduli are not co-prime, you may receive an incorrect result
    if you use ``check=False``:

    >>> crt([12, 6, 17], [3, 4, 2], check=False)
    (954, 1224)
    >>> [954 % m for m in [12, 6, 17]]
    [6, 0, 2]
    >>> crt([12, 6, 17], [3, 4, 2]) is None
    True
    >>> crt([3, 6], [2, 5])
    (5, 6)

    Notes
    =====

    Rather than checking that all pairs of moduli share
    no GCD (an O(n**2) test) and rather than factoring all moduli and seeing
    that there is no factor in common, a check that the result gives the
    indicated residuals is performed -- an O(n) operation.

    See Also
    ========

    solve_congruence

    """
    from ..domains import ZZ

    if check:
        M = list(map(as_int, M))
        U = list(map(as_int, U))

    p = math.prod(M, start=ZZ.one)
    v = ZZ.zero

    for u, m in zip(U, M):
        e = p // m
        s, _, _ = ZZ.gcdex(e, m)
        v += e*(u*s % m)

    result = v % p
    mm = math.prod(M)

    if check:
        if not all(v % m == result % m for v, m in zip(U, M)):
            result = solve_congruence(*list(zip(U, M)),
                                      check=False, symmetric=symmetric)
            if result is None:
                return result
            result, mm = result

    if symmetric:
        return symmetric_residue(result, mm), mm
    return result, mm


def crt1(M):
    """First part of Chinese Remainder Theorem, for multiple application.

    Examples
    ========

    >>> crt1([18, 42, 6])
    (4536, [252, 108, 756], [0, 2, 0])

    """
    from ..domains import ZZ

    E, S = [], []
    p = math.prod(M, start=ZZ.one)

    for m in M:
        E.append(p // m)
        S.append(ZZ.gcdex(E[-1], m)[0] % m)

    return p, E, S


def crt2(M, U, p, E, S, symmetric=False):
    """Second part of Chinese Remainder Theorem, for multiple application.

    Examples
    ========

    >>> mm, e, s = crt1([18, 42, 6])
    >>> crt2([18, 42, 6], [0, 0, 0], mm, e, s)
    (0, 4536)

    """
    from ..domains import ZZ

    v = ZZ.zero

    for u, m, e, s in zip(U, M, E, S):
        v += e*(u*s % m)

    result = v % p

    if symmetric:
        result = symmetric_residue(result, p)
    return result, p


def solve_congruence(*remainder_modulus_pairs, **hint):
    """Compute the integer ``n`` that has the residual ``ai`` when it is
    divided by ``mi`` where the ``ai`` and ``mi`` are given as pairs to
    this function: ((a1, m1), (a2, m2), ...). If there is no solution,
    return. Otherwise return ``n`` and its modulus.

    The ``mi`` values need not be co-prime. If it is known that the moduli are
    not co-prime then the hint ``check`` can be set to False (default=True) and
    the check for a quicker solution via crt() (valid when the moduli are
    co-prime) will be skipped.

    If the hint ``symmetric`` is True (default is False), the value of ``n``
    will be within 1/2 of the modulus, possibly negative.

    Examples
    ========

    What number is 2 mod 3, 3 mod 5 and 2 mod 7?

    >>> solve_congruence((2, 3), (3, 5), (2, 7))
    (23, 105)
    >>> [23 % m for m in [3, 5, 7]]
    [2, 3, 2]

    If you prefer to work with all remainder in one list and
    all moduli in another, send the arguments like this:

    >>> solve_congruence(*zip((2, 3, 2), (3, 5, 7)))
    (23, 105)

    The moduli need not be co-prime; in this case there may or
    may not be a solution:

    >>> solve_congruence((2, 3), (4, 6)) is None
    True

    >>> solve_congruence((2, 3), (5, 6))
    (5, 6)

    The symmetric flag will make the result be within 1/2 of the modulus:

    >>> solve_congruence((2, 3), (5, 6), symmetric=True)
    (-1, 6)

    See Also
    ========

    crt : high level routine implementing the Chinese Remainder Theorem

    """
    def combine(c1, c2):
        """Return the tuple (a, m) which satisfies the requirement
        that n = a + i*m satisfy n = a1 + j*m1 and n = a2 = k*m2.

        References
        ==========

        * https://en.wikipedia.org/wiki/Method_of_successive_substitution

        """
        a1, m1 = c1
        a2, m2 = c2
        a, b, c = m1, a2 - a1, m2
        g = functools.reduce(math.gcd, [a, b, c])
        a, b, c = (i//g for i in [a, b, c])
        if a != 1:
            inv_a, _, g = igcdex(a, c)
            if g != 1:
                return
            b *= inv_a
        a, m = a1 + m1*b, m1*c
        return a, m

    rm = remainder_modulus_pairs
    symmetric = hint.get('symmetric', False)

    if hint.get('check', True):
        rm = [(as_int(r), as_int(m)) for r, m in rm]

        # ignore redundant pairs but raise an error otherwise; also
        # make sure that a unique set of bases is sent to crt() if
        # they are all prime.
        #
        # The routine will work out less-trivial violations and
        # return None, e.g. for the pairs (1,3) and (14,42) there
        # is no answer because 14 mod 42 (having a gcd of 14) implies
        # (14/2) mod (42/2), (14/7) mod (42/7) and (14/14) mod (42/14)
        # which, being 0 mod 3, is inconsistent with 1 mod 3. But to
        # preprocess the input beyond checking of another pair with 42
        # or 3 as the modulus (for this example) is not necessary.
        uniq = {}
        for r, m in rm:
            r %= m
            if m in uniq:
                if r != uniq[m]:
                    return
            else:
                uniq[m] = r
        rm = [(r, m) for m, r in uniq.items()]
        del uniq

        # if the moduli are co-prime, the crt will be significantly faster;
        # checking all pairs for being co-prime gets to be slow but a prime
        # test is a good trade-off
        if all(isprime(m) for r, m in rm):
            r, m = list(zip(*rm))
            return crt(m, r, symmetric=symmetric, check=False)

    rv = (0, 1)
    for rmi in rm:
        rv = combine(rv, rmi)
        if rv is None:
            break
        n, m = rv
        n = n % m
    else:
        if symmetric:
            return symmetric_residue(n, m), m
        return n, m


def integer_rational_reconstruction(c, m, domain):
    r"""
    Reconstruct a rational number `\frac a b` from

    .. math::

        c = \frac a b \; \mathrm{mod} \, m,

    where `c` and `m` are integers.

    The algorithm is based on the Euclidean Algorithm. In general, `m` is
    not a prime number, so it is possible that `b` is not invertible modulo
    `m`. In that case ``None`` is returned.

    Parameters
    ==========

    c : Integer
        `c = \frac a b \; \mathrm{mod} \, m`
    m : Integer
        modulus, not necessarily prime
    domain : IntegerRing
        `a, b, c` are elements of ``domain``

    Returns
    =======

    frac : Rational
        either `\frac a b` in `\mathbb Q` or ``None``

    References
    ==========

    * :cite:`Wang1981partial`

    """
    if c < 0:
        c += m

    r0, s0 = m, domain.zero
    r1, s1 = c, domain.one

    bound = mpmath.sqrt(m / 2)  # still correct if replaced by ZZ.sqrt(m // 2) ?

    while r1 >= bound:
        quo = r0 // r1
        r0, r1 = r1, r0 - quo*r1
        s0, s1 = s1, s0 - quo*s1

    if abs(s1) >= bound:
        return

    if s1 < 0:
        r1, s1 = -r1, -s1

    field = domain.field

    return field(r1) / field(s1)
