"""
# Summable Primes

Geometric Invariant Hash Functions are relevant to grid worlds experiencing wraparound ([closed manifolds](https://en.wikipedia.org/wiki/Closed_manifold)) in [Conway's Reverse Game of Life](https://www.kaggle.com/c/conway-s-reverse-game-of-life) and [Halite](https://www.kaggle.com/c/halite) and are also relevant to the [Abstraction and Reasoning Corpus](https://www.kaggle.com/c/abstraction-and-reasoning-challenge). They have the property of being invariant to `np.roll()` and optionally `np.flip()` and `np.rot90()`

To achieve this, we also need access to a set of Summable Prime Numbers. 

According to the [Unique Factorization Theorem](https://en.wikipedia.org/wiki/Fundamental_theorem_of_arithmetic) the product of any combination of primes results in a unique number. This is not guaranteed to be true for summation, but it is possible through search to find a subset of prime numbers for which this property holds true. This property is important for preventing hash collisions.

However for the purpose of hashing functions, it may also be useful to have a set of summable prime numbers, such that the sum of any combination of them will also produce a unique number.

This notebook will use the [Z3 Theorm Prover](https://github.com/Z3Prover/z3) to prove that such a set exists.

This research is required for the following notebooks:  
- https://www.kaggle.com/jamesmcguigan/geometric-invariant-hash-functions
- https://www.kaggle.com/jamesmcguigan/game-of-life-hashmap-solver
"""
# TODO: add z3-solver to kaggle-docker image
! pip3 install -q z3-solver
import z3
import numpy as np
import itertools
import sys
import math
import time
import timeit
import signal
from typing import List
from numba import njit

notebook_start = time.perf_counter()
"""
# Generating Prime Numbers

A simple but consise way of generating small prime numbers
"""
def generate_primes(count):
    primes = [2]
    for n in range(3, sys.maxsize, 2):
        if len(primes) >= count: break
        if all( n % i != 0 for i in range(3, int(math.sqrt(n))+1, 2) ):
            primes.append(n)
    return primes


#      10 primes generated in:    0.01 ms
#     100 primes generated in:    0.32 ms
#    1000 primes generated in:    6.99 ms
#  10,000 primes generated in:  153.21 ms
# 100,000 primes generated in: 4831.97 ms
print( generate_primes(42) )
print( f'     10 primes generated in: {timeit.timeit(lambda: generate_primes(     10), number=10)/10*1000:7.2f} ms' )
print( f'    100 primes generated in: {timeit.timeit(lambda: generate_primes(    100), number=10)/10*1000:7.2f} ms' )
print( f'   1000 primes generated in: {timeit.timeit(lambda: generate_primes(  1_000), number=10)/10*1000:7.2f} ms' )
print( f' 10,000 primes generated in: {timeit.timeit(lambda: generate_primes( 10_000), number=10)/10*1000:7.2f} ms' )
# print( f'100,000 primes generated in: {timeit.timeit(lambda: generate_primes(100_000), number=10)/10*1000:7.2f} ms' )
"""
# Z3 Prime Numbers

This can also be implemented in Z3 using first order logic.

Note that Z3 only promises to return from the set of all primes, which are not always the first N primes. Z3 is also a couple of orders of magnitude slower than our simple loop solution.

Many of the additional constraints added to this solver, such as `x % 1 == 0` (is int), `x % 2 != 0` (is even), and iteratively increasing the `domain` where required to improve the performance of the solver. 

It is also much more performant, and less likely to produce `unsat`, to repeatedly solve for a single prime, and iteratively add constraints that exclude previously found primes, than it is to pass use a `z3.Distinct()` clause on an array of `Ints`.
"""
def z3_is_prime(x):
    y = z3.Int("y")
    return z3.Or([
        x == 2,              # 2 is the only even prime
        z3.And(
            x > 1,           # x is positive
            x % 1 == 0,      # x is int
            x % 2 != 0,      # x is odd
            z3.Not(z3.Exists([y], z3.And(
                y   > 1,     # y is positive
                y*y < x,     # y < sqrt(x)
                y % 2 != 0,  # y is odd
                x % y == 0   # y is not a divisor of x
            )))
        )
    ])


def z3_generate_primes(count):
    output = []

    number = z3.Int('prime')
    solver = z3.Solver()
    solver.add([ number > 1, number % 1 == 0 ])  # is positive int
    solver.add([ z3_is_prime(number) ])          # must be a prime, obviously
    solver.push()

    domain = 2
    while len(output) < count:
        solver.pop()
        solver.push()
        solver.add([ number < domain ])  # this helps prevent unsat
        solver.add(z3.And([ number != value for value in output ]))
        while solver.check() == z3.sat:
            value = solver.model()[number].as_long()
            solver.add([ number != value ])
            output.append( value )
            if len(output) >= count: break
        domain *= 2  # increment search space
    return sorted(output)

assert len(z3_generate_primes(24)) == 24
print( 'z3: ', z3_generate_primes(24) )
print( 'std:', generate_primes(24) )


