src/collatz/inverse_tree.py

from __future__ import annotations

from typing import List, Tuple

import pandas as pd


def generate_collatz_tree(
    backbone_length: int = 11,
    branch_length: int = 5,
    max_depth: int = 3,
) -> pd.DataFrame:
    """
    Generate an inverse Collatz tree structure starting from 1, with controlled
    depth and branch expansion, based on the reverse Collatz rule:

        If (n - 1) % 3 == 0 and n is even, then (n - 1) / 3 is a valid preimage.

    The construction uses:
      - a "backbone" of powers of 2 starting at 1,
      - branches that grow from specific backbone nodes and from later branch nodes,
      - recursive expansion up to exactly `max_depth` levels.

    Parameters
    ----------
    backbone_length : int, default=11
        Length of the backbone (powers of 2).
        Backbone nodes are exactly: 1, 2, 4, ..., 2^(backbone_length - 1).
    branch_length : int, default=5
        Length of each branch created from valid reverse steps.
        Branch nodes are formed by repeated doubling from the branch root.
        Branches always have exactly `branch_length` nodes.
    max_depth : int, default=3
        Number of recursive branch-expansion levels.

    Returns
    -------
    pd.DataFrame
        DataFrame with two columns:
            - "Source": child node (preimage under the Collatz map),
            - "Target": parent node (image under the Collatz map).

        These edges define a directed tree (or forest) rooted at 1.
    """

    if backbone_length < 2:
        raise ValueError("backbone_length must be at least 2.")

    if branch_length < 1:
        raise ValueError("branch_length must be at least 1.")

    if max_depth < 0:
        raise ValueError("max_depth must be non-negative.")

    # Backbone nodes: 1, 2, 4, ..., 2^(backbone_length - 1)
    branches: List[List[int]] = [[2 ** i for i in range(backbone_length)]] 

    # Backbone edges: 2 -> 1, 4 -> 2, 8 -> 4, ...
    edges: List[Tuple[int, int]] = [(2 ** i, 2 ** (i - 1)) for i in range(1, backbone_length)] + [(1,4)]

    # Backbone nodes that will create branches: 2^4, 2^6, 2^8, ...
    branch_creating_numbers: List[int] = [
        2 ** i for i in range(4, backbone_length, 2)
    ]

    # Recursive branch expansion with EXACT branch_length and max_depth
    for _ in range(max_depth):
        next_level_branch_creators: List[int] = []

        for branch_num in branch_creating_numbers:
            # Reverse Collatz: (branch_num - 1) / 3
            parent = (branch_num - 1) // 3
            edges.append((parent, branch_num))

            # Create a branch of EXACT length = branch_length via repeated doubling
            new_branch = [parent * (2 ** i) for i in range(branch_length)]
            branches.append(new_branch)

            # Link within the branch (2a -> a, 4a -> 2a, ...)
            edges.extend(
                (new_branch[i + 1], new_branch[i])
                for i in range(len(new_branch) - 1)
            )

            # Candidates for further branching on next level
            next_level_branch_creators.extend(
                n for n in new_branch[1:] if (n - 1) % 3 == 0
            )

        branch_creating_numbers = next_level_branch_creators

    df_edges = pd.DataFrame(edges, columns=["Source", "Target"]).drop_duplicates()

    return df_edges