#!/usr/bin/env python3
"""Audit the global-depth edge of THE PRIMITIVE-ROOT ENGINE.

The earlier tower instrument pools branch values from every explored depth.
That is a useful census, but pooling can hide a biased outer layer.  This
instrument measures each layer separately and prints the pooled census beside
the last layer, so the two claims cannot be confused.

The reverse Collatz tree is rooted at 1.  Every value v has the child 2v; a
value v == 4 (mod 6) also has (v-1)/3, except that v=4 would recreate the
root's 1-2-4 cycle and is therefore cut.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass

import numpy as np


DEFAULT_MODULI = (18, 54, 162, 486, 1458)
UINT64_MAX = np.iinfo(np.uint64).max


@dataclass(frozen=True)
class Weather:
    minimum: float
    maximum: float
    total_variation: float


def weather(branches: np.ndarray, modulus: int) -> Weather:
    """Return multiplicative min/max and TV distance from uniform."""
    if modulus < 6 or modulus % 6:
        raise ValueError(f"modulus must be a positive multiple of 6: {modulus}")
    residues = np.asarray(branches % modulus, dtype=np.int64)
    counts = np.bincount(residues, minlength=modulus)[4::6]
    if counts.sum() == 0:
        return Weather(float("nan"), float("nan"), float("nan"))
    scaled = counts / counts.sum() * counts.size
    tv = 0.5 * np.abs(scaled - 1).sum() / counts.size
    return Weather(float(scaled.min()), float(scaled.max()), float(tv))


def next_layer(level: np.ndarray) -> np.ndarray:
    if int(level.max()) > UINT64_MAX // 2:
        raise OverflowError("uint64 horizon reached; lower --depth")
    branch = (level % 6 == 4) & (level != 4)
    doubled = level * np.uint64(2)
    odd_children = ((level[branch] - 1) // 3).astype(np.uint64)
    return np.concatenate((doubled, odd_children))


def audit_small_tree(depth: int) -> None:
    """Independently reproduce the early tree with Python integer sets."""
    exact = {1}
    fast = np.array([1], dtype=np.uint64)
    for _ in range(depth):
        nxt: set[int] = set()
        for value in exact:
            nxt.add(2 * value)
            if value % 6 == 4 and value != 4:
                nxt.add((value - 1) // 3)
        fast = next_layer(fast)
        assert len(fast) == len(set(map(int, fast))), "duplicate child in array tree"
        assert set(map(int, fast)) == nxt, "array tree differs from exact set tree"
        exact = nxt


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--depth", type=int, default=63)
    parser.add_argument("--tail", type=int, default=6)
    parser.add_argument("--audit-depth", type=int, default=24)
    parser.add_argument("--moduli", type=int, nargs="+", default=DEFAULT_MODULI)
    args = parser.parse_args()
    if not 1 <= args.depth <= 63:
        parser.error("--depth must be between 1 and 63 for exact uint64 arithmetic")
    if args.tail < 1:
        parser.error("--tail must be positive")

    audit_small_tree(min(args.audit_depth, args.depth))
    print(f"independent set audit passed through depth {min(args.audit_depth, args.depth)}")

    level = np.array([1], dtype=np.uint64)
    pooled: list[np.ndarray] = []
    history: list[tuple[int, int, int, dict[int, Weather]]] = []
    for depth in range(args.depth + 1):
        branch = level[(level % 6 == 4) & (level != 4)]
        if branch.size:
            pooled.append(branch.copy())
        if depth >= max(0, args.depth - args.tail + 1):
            history.append(
                (
                    depth,
                    int(level.size),
                    int(branch.size),
                    {modulus: weather(branch, modulus) for modulus in args.moduli},
                )
            )
        if depth < args.depth:
            level = next_layer(level)

    print("depth,nodes,branches,modulus,min*K,max*K,TV")
    for depth, nodes, branches, observations in history:
        for modulus, result in observations.items():
            print(
                f"{depth},{nodes},{branches},{modulus},"
                f"{result.minimum:.9f},{result.maximum:.9f},"
                f"{result.total_variation:.9f}"
            )

    all_branches = np.concatenate(pooled)
    final_branches = level[(level % 6 == 4) & (level != 4)]
    print("\nouter layer versus pooled depths (pooling is the older statistic)")
    print("modulus,layer_TV,pooled_TV,layer_min*K,layer_max*K")
    for modulus in args.moduli:
        layer_result = weather(final_branches, modulus)
        pooled_result = weather(all_branches, modulus)
        print(
            f"{modulus},{layer_result.total_variation:.9f},"
            f"{pooled_result.total_variation:.9f},"
            f"{layer_result.minimum:.9f},{layer_result.maximum:.9f}"
        )

    print(
        "\nVERDICT: pooled calm and layerwise calm are different claims. "
        "The direct outer layers move near uniformity in this finite survey, "
        "but the depth-limit remains unproved."
    )


if __name__ == "__main__":
    main()