src/utils.py

"""
High-level utilities for Collatz operations.

This module provides glue functions used by the web interface (app.py),
combining steps such as:

  - generating an inverse Collatz tree,
  - generating the minimal subtree up to N,
  - computing basic statistics,
  - rendering graphs to Graphviz PNG images.
"""

from __future__ import annotations

from typing import Optional, Tuple
from pathlib import Path

import pandas as pd

from src.collatz.inverse_tree import generate_collatz_tree
from src.collatz.minimal_subtree import generate_minimal_collatz_subtree_edges
from src.visual.render import render_collatz_tree_graphviz
from src.collatz.metrics import compute_basic_graph_stats


def build_and_render_collatz_tree(
    backbone_length: int,
    branch_length: int,
    max_depth: int,
    *,
    filename: str = "collatz_tree",
    output_dir: Optional[str] = "outputs",
    return_edges: bool = False,
) -> Tuple[str, Optional[pd.DataFrame]]:
    """
    Generate an inverse Collatz tree, render it as a PNG image, and optionally
    return the underlying edge DataFrame.
    """

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    df_edges: pd.DataFrame = generate_collatz_tree(
        backbone_length=backbone_length,
        branch_length=branch_length,
        max_depth=max_depth,
    )

    image_path = render_collatz_tree_graphviz(
        df_edges=df_edges,
        filename=filename,
        directory=output_path,
        image_format="png",
    )

    if return_edges:
        return image_path, df_edges

    return image_path, None


def build_and_render_minimal_subtree(
    N: int,
    *,
    filename: str = "collatz_minimal_subtree",
    output_dir: Optional[str] = "outputs",
    return_edges: bool = False,
) -> Tuple[str, Optional[pd.DataFrame]]:
    """
    Generate the minimal Collatz subtree containing all integers 1..N,
    render it as a PNG image, and optionally return the underlying
    edge DataFrame.

    This uses the structural branch construction (Algorithms 5 and 6)
    implemented in `generate_minimal_collatz_subtree_edges`.
    """

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    df_edges: pd.DataFrame = generate_minimal_collatz_subtree_edges(N)

    image_path = render_collatz_tree_graphviz(
        df_edges=df_edges,
        filename=filename,
        directory=output_path,
        image_format="png",
    )

    if return_edges:
        return image_path, df_edges

    return image_path, None


def compute_collatz_tree_stats(
    backbone_length: int,
    branch_length: int,
    max_depth: int,
) -> dict:
    """
    Convenience helper: generate the inverse Collatz tree and return
    basic graph stats (number of nodes, edges, parity counts, etc.).
    """
    df_edges = generate_collatz_tree(
        backbone_length=backbone_length,
        branch_length=branch_length,
        max_depth=max_depth,
    )
    return compute_basic_graph_stats(df_edges)


def safe_int(x, default: int = 1) -> int:
    """
    Convert a value to int safely.
    This protects the Gradio app from crashing when users type invalid input.
    """
    try:
        return int(x)
    except Exception:
        return default