#   2 primes generated in:    12.8 ms
#   4 primes generated in:    17.8 ms
#   8 primes generated in:    38.0 ms
#  16 primes generated in:    86.8 ms
#  32 primes generated in:   246.1 ms
#  64 primes generated in:  1655.1 ms
# 128 primes generated in: 14740.5 ms
for n in [2,4,8,16,32,64,128]:
    print( f'{n:3d} primes generated in: {timeit.timeit(lambda: z3_generate_primes(n), number=1)/1*1000:7.1f} ms' )
"""
# Summable Primes

This generates a list of prime numbers, such that all combinations of the sum of these numbers produces a unique value.

Previous results:
```
size =  2 | combinations = 2 | time =     0.0s |  [2, 3]

size =  3 | combinations = 2 | time =     0.0s |  [2, 3, 7]
size =  3 | combinations = 3 | time =     0.0s |  [2, 3, 7]

size =  4 | combinations = 2 | time =     0.0s |  [3, 7, 11, 13]
size =  4 | combinations = 3 | time =     0.0s |  [2, 3, 7, 13]
size =  4 | combinations = 4 | time =     0.0s |  [2, 3, 7, 13]

size =  5 | combinations = 2 | time =     0.1s |  [7, 17, 19, 23, 31]
size =  5 | combinations = 3 | time =     0.1s |  [7, 17, 23, 29, 31]
size =  5 | combinations = 4 | time =     0.2s |  [7, 13, 17, 19, 31]
size =  5 | combinations = 5 | time =     0.3s |  [2, 7, 19, 23, 31]

size =  6 | combinations = 2 | time =     0.1s |  [2, 3, 11, 19, 23, 29]
size =  6 | combinations = 3 | time =     0.1s |  [3, 13, 17, 23, 29, 31]
size =  6 | combinations = 4 | time =     0.1s |  [3, 13, 17, 23, 29, 31]
size =  6 | combinations = 5 | time =     0.1s |  [3, 13, 17, 23, 29, 31]
size =  6 | combinations = 6 | time =     0.2s |  [3, 13, 17, 23, 29, 31]

size =  7 | combinations = 2 | time =     1.1s |  [5, 7, 13, 17, 29, 43, 61]
size =  7 | combinations = 3 | time =     1.7s |  [3, 17, 31, 37, 41, 43, 59]
size =  7 | combinations = 4 | time =     3.1s |  [2, 13, 67, 103, 107, 113, 127]
size =  7 | combinations = 5 | time =     3.3s |  [2, 3, 23, 31, 79, 83, 127]
size =  7 | combinations = 6 | time =     3.4s |  [2, 5, 11, 29, 59, 103, 113]
size =  7 | combinations = 7 | time =     3.8s |  [3, 11, 13, 17, 29, 61, 113] 

size =  8 | combinations = 2 | time =     0.3s |  [2, 3, 7, 17, 23, 29, 41, 59]
size =  8 | combinations = 3 | time =     1.9s |  [3, 5, 7, 11, 17, 53, 83, 109]
size =  8 | combinations = 4 | time =  1213.3s |  [3, 31, 53, 67, 71, 97, 103, 109]
size =  8 | combinations = 5 | time =  1221.1s |  [3, 29, 61, 71, 89, 101, 107, 113]
size =  8 | combinations = 6 | time =   492.7s |  [3, 31, 53, 71, 89, 97, 103, 109]
size =  8 | combinations = 7 | time =   355.4s |  [3, 31, 53, 71, 89, 97, 103, 109]
size =  8 | combinations = 8 | time =  4217.6s |  [3, 61, 71, 89, 101, 107, 113, 127]

size =  9 | combinations = 2 | time =     1.8s |  [11, 19, 23, 71, 73, 97, 107, 113, 127]
size =  9 | combinations = 3 | time =  2302.3s |  [3, 67, 71, 89, 97, 131, 197, 241, 251]
size =  9 | combinations = 4 | time =  8796.2s |  [2, 5, 37, 331, 337, 347, 359, 431, 509]
size =  9 | combinations = 5 | time =  5388.6s |  [2, 3, 11, 23, 89, 137, 307, 487, 491]
size =  9 | combinations = 6 | time =  3644.8s |  [2, 31, 43, 61, 163, 167, 263, 449, 509]
size =  9 | combinations = 7 | time =   732.5s |  [3, 5, 7, 43, 71, 277, 421, 503, 509]
size =  9 | combinations = 9 | time =    15.1s |  [2, 3, 223, 337, 367, 449, 491, 499, 503]
size =  9 | combinations = 9 | time = 10130.9s |  [2, 29, 37, 73, 97, 101, 149, 293, 487]
```
"""
def z3_generate_summable_primes(size=50, combinations=2, timeout=0) -> List[int]:
    candidates = [ z3.Int(n) for n in range(size) ]
    summations = []
    for n_combinations in range(1,combinations+1):
        summations += [
            z3.Sum(group)
            for group in itertools.combinations(candidates, n_combinations)
        ]

    solver = z3.Solver()
    if timeout: solver.set("timeout", int(timeout * 1000/2.5))  # timeout is in milliseconds, but inexact and ~2.5x slow

    solver.add([ num > 1      for num in candidates ])
    solver.add([ num % 1 == 0 for num in candidates ])
    solver.add([ candidates[n] < candidates[n+1] for n in range(len(candidates)-1) ])  # sorted

    # solver.add([ z3_is_prime(num) for num in candidates ])
    # primes = z3_generate_primes(128)

    solver.add( z3.Distinct(candidates) )
    solver.add( z3.Distinct(summations) )
    solver.push()
    domain = 2
    while True:
        solver.pop()
        solver.push()

        primes = generate_primes(domain)
        solver.add([ z3.Or([ num == prime for prime in primes ])
                     for num in candidates ])
        solver.add([ num < domain for num in candidates ])

        if solver.check() != z3.sat:
            domain *= 2
        else:
            values = sorted([ solver.model()[num].as_long() for num in candidates ])
            return list(values)
