#!/usr/bin/env python3
"""Finite regression checks for the common-horizon block-selector theorem.

The script is deliberately not a proof assistant.  It exhaustively checks the
small instances that are most useful for catching indexing, spacing, and
quantifier mistakes in the paper proof:

* every special horizon has 2**d zero-track rows and one-track width one;
* the tuple of all special-horizon rows recovers all d*R data bits;
* an H-cone meets at most one separated block;
* k cones expose at most d*min(k,R) data bits;
* the claimed upper construction is simultaneously resolving; and
* the staged total target has disjoint output-length windows and a
  two-snapshot/two-track plan at every exact cut.

Run with:

    python3 check_horizon_join.py
"""

from __future__ import annotations

from itertools import combinations_with_replacement, product
from math import log2


def blocks_for(d: int, count: int, gap: int) -> tuple[tuple[int, ...], ...]:
    """Return count length-d blocks whose closest cells are gap+1 apart."""
    stride = d + gap
    return tuple(
        tuple(range(t * stride, t * stride + d))
        for t in range(count)
    )


def cone(root: int, horizon: int) -> frozenset[int]:
    if root < 0:
        return frozenset()
    return frozenset(range(max(0, root - horizon), root + 1))


def selector_row(pattern: tuple[int, ...], block: tuple[int, ...]) -> tuple[int, ...]:
    return tuple(pattern[position] for position in block)


def joint_row(
    pattern: tuple[int, ...],
    blocks: tuple[tuple[int, ...], ...],
) -> tuple[int, ...]:
    return tuple(bit for block in blocks for bit in selector_row(pattern, block))


def visible_data_positions(
    roots: tuple[int, ...],
    horizon: int,
    data_positions: frozenset[int],
) -> frozenset[int]:
    seen: set[int] = set()
    for root in roots:
        seen.update(cone(root, horizon) & data_positions)
    return frozenset(seen)


def check_instance(d: int, count: int, max_k: int) -> None:
    assert d >= 1 and count >= 1
    # The special horizons are d, ..., d+count-1.
    horizons = tuple(d + t for t in range(count))
    horizon = max(horizons)
    # Closest informative cells in different blocks are horizon+1 apart.
    blocks = blocks_for(d, count, horizon)
    data_positions = frozenset(position for block in blocks for position in block)
    last_position = max(data_positions)

    for left, right in zip(blocks, blocks[1:]):
        assert min(abs(x - y) for x in left for y in right) > horizon

    patterns = []
    for bits in product((0, 1), repeat=d * count):
        pattern = [0] * (last_position + 1)
        for bit, position in zip(bits, sorted(data_positions)):
            pattern[position] = bit
        patterns.append(tuple(pattern))

    marginal_counts = []
    for block, exact_horizon in zip(blocks, horizons):
        rows = {selector_row(pattern, block) for pattern in patterns}
        marginal_counts.append(len(rows))
        assert len(rows) == 2**d
        assert set(block) <= cone(max(block), exact_horizon)

    joint_rows = {joint_row(pattern, blocks) for pattern in patterns}
    assert len(joint_rows) == 2 ** (d * count)

    roots = tuple(range(-1, last_position + 1))
    for root in roots:
        touched = [block for block in blocks if cone(root, horizon) & set(block)]
        assert len(touched) <= 1

    observed_frontier = []
    for k in range(max_k + 1):
        max_visible = 0
        # Head labels do not affect union size, so combinations suffice here.
        for layout in combinations_with_replacement(roots, k):
            visible = visible_data_positions(layout, horizon, data_positions)
            max_visible = max(max_visible, len(visible))
        expected_visible = d * min(k, count)
        assert max_visible == expected_visible

        expected_width = 2 ** (d * (count - min(k, count)))
        observed_frontier.append(expected_width)

        # Check the explicit upper construction.  Track the first min(k,R)
        # blocks and let the snapshot identity be the remaining block bits.
        tracked = blocks[: min(k, count)]
        hidden = blocks[min(k, count) :]
        cells: dict[tuple[int, ...], list[tuple[int, ...]]] = {}
        for pattern in patterns:
            key = tuple(pattern[position] for block in hidden for position in block)
            cells.setdefault(key, []).append(pattern)
        assert len(cells) == expected_width

        tracked_positions = frozenset(position for block in tracked for position in block)
        for members in cells.values():
            seen: set[tuple[int, ...]] = set()
            for pattern in members:
                projection = tuple(pattern[position] for position in sorted(tracked_positions))
                assert projection not in seen
                seen.add(projection)

    print(
        f"d={d}, R={count}, H={horizon}: "
        f"marginals={marginal_counts}, common={observed_frontier}"
    )


def check_staged_target(stages: int = 8) -> None:
    previous_end = -1
    for r in range(1, stages + 1):
        d = 2**r
        m = 4 * d * d
        horizon = 2 * d - 1
        blocks = tuple(
            tuple(range(3 * d * t, 3 * d * t + d))
            for t in range(d)
        )
        assert max(blocks[-1]) < m
        for left, right in zip(blocks, blocks[1:]):
            closest = min(right) - max(left)
            assert closest > horizon
        output_window = (m + d, m + 2 * d - 1)
        assert output_window[0] > previous_end
        previous_end = output_window[1]

        fixed_k_lower = d * (d - 1) - log2(m + 1)
        if r >= 2:
            assert fixed_k_lower > 0
    print(f"staged target: first {stages} windows are disjoint and geometrically legal")


def check_global_marginal_plans(stages: int = 6) -> None:
    checked = 0
    for r in range(1, stages + 1):
        d = 2**r
        cut_stage = 4 * d * d
        for t in range(d):
            total_length = cut_stage + d + t
            block_start = 3 * d * t
            block_end = block_start + d - 1
            for cut_position in range(total_length + 1):
                horizon = total_length - cut_position
                if cut_position <= cut_stage:
                    old_left = block_start
                    old_right = min(block_end, cut_position - 1)
                    old_count = max(0, old_right - old_left + 1)
                    assert old_count <= d
                    assert old_count == 0 or horizon + 1 >= old_count
                    snapshots, tracks = 1, int(old_count > 0)
                elif cut_position < cut_stage + r:
                    old_selector_count = cut_position - cut_stage
                    assert 1 <= old_selector_count < r
                    assert horizon + 1 >= d // 2
                    assert horizon + 1 >= old_selector_count
                    snapshots, tracks = 2, 2
                else:
                    snapshots, tracks = 2, 0
                assert snapshots <= 2 and tracks <= 2
                checked += 1
    print(
        f"global marginal plans: {checked} staged cut/horizon cases "
        f"fit within two snapshots and two tracks"
    )


def main() -> None:
    # Full row enumeration is 2**(dR), so keep these deliberately small.
    check_instance(d=1, count=2, max_k=2)
    check_instance(d=2, count=2, max_k=2)
    check_instance(d=2, count=3, max_k=3)
    check_instance(d=3, count=3, max_k=3)
    check_staged_target()
    check_global_marginal_plans()
    print("all common-horizon checks passed")


if __name__ == "__main__":
    main()