"""Canonical orbit search for subset masks under exact channel symmetries."""
from itertools import combinations, permutations, product
from time import perf_counter


def lex_less(a, b):
    """The paper's order: min(a triangle b) belongs to a."""
    d = set(a).symmetric_difference(b)
    return bool(d) and min(d) in set(a)


def apply_perm(mask, perm):
    return tuple(sorted(perm[i] for i in mask))


def all_permutations(n):
    return list(permutations(range(n)))


def product_group(blocks):
    """All permutations independently permuting the supplied blocks."""
    n = max(x for b in blocks for x in b) + 1
    block_perms = [list(permutations(b)) for b in blocks]
    group = []
    for choices in product(*block_perms):
        p = list(range(n))
        for block, image in zip(blocks, choices):
            for x, y in zip(block, image):
                p[x] = y
        group.append(tuple(p))
    return group


def canonical(mask, group):
    mask = tuple(sorted(mask))
    return min(apply_perm(mask, p) for p in group)


def is_canonical(mask, group):
    mask = tuple(sorted(mask))
    return canonical(mask, group) == mask


def exhaustive(n, k):
    return list(combinations(range(n), k))


def canonical_dfs(n, k, group):
    found = []
    cache = {}
    def keep(mask):
        mask = tuple(mask)
        if mask not in cache:
            cache[mask] = is_canonical(mask, group)
        return cache[mask]
    def dfs(mask):
        if len(mask) == k:
            found.append(mask)
            return
        start = mask[-1] + 1 if mask else 0
        for i in range(start, n):
            child = mask + (i,)
            if keep(child):
                dfs(child)
    dfs(())
    return found


def orbit(mask, group):
    return {apply_perm(mask, p) for p in group}


def verify(n=7, k=3):
    group = all_permutations(n)
    masks = exhaustive(n, k)
    reps = {canonical(s, group) for s in masks}
    dfs = set(canonical_dfs(n, k, group))
    closure = all(is_canonical(s[:-1], group) for s in reps if s)
    covered = set().union(*(orbit(r, group) for r in reps))
    # Also directly check the stated symmetric-difference definition.
    order_check = all((a == b or lex_less(a, b) != lex_less(b, a))
                      for a in masks for b in masks)
    return {
        'n_masks': len(masks), 'group_size': len(group),
        'orbit_count': len(reps), 'dfs_count': len(dfs),
        'all_orbits_covered': covered == set(masks),
        'dfs_equals_reps': dfs == reps,
        'predecessor_closure': closure,
        'lex_order_antisymmetric_check': order_check,
    }


def benchmark(n=8, k=4, repeats=2):
    group = all_permutations(n)
    masks = exhaustive(n, k)
    # Symmetry-invariant one-shot score; this isolates search reduction.
    scores = {s: float(len(s)) for s in masks}
    t0 = perf_counter()
    best_ex = max((scores[s], s) for s in masks)
    exhaustive_time = perf_counter() - t0
    times = []
    for _ in range(repeats):
        t0 = perf_counter()
        reps = canonical_dfs(n, k, group)
        best_can = max((scores[s], s) for s in reps)
        times.append(perf_counter() - t0)
    return {
        'n': n, 'k': k, 'total_masks': len(masks), 'canonical_masks': len(reps),
        'evaluated_baseline': len(masks), 'evaluated_idea': len(reps),
        'best_baseline': best_ex, 'best_idea': best_can,
        'exhaustive_score_time_sec': exhaustive_time,
        'canonical_search_times_sec': times,
        'speedup_in_evaluations': len(masks) / len(reps),
    }

if __name__ == '__main__':
    print('VERIFY', verify())
    print('BENCHMARK', benchmark())