"""

Triangle, square, pentagonal, hexagonal, heptagonal, and octagonal numbers are all figurate (polygonal) numbers and are generated by the following formulae:
Triangle
 
P3,n=n(n+1)/2
 
1, 3, 6, 10, 15, ...
Square
 
P4,n=n2
 
1, 4, 9, 16, 25, ...
Pentagonal
 
P5,n=n(3n−1)/2
 
1, 5, 12, 22, 35, ...
Hexagonal
 
P6,n=n(2n−1)
 
1, 6, 15, 28, 45, ...
Heptagonal
 
P7,n=n(5n−3)/2
 
1, 7, 18, 34, 55, ...
Octagonal
 
P8,n=n(3n−2)
 
1, 8, 21, 40, 65, ...
The ordered set of three 4-digit numbers: 8128, 2882, 8281, has three interesting properties.
The set is cyclic, in that the last two digits of each number is the first two digits of the next number (including the last number with the first).
Each polygonal type: triangle (P3,127=8128), square (P4,91=8281), and pentagonal (P5,44=2882), is represented by a different number in the set.
This is the only set of 4-digit numbers with this property.
Find the sum of the only ordered set of six cyclic 4-digit numbers for which each polygonal type: triangle, square, pentagonal, hexagonal, heptagonal, and octagonal, is represented by a different number in the set.

"""

import itertools as it


def take_between(predicate, iterable):
    """Dress iterator starting when first true and stopping thereafter when first false."""
    dropped = it.dropwhile(lambda x: not predicate(x), iterable)
    return it.takewhile(predicate, dropped)

pred = lambda x: len(str(x)) == 4

polys = {
    'triangles': set(take_between(pred, ((n * (n + 1)) // 2 for n in it.count()))),
    'squares': set(take_between(pred, (n ** 2 for n in it.count()))),
    'pents': set(take_between(pred, ((n * (3 * n - 1)) // 2 for n in it.count()))),
    'hexs': set(take_between(pred, (n * (2 * n - 1) for n in it.count()))),
    'hepts': set(take_between(pred, ((n * (5 * n - 3)) // 2 for n in it.count()))),
    'octs': set(take_between(pred, (n * (3 * n - 2) for n in it.count()))),
}

firstdigits = [str(i) for i in range(1, 10)]
seconddigits =[str(i) for i in range(10)]


def gen_cycles(num, ncycles=3):
    A, B, C, D = str(num)
    ABCD = num

    for firsts in it.product(firstdigits, repeat=ncycles-2):
        for second in it.product(seconddigits, repeat=ncycles-2):
            xyAB = int(firsts[0] + second[0] + A + B)
            CDxy = int(C + D + firsts[0] + second[0])
            yield (ABCD, xyAB, CDxy)


def one_to_one(sets, vals):
    """True if each val is in exactly one set and each set contains exactly one value"""
    for s in sets:
        if len(s & set(vals)) != 1:
            return False
    for v in vals:
        if sum(v in s for s in sets) != 1:
            return False
    return True

for tri in polys['triangles']:
    for cycle in gen_cycles(tri):
        if one_to_one(polys.values(), cycle):
            print(cycle)

