from lib_prime import primes, is_prime_rabin_miller
from functools import lru_cache


def first_approach():
    """ This was my original naiv approach that worked, but was way to slow
    because we computed the whole hex tile grid. """
    ps = set(primes(1000000))
    hextiles = {(0, 0): 1}
    add_ring(1, hextiles)
    add_ring(2, hextiles)

    # Visualization of how the coord system for the hex tile grid works.
    #                 
    #     2 1 0 1 2  
    #  4      8          
    #  3    9  19        
    #  2 10   2  18      
    #  1    3   7        
    #  0 11   1  17      
    #  1    4   6        
    #  2 12   5  16      
    #  3   13  15        
    #  4     14          
    #                 
    #                 

    pds = [1, 2, 8]
    for n in range(0, 100):
        ring_coord = ring_start(n)
        print("  ", ring_coord, hextiles[ring_coord])
        if n % 10 == 0:
            print(n)
        add_ring(n + 1, hextiles)
        for coord in ring_coords(n):
            pdv = pd(coord, hextiles, ps)
            if pdv == 3:
                v = hextiles[coord]
                assert v == first_number_ring(n) or v == first_number_ring(n + 1) - 1
                pds.append(hextiles[coord])
    print(len(pds))
    print(sorted(pds))
    target = 10
    if len(pds) >= target:
        return sorted(pds)[target - 1]
    return None


def ring_start(n):
    return (-n * 2, 0)


@lru_cache
def first_number_ring(n):
    if n == 0:
        return 1
    elif n == 1:
        return 2
    else:
        return first_number_ring(n - 1) + (n - 1) * 6


def get_nbvs_ring_start(n):
    #    8        1
    #  9   19   2   6   
    #    2        0  
    #  3   7    3   5
    #    l        4 
    v0 = first_number_ring(n)
    v1 = first_number_ring(n + 1)
    v2 = v1 + 1
    v3 = v0 + 1
    v4 = first_number_ring(n - 1)
    v5 = v1 - 1
    v6 = first_number_ring(n + 2) - 1
    nbvs = [v1, v2, v3, v4, v5, v6]
    # print(v0, nbvs)
    return nbvs


def pd_ring_start(n):
    v0 = first_number_ring(n)
    nbvs = get_nbvs_ring_start(n)
    r = 0
    for v in nbvs:
        if is_prime_rabin_miller(abs(v0 - v)):
            r += 1
    return r


def get_nbvs_ring_prev(n):
    """ Get neighbors and values for tile before top tile. """
    v0 = first_number_ring(n + 1) - 1
    v1 = first_number_ring(n + 2) - 1
    v2 = first_number_ring(n)
    v3 = first_number_ring(n - 1)
    v4 = v2 - 1
    v5 = first_number_ring(n + 1) - 2
    v6 = v1 - 1
    nbvs = [v1, v2, v3, v4, v5, v6]
    # print(v0, nbvs)
    return nbvs


def pd_ring_prev(n):
    v0 = first_number_ring(n + 1) - 1
    nbvs = get_nbvs_ring_prev(n)
    r = 0
    for v in nbvs:
        if is_prime_rabin_miller(abs(v0 - v)):
            r += 1
    return r


def add_ring(n, numbers):
    """ Adds ring n coords and values to numbers. """
    if n == 0:
        return
    first_coord = ring_start(n)
    current_number = first_number_ring(n)
    numbers[first_coord] = current_number
    current_coord = tuple(first_coord)

    for ro, co in [(1, -1), (2, 0), (1, 1), (-1, 1), (-2, 0), (-1, -1)]:
        for _ in range(n):
            current_coord = (current_coord[0] + ro, current_coord[1] + co)
            current_number += 1
            numbers[current_coord] = current_number

    # Reset first coord which is overriden.
    numbers[first_coord] = first_number_ring(n)


def get_neighbor_values(coord, numbers):
    neighbors = []

    for ro, co in [(-2, 0), (-1, 1), (1, 1), (2, 0), (1, -1), (-1, -1)]:
        nc = (coord[0] + ro, coord[1] + co)
        neighbors.append(numbers[nc])
    return neighbors


def ring_coords(n):
    """ Returns coords for ring n. """
    if n == 0:
        yield (0, 0)
    current_coord = ring_start(n)
    for ro, co in [(1, -1), (2, 0), (1, 1), (-1, 1), (-2, 0), (-1, -1)]:
        for _ in range(n):
            current_coord = (current_coord[0] + ro, current_coord[1] + co)
            yield current_coord


def pd(coord, numbers, primes):
    prime_delta = 0
    v = numbers[coord] 
    for nbv in get_neighbor_values(coord, numbers):
        if abs(v - nbv) in primes:
            prime_delta += 1
    return prime_delta


def test_get_nbvs_functions():
    """ Make sure that the function to compute the neighbor values works
    correctly by comparing the values to the original grid based approach. """
    hextiles = {(0, 0): 1}
    add_ring(1, hextiles)
    add_ring(2, hextiles)

    for n in range(2, 30):
        add_ring(n + 1, hextiles)
        nbvs1 = get_nbvs_ring_start(n)
        ring_coord = ring_start(n)
        nbvs2 = get_neighbor_values(ring_coord, hextiles)
        assert sorted(nbvs1) == sorted(nbvs2)

        nbvs1 = get_nbvs_ring_prev(n)
        ring_coord = (ring_start(n)[0] + 1, ring_start(n)[1] + 1)
        nbvs2 = get_neighbor_values(ring_coord, hextiles)
        assert sorted(nbvs1) == sorted(nbvs2)


def euler_128():
    # Conjecture: PD3s only occur at ring starts or at ring start minus one.
    # Under this assumption, we can significantly reduce the search space.
    # The only challenge is to find a way to directly compute the relevant
    # coords and neighbor values.

    test_get_nbvs_functions()

    target = 2000
    pds = [1, 2]
    for n in range(2, 80000):
        if pd_ring_start(n) == 3:
            v0 = first_number_ring(n)
            pds.append(v0)

        if pd_ring_prev(n) == 3:
            v0 = first_number_ring(n + 1) - 1
            pds.append(v0)

    assert len(pds) > target
    return sorted(pds)[target - 1]


if __name__ == "__main__":
    solution = euler_128()
    print("e128.py: " + str(solution))
    assert(solution == 14516824220)