import bisect
import dataclasses
import itertools
import timeit
import math
from typing import Generator, Callable, Iterator

@dataclasses.dataclass
class ConsecutivePrimeSum:
    prime: int
    consecutive_primes: list

    @property
    def chain_length(self):
        return len(self.consecutive_primes)

def eratosthenes_sieve(n: int) -> Callable[[int], bool]:
    """Returns a function that is an efficient implementation of is_prime

        Generates a Eratosthenes' sieve for primes, and uses it as a cache for a wrapped function
    """
    is_prime = [True] * (n + 1)
    is_prime[0] = is_prime[1] = False

    for i in range(2, int(math.sqrt(n)) + 1):
        # if i is marked as "not prime", its multiple must have been too
        # so nothing to do here
        if not is_prime[i]:
            continue

        # If i is not marked as "not prime", it MUST be prime
        # Mark all multiples of i as not prime
        multiples_of_i = range(i**2, n + 1, i)
        for j in multiples_of_i:
            is_prime[j] = False

    return lambda x: is_prime[x]

def primes_under(n: int) -> Iterator[int]:
    is_prime = eratosthenes_sieve(n)
    return [i for i in range(n) if is_prime(i)]

primes_under_100 = primes_under(100)
print(primes_under_100)

prefix_sums = [0] + list(itertools.accumulate(primes_under_100))
print(list(prefix_sums))

end = primes_under_100.index(7)
start = primes_under_100.index(3)
prefix_sums[end + 1] - prefix_sums[start]

def consecutive_prime_sums(n, debug=False) -> Generator[ConsecutivePrimeSum, None, None]:
    is_prime = eratosthenes_sieve(n)
    primes_under_n = primes_under(n)

    # prepend a 0 for nicer indexing
    prefix_sums = [0] + list(itertools.accumulate(primes_under_n))

    num_primes = len(primes_under_n)

    max_length_until_now = 0

    # Consider all chains starting from the ith prime
    for i in range(num_primes):
        # We iterate backwards for the end, since we are looking for the longest chain
        # Once we find a chain from [i,j], there is no point in looking for smaller j
        j_start = num_primes - 1

        # chains with a length smaller than the longest one we have found so far
        # will not contribute to the answer, so we can and should skip them
        j_end = max_length_until_now + i - 1

        for j in range(j_start, j_end, -1):
            chain_length = j - i + 1

            # prime_sum includes both the ith and jth prime
            consecutive_prime_sum = prefix_sums[j + 1] - prefix_sums[i]

            # if this sum is greater than the highest acceptable prime for this problem
            # skip it
            if consecutive_prime_sum > n:
                continue

            if is_prime(consecutive_prime_sum):
                max_length_until_now = max(max_length_until_now, chain_length)
                yield ConsecutivePrimeSum(consecutive_prime_sum, primes_under_n[i:j + 1])

                # Smaller j's with the same i will never give a better answer
                break

def solve(n, solver):
    return max(solver(n), key=lambda x: x.chain_length)

solve(100, solver=consecutive_prime_sums)

solve(1000, solver=consecutive_prime_sums)

from tabulate import tabulate

def test_performance(solver):
    # Simple geometric progression to see how the algorithm scales
    ns = [100 * 2**(n - 1) for n in range(1, 14)]

    data = []
    for n in ns:
        start = timeit.default_timer()
        ans = solve(n, solver)
        end = timeit.default_timer()
        data.append([n, round(end-start, 6), ans.prime])
    return data

data = test_performance(consecutive_prime_sums)

print(tabulate(data, headers=["n", "time (s)", "answer"], tablefmt="fancy_grid"))

primes_under_million = primes_under(1000_000)

len(primes_under_million)

primes_under_million[-10:]

k = 0
running_sum = 0
while running_sum < 1000_000:
    running_sum += primes_under_million[k]
    k += 1

k

def faster_consecutive_prime_sums(n, debug=False) -> Generator[ConsecutivePrimeSum, None, None]:
    is_prime = eratosthenes_sieve(n)
    primes_under_n = primes_under(n)

    prefix_sums = [0] + list(itertools.accumulate(primes_under_n))

    # Since prefix_sums is sorted by nature, use binary search to find n
    max_possible_chain_length = bisect.bisect_left(prefix_sums, n) + 1
    if debug:
        print(f"Max possible chain length for {n} is {max_possible_chain_length}")

    num_primes = len(primes_under_n)

    max_length_until_now = 0

    # Consider all chains starting from the ith prime
    for i in range(num_primes):
        j_start = min(i + max_possible_chain_length, num_primes - 1)
        j_end = max(max_length_until_now + i - 1, 0)

        for j in range(j_start, j_end, -1):
            chain_length = j - i + 1

            consecutive_prime_sum = prefix_sums[j + 1] - prefix_sums[i]

            if consecutive_prime_sum > n:
                continue

            if is_prime(consecutive_prime_sum):
                max_length_until_now = max(max_length_until_now, chain_length)
                yield ConsecutivePrimeSum(consecutive_prime_sum, primes_under_n[i:j + 1])
                break

data = test_performance(faster_consecutive_prime_sums)
print(tabulate(data, headers=["n", "time (s)", "answer"], tablefmt="fancy_grid"))

%%timeit -n3 -r3
solve(1_000_000, solver=faster_consecutive_prime_sums)