try:
    timeout = 5*60*60
    def raise_timeout(signum, frame): raise TimeoutError    # DOC: https://docs.python.org/3.6/library/signal.html
    signal.signal(signal.SIGALRM, raise_timeout)            # Register a function to raise a TimeoutError on the signal.
    signal.alarm(timeout)                                   # Schedule the signal to be sent after ``time``.
    
    for size in range(10,50+1):
        for combinations in [size]:  # range(2,size+1):
            time_start      = time.perf_counter()
            hashable_primes = z3_generate_summable_primes(size=size, combinations=combinations, timeout=timeout)
            time_taken      = time.perf_counter() - time_start
            print(f'size = {size:2d} | combinations = {combinations} | time = {time_taken:7.1f}s | ', hashable_primes)        
        print()
except   TimeoutError as exception: print('timeout')
finally: signal.alarm(0)
"""
# Hashable Primes

As Z3 is very slow at generating large lists of Summable Primes, and that for [Geometric Invariant Hash Functions](https://www.kaggle.com/jamesmcguigan/geometric-invariant-hash-functions) we need a list at least twice as large as the board size (ie 50 to hash a 25x25 board), we need a quicker method to generate useful numbers.

For the purposes of this implemention, we reduce the constraint to that of being no pairwise sum collisions between numbers in the set, rather than requiring the sum of any combinations of numbers to be unique.
"""
primes_np = np.array( generate_primes(100_000), dtype=np.int64 )

@njit()
def generate_hashable_primes(size=50, combinations=2) -> np.ndarray:
    """
    Return a list of primes that have no summation collisions for N=2, for use in hashing
    NOTE: size > 50 or combinations > 2 produces no results
    """
    domain     = primes_np
    candidates = np.zeros((0,), dtype=np.int64)
    exclusions = set()
    while len(candidates) < size:  # loop until we have successfully filled the buffer
        # refill candidates ignoring exclusions
        for n in range(len(domain)):
            prime = np.int64( domain[n] )
            if not np.any( candidates == prime ):
                if prime not in exclusions:
                    candidates = np.append( candidates, prime )
                    if len(candidates) >= size: break
        else:
            return np.zeros((0,), dtype=np.int64)  # prevent infinite loop if we run out of primes

        # This implements itertools.product(*[ candidates, candidates, candidates ]
        collisions = set(candidates)
        indexes    = np.array([ 0 ] * combinations)
        while np.min(indexes) < len(candidates)-1:

            # Sum N=combinations candidate primes and check for duplicate collisions, then exclude and try again
            values = candidates[indexes]
            summed = np.sum(values)
            if summed in collisions:  # then remove the largest conflicting number from the list of candidates
                exclude    = np.max(candidates[indexes])
                candidates = candidates[ candidates != exclude ]
                exclusions.add(exclude)
                break  # pick a new set of candidates and try again
            collisions.add(summed)

            # This implements itertools.product(*[ candidates, candidates, candidates ]
            indexes[0] += 1
            for i in range(len(indexes)-1):
                while np.count_nonzero( indexes == indexes[i] ) >= 2:
                    indexes[i] += 1                    # ensure indexes are unique
                if indexes[i] >= len(candidates):      # overflow to next buffer
                    indexes[i]    = np.min(indexes)+1  # for triangular iteration
                    indexes[i+1] += 1
    return candidates
combinations = 2
for size in [5, 25, 50, 100, 150]:
    time_start = time.perf_counter()
    hashable_primes = generate_hashable_primes(size=size, combinations=combinations)
    time_taken = time.perf_counter() - time_start
    print(f'combinations = {combinations} | size = {size} | time = {time_taken:.0f}s\n', hashable_primes, '\n')
    if len(hashable_primes) == 0: break
        
combinations = 3
for size in range(2,25):
    time_start = time.perf_counter()
    hashable_primes = generate_hashable_primes(size=size, combinations=combinations)
    time_taken = time.perf_counter() - time_start
    print(f'combinations = {combinations} | size = {size} | time = {time_taken:.0f}s\n', hashable_primes, '\n')
    if len(hashable_primes) == 0: break