# -*- coding: utf-8 -*-

from ..syntax import macros, test, test_raises, fail, the  # noqa: F401
from ..test.fixtures import session, testset

from itertools import count, takewhile, chain
from collections import Counter

from ..gmemo import gmemoize, imemoize, fimemoize

from ..it import take, drop, last
from ..fold import prod
from ..funutil import call
from ..misc import timer

def runtests():
    with testset("multiple instances, interleaved"):
        total_evaluations = 0
        @gmemoize
        def gen():
            nonlocal total_evaluations
            j = 1
            while True:
                total_evaluations += 1
                yield j
                j += 1
        g1 = gen()
        g2 = gen()
        test[next(g1) == 1]
        test[next(g1) == 2]
        test[next(g2) == 1]
        test[next(g1) == 3]
        test[next(g2) == 2]
        test[next(g2) == 3]
        test[next(g2) == 4]
        test[next(g1) == 4]
        g3 = gen()
        test[next(g3) == 1]
        test[next(g3) == 2]
        test[next(g3) == 3]
        test[next(g3) == 4]
        test[total_evaluations == 4]

    with testset("multiple instances, exhaust one first"):
        total_evaluations = 0
        @gmemoize
        def gen():
            nonlocal total_evaluations
            for j in range(3):
                total_evaluations += 1
                yield j
        g1 = gen()
        g2 = gen()
        test[total_evaluations == 0]
        test[tuple(x for x in g1) == (0, 1, 2)]
        test[total_evaluations == 3]
        test[tuple(x for x in g2) == (0, 1, 2)]
        test[total_evaluations == 3]

    with testset("@gmemoize caches exceptions"):
        class AllOkJustTesting(Exception):
            pass
        total_evaluations = 0
        @gmemoize
        def gen():
            nonlocal total_evaluations
            total_evaluations += 1
            yield 1
            total_evaluations += 1
            raise AllOkJustTesting("ha ha only serious")

        g1 = gen()
        test[total_evaluations == 0]
        try:
            next(g1)
            test[total_evaluations == 1]
            next(g1)
        except AllOkJustTesting as err:
            exc_instance = err
        else:
            fail["Should have raised at the second next() call."]  # pragma: no cover
        test[total_evaluations == 2]

        g2 = gen()
        next(g2)
        test[total_evaluations == 2]  # still just two, it's memoized
        try:
            next(g2)
        except AllOkJustTesting as err2:
            test[the[err2 is exc_instance], "should be the same cached exception instance"]
        else:
            fail["Should have raised at the second next() call."]  # pragma: no cover
        test[total_evaluations == 2]

    with testset("subscripting to get already computed items"):
        @gmemoize
        def gen():
            yield from range(5)
        g3 = gen()

        # Any item that has entered the memo can be retrieved by subscripting.
        # len() is the current length of the memo.
        test[len(g3) == 0]
        next(g3)
        test[len(g3) == 1]
        next(g3)
        test[len(g3) == 2]
        next(g3)
        test[len(g3) == 3]
        test[g3[0] == 0]
        test[g3[1] == 1]
        test[g3[2] == 2]

        # Items not yet memoized cannot be retrieved from the memo.
        test_raises[IndexError, g3[3]]

        # Negative indices work too, counting from the current end of the memo.
        test[g3[-1] == 2]
        test[g3[-2] == 1]
        test[g3[-3] == 0]

        # Counting back past the start is an error, just like in `list`.
        test_raises[IndexError, g3[-4]]

        # Slicing is supported.
        test[g3[0:3] == [0, 1, 2]]
        test[g3[0:2] == [0, 1]]
        test[g3[::-1] == [2, 1, 0]]
        test[g3[0::2] == [0, 2]]
        test[g3[2::-2] == [2, 0]]

        # Out-of-range slices produce the empty list, like in `list`.
        test[g3[3:] == []]
        test[g3[-4::-1] == []]

    with testset("memoizing a sequence partially"):
        # To do this, build a chain of generators, then memoize only the last one:
        evaluations = Counter()
        def orig():
            yield from range(100)
        def evens():
            yield from (x for x in orig() if x % 2 == 0)
        @gmemoize
        def some_evens(n):  # drop n first terms
            evaluations[n] += 1
            yield from drop(n, evens())
        last(some_evens(25))
        last(some_evens(25))
        last(some_evens(20))
        test[all(v == 1 for k, v in the[evaluations].items())]

        # Or use lambda for a more compact presentation:
        se = gmemoize(lambda n: (yield from drop(n, evens())))
        test[the[last(se(25))] == the[last(se(25))]]  # iterating twice!

        # Using fimemoize, we can omit the "yield from" (specifying a regular
        # factory function that makes an iterable, instead of a gfunc):
        se = fimemoize(lambda n: drop(n, evens()))
        test[the[last(se(25))] == the[last(se(25))]]  # iterating twice!

        # In the nonparametric case, we can memoize the iterable directly:
        se = imemoize(drop(25, evens()))
        test[the[last(se())] == the[last(se())]]  # iterating twice!

        # DANGER: WRONG! Now we get a new instance of evens() also for the same n,
        # so each call to se(n) caches separately. (This is why we have fimemoize.)
        se = lambda n: call(imemoize(drop(n, evens())))  # call() invokes the gfunc
        test[the[last(se(25))] == the[last(se(25))]]
        test[the[last(se(20))] == the[last(se(20))]]

    with testset("FP sieve of Eratosthenes"):
        def primes():  # no memoization, recomputes unnecessarily, very slow
            yield 2
            for n in count(start=3, step=2):
                if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, primes())):
                    yield n

        @gmemoize  # <-- the only change (beside the function name)
        def mprimes():  # external memo for users, re-use it internally - simplest code
            yield 2
            for n in count(start=3, step=2):
                if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, mprimes())):
                    yield n

        def memo_primes():  # manual internal memo only - fastest, no caching for users
            memo = []
            def manual_mprimes():
                memo.append(2)
                yield 2
                for n in count(start=3, step=2):
                    if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, memo)):
                        memo.append(n)
                        yield n
            return manual_mprimes()

        # external memo for users, separate manual internal memo
        # doubles memory usage due to exactly one internal memo; almost as fast as memo_primes
        # since the tight inner loop skips the very general gmemoize machinery
        #
        # This version wins in speed for moderate n (1e5) on typical architectures where
        # the memory bus is a bottleneck, since the rule for generating new candidates is
        # simple arithmetic. Contrast memo_primes3, which needs to keep a table that gets
        # larger as n grows (so memory transfers dominate for large n). That strategy
        # seems faster for n ~ 1e3, though.
        @gmemoize
        def memo_primes2():
            memo = []
            def manual_mprimes2():
                memo.append(2)
                yield 2
                for n in count(start=3, step=2):
                    if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, memo)):
                        memo.append(n)
                        yield n
            return manual_mprimes2()

        # small refinement: skip testing 15, 25, 35, ...
        # - we know that in base-10, for any prime > 10 the last digit must be 1, 3, 7 or 9;
        #   if it is 0, 2 or 5, the number is divisible by at least one factor of 10 (namely 2 or 5)
        # - n < 10 must be checked separately; the primes are 2, 3, 5, 7
        #   (note the factors of 10 are there, plus some unrelated primes)
        @gmemoize
        def mprimes2():
            yield 2
            for n in chain([3, 5, 7], (d + k for d in count(10, step=10)
                                             for k in [1, 3, 7, 9])):
                if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, mprimes2())):
                    yield n

        # generalization: let's not be limited by base-10
        # base-b representation, switch b when appropriate:
        #   n = k*b + m
        #   b = 2*3, 2*3*5, 2*3*5*7, ...
        # k: integer,  1, 2, ..., {next factor to account for in b} - 1
        #    so e.g. when b=6, we check from 6 to 29; when b=30, from 30 to 209, ...
        # m: last digit in base-b representation of n, note m < b
        #    for a number represented in base-b to be prime, m must not be divisible by any factor of b
        # Only the numbers up to b must be checked separately (and already have when we reach the next b).
        #
        # For the first 5e4 primes, about 20% of the integers within each range are candidates.
        # If you want the details, add this just before "for n in ns:":
        #    print(b, ns[-1]**(1/2), len(ns), (nextp-1)*b, len(ns)/((nextp-1)*b))
        @gmemoize
        def mprimes3():
            # minimal init takes three terms; b = 2*3 = 6 > 5, so no overlap in output of init and general loop
            # (and this init yields all primes up to b = 6)
            yield from (2, 3, 5)
            theprimes = mprimes3()
            ps = list(take(2, theprimes))  # factors of b; b is chosen s.t. each factor is a different prime
            p, b, np = ps[-1], prod(ps), len(ps)
            lastdigits = [1, 3, 5]  # last digits in base-6 that are not divisible by 2

            while True:
                nextp = next(theprimes)
                lastdigits = [n for n in lastdigits if n % p != 0]
                ns = [k * b + m for k in range(1, nextp)
                                for m in lastdigits]
                # in ns, we have already eliminated the first np primes as possible factors, so skip checking them
                for n in ns:
                    if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, drop(np, mprimes3()))):
                        yield n
                ps.append(nextp)
                b *= nextp
                p = nextp
                np += 1
                lastdigits = lastdigits + ns
        test[the[tuple(take(500, mprimes3())) == tuple(take(500, mprimes2()))]]  # de-spam: don't capture the LHS.

        @gmemoize
        def memo_primes3():
            memo = []
            def manual_mprimes3():
                for p in (2, 3, 5):
                    memo.append(p)
                    yield p
                p, b, np = 3, 6, 2
                lastdigits = [1, 3, 5]

                while True:
                    nextp = memo[np]
                    lastdigits = [n for n in lastdigits if n % p != 0]
                    ns = [k * b + m for k in range(1, nextp)
                                    for m in lastdigits]
                    for n in ns:
                        if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, drop(np, memo))):
                            memo.append(n)
                            yield n
                    b *= nextp
                    p = nextp
                    np += 1
                    lastdigits += ns
            return manual_mprimes3()
        test[the[tuple(take(500, memo_primes3())) == tuple(take(500, mprimes2()))]]

        @gmemoize
        def memo_primes4():
            memo = []
            def manual_mprimes4():
                for p in (2, 3, 5):
                    memo.append(p)
                    yield p
                p, b, np = 3, 6, 2
                lastdigits = [1, 3, 5]
                maxnp = 5  # --> b = 2*3*5*7*11 = 2310; optimal setting depends on CPU cache size

                while True:
                    nextp = memo[np]
                    lastdigits = [n for n in lastdigits if n % p != 0]
                    ns = [k * b + m for k in range(1, nextp)
                                    for m in lastdigits]
                    for n in ns:
                        if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, drop(np, memo))):
                            memo.append(n)
                            yield n
                    if np == maxnp:  # avoid table becoming too big (leading to memory bus dominated run time)
                        break
                    b *= nextp
                    p = nextp
                    np += 1
                    lastdigits += ns
                # once maximum b reached, stay at that b, using the final table of lastdigits
                for kb in count(nextp * b, step=b):
                    for n in (kb + m for m in lastdigits):
                        if not any(n % p == 0 for p in takewhile(lambda x: x * x <= n, drop(np, memo))):
                            memo.append(n)
                            yield n
            return manual_mprimes4()
        test[the[tuple(take(500, memo_primes4())) == tuple(take(500, mprimes2()))]]
        test[last(take(5000, memo_primes4())) == 48611]  # trigger the maxnp case

        test[tuple(take(10, primes())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]
        test[tuple(take(10, mprimes())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]
        test[tuple(take(10, memo_primes())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]
        test[tuple(take(10, mprimes2())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]
        test[tuple(take(10, memo_primes2())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]
        test[tuple(take(10, mprimes3())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]
        test[tuple(take(10, memo_primes3())) == (2, 3, 5, 7, 11, 13, 17, 19, 23, 29)]

    # TODO: need some kind of benchmarking tools to do this properly.
    with testset("performance benchmark"):
        n = 2500
        print(f"Performance for first {n:d} primes:")
        for g in (mprimes(), memo_primes(), mprimes2(), memo_primes2(), mprimes3(), memo_primes3(), memo_primes4()):
            with timer() as tictoc:
                last(take(n, g))
            print(g, tictoc.dt)

if __name__ == '__main__':  # pragma: no cover
    with session(__file__):
        runtests()