import operator
from six import print_
from functools import reduce
import operator
try:
    import unittest2 as unittest
except ImportError:
    import unittest
import hypothesis.strategies as st
import pytest
from hypothesis import given, settings, example
try:
    from hypothesis import HealthCheck
    HC_PRESENT=True
except ImportError:
    HC_PRESENT=False
from .numbertheory import (SquareRootError, factorization, gcd, lcm,
                           jacobi, inverse_mod,
                           is_prime, next_prime, smallprimes,
                           square_root_mod_prime)


BIGPRIMES = (999671,
             999683,
             999721,
             999727,
             999749,
             999763,
             999769,
             999773,
             999809,
             999853,
             999863,
             999883,
             999907,
             999917,
             999931,
             999953,
             999959,
             999961,
             999979,
             999983)


@pytest.mark.parametrize(
    "prime, next_p",
    [(p, q) for p, q in zip(BIGPRIMES[:-1], BIGPRIMES[1:])])
def test_next_prime(prime, next_p):
    assert next_prime(prime) == next_p


@pytest.mark.parametrize("prime", smallprimes)
def test_square_root_mod_prime_for_small_primes(prime):
    squares = set()
    for num in range(0, 1 + prime // 2):
        sq = num * num % prime
        squares.add(sq)
        root = square_root_mod_prime(sq, prime)
        # tested for real with TestNumbertheory.test_square_root_mod_prime
        assert root * root % prime == sq

    for nonsquare in range(0, prime):
        if nonsquare in squares:
            continue
        with pytest.raises(SquareRootError):
            square_root_mod_prime(nonsquare, prime)


@st.composite
def st_two_nums_rel_prime(draw):
    # 521-bit is the biggest curve we operate on, use 1024 for a bit
    # of breathing space
    mod = draw(st.integers(min_value=2, max_value=2**1024))
    num = draw(st.integers(min_value=1, max_value=mod-1)
               .filter(lambda x: gcd(x, mod) == 1))
    return num, mod


@st.composite
def st_primes(draw, *args, **kwargs):
    if "min_value" not in kwargs:
        kwargs["min_value"] = 1
    prime = draw(st.sampled_from(smallprimes) |
                 st.integers(*args, **kwargs)
                 .filter(is_prime))
    return prime


@st.composite
def st_num_square_prime(draw):
    prime = draw(st_primes(max_value=2**1024))
    num = draw(st.integers(min_value=0, max_value=1 + prime // 2))
    sq = num * num % prime
    return sq, prime


@st.composite
def st_comp_with_com_fac(draw):
    """
    Strategy that returns lists of numbers, all having a common factor.
    """
    primes = draw(st.lists(st_primes(max_value=2**512), min_size=1,
                           max_size=10))
    # select random prime(s) that will make the common factor of composites
    com_fac_primes = draw(st.lists(st.sampled_from(primes),
                                   min_size=1, max_size=20))
    com_fac = reduce(operator.mul, com_fac_primes, 1)

    # select at most 20 lists (returned numbers),
    # each having at most 30 primes (factors) including none (then the number
    # will be 1)
    comp_primes = draw(
        st.integers(min_value=1, max_value=20).
        flatmap(lambda n: st.lists(st.lists(st.sampled_from(primes),
                                            max_size=30),
                                   min_size=1, max_size=n)))

    return [reduce(operator.mul, nums, 1) * com_fac for nums in comp_primes]


@st.composite
def st_comp_no_com_fac(draw):
    """
    Strategy that returns lists of numbers that don't have a common factor.
    """
    primes = draw(st.lists(st_primes(max_value=2**512),
                           min_size=2, max_size=10, unique=True))
    # first select the primes that will create the uncommon factor
    # between returned numbers
    uncom_fac_primes = draw(st.lists(
        st.sampled_from(primes),
        min_size=1, max_size=len(primes)-1, unique=True))
    uncom_fac = reduce(operator.mul, uncom_fac_primes, 1)

    # then build composites from leftover primes
    leftover_primes = [i for i in primes if i not in uncom_fac_primes]

    assert leftover_primes
    assert uncom_fac_primes

    # select at most 20 lists, each having at most 30 primes
    # selected from the leftover_primes list
    number_primes = draw(
        st.integers(min_value=1, max_value=20).
        flatmap(lambda n: st.lists(st.lists(st.sampled_from(leftover_primes),
                                            max_size=30),
                                   min_size=1, max_size=n)))

    numbers = [reduce(operator.mul, nums, 1) for nums in number_primes]

    insert_at = draw(st.integers(min_value=0, max_value=len(numbers)))
    numbers.insert(insert_at, uncom_fac)
    return numbers


HYP_SETTINGS = {}
if HC_PRESENT:
    HYP_SETTINGS['suppress_health_check']=[HealthCheck.filter_too_much,
                                           HealthCheck.too_slow]
    # the factorization() sometimes takes a long time to finish
    HYP_SETTINGS['deadline'] = 5000


class TestNumbertheory(unittest.TestCase):
    def test_gcd(self):
        assert gcd(3 * 5 * 7, 3 * 5 * 11, 3 * 5 * 13) == 3 * 5
        assert gcd([3 * 5 * 7, 3 * 5 * 11, 3 * 5 * 13]) == 3 * 5
        assert gcd(3) == 3

    @unittest.skipUnless(HC_PRESENT,
                         "Hypothesis 2.0.0 can't be made tolerant of hard to "
                         "meet requirements (like `is_prime()`), the test "
                         "case times-out on it")
    @settings(**HYP_SETTINGS)
    @given(st_comp_with_com_fac())
    def test_gcd_with_com_factor(self, numbers):
        n = gcd(numbers)
        assert 1 in numbers or n != 1
        for i in numbers:
            assert i % n == 0

    @unittest.skipUnless(HC_PRESENT,
                         "Hypothesis 2.0.0 can't be made tolerant of hard to "
                         "meet requirements (like `is_prime()`), the test "
                         "case times-out on it")
    @settings(**HYP_SETTINGS)
    @given(st_comp_no_com_fac())
    def test_gcd_with_uncom_factor(self, numbers):
        n = gcd(numbers)
        assert n == 1

    @given(st.lists(st.integers(min_value=1, max_value=2**8192),
                    min_size=1, max_size=20))
    def test_gcd_with_random_numbers(self, numbers):
        n = gcd(numbers)
        for i in numbers:
            # check that at least it's a divider
            assert i % n == 0

    def test_lcm(self):
        assert lcm(3, 5 * 3, 7 * 3) == 3 * 5 * 7
        assert lcm([3, 5 * 3, 7 * 3]) == 3 * 5 * 7
        assert lcm(3) == 3

    @given(st.lists(st.integers(min_value=1, max_value=2**8192),
                    min_size=1, max_size=20))
    def test_lcm_with_random_numbers(self, numbers):
        n = lcm(numbers)
        for i in numbers:
            assert n % i == 0

    @unittest.skipUnless(HC_PRESENT,
                         "Hypothesis 2.0.0 can't be made tolerant of hard to "
                         "meet requirements (like `is_prime()`), the test "
                         "case times-out on it")
    @settings(**HYP_SETTINGS)
    @given(st_num_square_prime())
    def test_square_root_mod_prime(self, vals):
        square, prime = vals

        calc = square_root_mod_prime(square, prime)
        assert calc * calc % prime == square

    @settings(**HYP_SETTINGS)
    @given(st.integers(min_value=1, max_value=10**12))
    @example(265399 * 1526929)
    @example(373297 ** 2 * 553991)
    def test_factorization(self, num):
        factors = factorization(num)
        mult = 1
        for i in factors:
            mult *= i[0] ** i[1]
        assert mult == num

    @settings(**HYP_SETTINGS)
    @given(st.integers(min_value=3, max_value=1000).filter(lambda x: x % 2))
    def test_jacobi(self, mod):
        if is_prime(mod):
            squares = set()
            for root in range(1, mod):
                assert jacobi(root * root, mod) == 1
                squares.add(root * root % mod)
            for i in range(1, mod):
                if i not in squares:
                    assert jacobi(i, mod) == -1
        else:
            factors = factorization(mod)
            for a in range(1, mod):
                c = 1
                for i in factors:
                    c *= jacobi(a, i[0]) ** i[1]
                assert c == jacobi(a, mod)

    @given(st_two_nums_rel_prime())
    def test_inverse_mod(self, nums):
        num, mod = nums

        inv = inverse_mod(num, mod)

        assert 0 < inv < mod
        assert num * inv % mod == 1
