#!/usr/bin/env python3
"""Audit THE FERTILE RAY and survey the reverse Collatz tree by depth.

The forward map is T(n)=n/2 for even n and T(n)=3n+1 for odd n.
Reverse children of v are therefore 2v and, only when v == 4 (mod 6),
(v-1)/3.  The familiar 1-2-4 cycle is cut by refusing already-seen nodes,
so the survey is the first-discovery tree rooted at 1.

Two exact facts checked below also have short derivations:

1. On a doubling ray, multiples of 3 stay 0 (mod 6).  Every other ray's
   even residues alternate 2,4 (mod 6), hence branch every other step.
   For a branch value v == 4,10,16 (mod 18), (v-1)/3 == 1,3,5 (mod 6).
2. After cutting the root cycle, every depth-d node has one unique even
   child 2v, and every branch node has one additional unique odd child.
   Even and odd children cannot collide.  Thus N[d+1] = N[d] + B[d].
"""

from __future__ import annotations

import argparse
from collections import Counter


def children(v: int) -> tuple[int, ...]:
    out = [2 * v]
    if v % 6 == 4:
        out.append((v - 1) // 3)
    return tuple(out)


def audit_local_law(limit: int, ray_steps: int) -> None:
    for v in range(1, limit + 1):
        x = v
        flags = []
        for _ in range(ray_steps):
            x *= 2
            flags.append(x % 6 == 4)
        if v % 3 == 0:
            assert not any(flags)
        else:
            assert all(flags[i] == flags[i + 2] for i in range(ray_steps - 2))
            assert abs(sum(flags) * 2 - ray_steps) <= 1

    for v in range(4, limit + 1, 6):
        u = (v - 1) // 3
        expected = {4: 1, 10: 3, 16: 5}[v % 18]
        assert u % 6 == expected


def survey(depth: int) -> list[dict[str, int | float]]:
    layer = {1}
    seen = {1}
    rows: list[dict[str, int | float]] = []

    for d in range(depth + 1):
        counts = Counter(v % 3 for v in layer)
        branches = sum(v % 6 == 4 for v in layer)
        total = len(layer)
        rows.append(
            {
                "depth": d,
                "nodes": total,
                "fertile": counts[1] + counts[2],
                "sterile": counts[0],
                "branch_nodes": branches,
                "fertile_share": (counts[1] + counts[2]) / total,
                "branch_share": branches / total,
            }
        )

        nxt: set[int] = set()
        for v in layer:
            for child in children(v):
                if child not in seen:
                    seen.add(child)
                    nxt.add(child)
        if d >= 3:
            assert len(nxt) == total + branches
        layer = nxt

    return rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--depth", type=int, default=60)
    parser.add_argument("--audit-limit", type=int, default=100_000)
    parser.add_argument("--ray-steps", type=int, default=20)
    parser.add_argument("--tail", type=int, default=20)
    args = parser.parse_args()

    audit_local_law(args.audit_limit, args.ray_steps)
    rows = survey(args.depth)

    print("depth,nodes,fertile,sterile,branch_nodes,fertile_share,branch_share")
    for row in rows[-args.tail :]:
        print(
            f"{row['depth']},{row['nodes']},{row['fertile']},"
            f"{row['sterile']},{row['branch_nodes']},"
            f"{row['fertile_share']:.9f},{row['branch_share']:.9f}"
        )

    even = [r["nodes"] for r in rows if r["depth"] >= 20 and r["depth"] % 2 == 0]
    odd = [r["nodes"] for r in rows if r["depth"] >= 21 and r["depth"] % 2 == 1]
    tail_rows = rows[max(1, len(rows) - args.tail) :]
    fertile_shares = [float(r["fertile_share"]) for r in tail_rows]
    branch_shares = [float(r["branch_share"]) for r in tail_rows]
    growth = [
        int(rows[i]["nodes"]) / int(rows[i - 1]["nodes"])
        for i in range(max(1, len(rows) - args.tail), len(rows))
    ]
    print(f"audited local law through v={args.audit_limit}, {args.ray_steps} doublings")
    print(f"nodes through depth {args.depth}: {sum(int(r['nodes']) for r in rows)}")
    print(f"last even-depth counts: {even[-8:]}")
    print(f"last odd-depth counts: {odd[-8:]}")
    print(
        "tail fertile share min/mean/max: "
        f"{min(fertile_shares):.9f}/"
        f"{sum(fertile_shares)/len(fertile_shares):.9f}/"
        f"{max(fertile_shares):.9f}"
    )
    print(
        "tail branch share min/mean/max: "
        f"{min(branch_shares):.9f}/"
        f"{sum(branch_shares)/len(branch_shares):.9f}/"
        f"{max(branch_shares):.9f}"
    )
    print(
        "tail layer-growth min/mean/max: "
        f"{min(growth):.9f}/{sum(growth)/len(growth):.9f}/{max(growth):.9f}"
    )
    for parity, name in ((0, "even"), (1, "odd")):
        selected = [r for r in tail_rows if int(r["depth"]) % 2 == parity]
        fmean = sum(float(r["fertile_share"]) for r in selected) / len(selected)
        bmean = sum(float(r["branch_share"]) for r in selected) / len(selected)
        print(f"tail {name}-depth means: fertile={fmean:.9f}, branch={bmean:.9f}")


if __name__ == "__main__":
    main()