tree_visualizer.py

"""
Parse Tree Visualizer
Draws the AST as a proper tree using matplotlib with a Reingold-Tilford layout.
"""

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch
import matplotlib.patheffects as pe
import numpy as np
from typing import Dict, List, Optional, Tuple
from ast_nodes import ASTNode


# ── Color palette per node category ──────────────────────────────────────────

NODE_COLORS = {
    'Program':     '#1a1a2e',
    'Block':       '#16213e',
    'FuncDef':     '#0f3460',
    'FunctionCall':'#533483',
    'If':          '#e94560',
    'While':       '#c84b31',
    'For':         '#b25068',
    'Return':      '#774898',
    'Print':       '#4f8a8b',
    'VarDecl':     '#2b6cb0',
    'Assign':      '#2c7a7b',
    'BinOp':       '#d4a017',
    'UnaryOp':     '#c05621',
    'Number':      '#276749',
    'String':      '#285e61',
    'Boolean':     '#553c9a',
    'ID':          '#2c5282',
    'default':     '#4a5568',
}


def get_node_color(label: str) -> str:
    for key, color in NODE_COLORS.items():
        if key.lower() in label.lower():
            return color
    return NODE_COLORS['default']


# ── Reingold-Tilford layout ───────────────────────────────────────────────────

class LayoutNode:
    def __init__(self, node: ASTNode, depth: int = 0):
        self.node = node
        self.depth = depth
        self.x: float = 0.0
        self.y: float = -depth
        self.children: List['LayoutNode'] = [
            LayoutNode(c, depth + 1) for c in node.children()
        ]
        self.mod: float = 0.0
        self.thread: Optional['LayoutNode'] = None
        self.ancestor: Optional['LayoutNode'] = None
        self.change: float = 0.0
        self.shift: float = 0.0
        self.number: int = 0
        self.prelim: float = 0.0


def _left_contour(node: LayoutNode, v: float, contour: Dict[int, float]):
    if node.depth not in contour:
        contour[node.depth] = node.x + v
    else:
        contour[node.depth] = min(contour[node.depth], node.x + v)
    for c in node.children:
        _left_contour(c, v + node.mod, contour)


def _right_contour(node: LayoutNode, v: float, contour: Dict[int, float]):
    if node.depth not in contour:
        contour[node.depth] = node.x + v
    else:
        contour[node.depth] = max(contour[node.depth], node.x + v)
    for c in node.children:
        _right_contour(c, v + node.mod, contour)


def _buchheim(node: LayoutNode, h_gap: float = 1.8):
    """Simplified Buchheim (Reingold-Tilford) layout."""
    _first_walk(node, h_gap)
    _second_walk(node, -node.prelim)


def _first_walk(v: LayoutNode, h_gap: float):
    if not v.children:
        if v.number > 0 and v.parent_sibling:
            v.prelim = v.parent_sibling.prelim + h_gap
        else:
            v.prelim = 0.0
        return

    for i, child in enumerate(v.children):
        child.number = i
        child.parent_sibling = v.children[i - 1] if i > 0 else None
        _first_walk(child, h_gap)

    left_child  = v.children[0]
    right_child = v.children[-1]
    midpoint = (left_child.prelim + right_child.prelim) / 2

    if v.number > 0 and v.parent_sibling:
        v.prelim = v.parent_sibling.prelim + h_gap
        v.mod = v.prelim - midpoint
    else:
        v.prelim = midpoint


# Simplified version without full ancestor tracking for robustness
def simple_layout(node: LayoutNode, depth: int = 0,
                  counter: List[float] = None, h_gap: float = 2.0):
    """Assign x positions using in-order traversal (simple, clean)."""
    if counter is None:
        counter = [0.0]

    for child in node.children:
        simple_layout(child, depth + 1, counter, h_gap)

    if not node.children:
        node.x = counter[0]
        counter[0] += h_gap
    else:
        node.x = (node.children[0].x + node.children[-1].x) / 2

    node.y = -depth * 1.8


def collect_nodes(layout: LayoutNode) -> List[LayoutNode]:
    result = [layout]
    for c in layout.children:
        result.extend(collect_nodes(c))
    return result


# ── Drawing ───────────────────────────────────────────────────────────────────

def draw_tree(ax: plt.Axes, root: ASTNode, title: str = "Parse Tree"):
    ax.clear()
    ax.set_facecolor('#0d1117')

    if root is None:
        ax.text(0.5, 0.5, "No tree to display", ha='center', va='center',
                color='white', fontsize=14, transform=ax.transAxes)
        return

    layout_root = LayoutNode(root)
    simple_layout(layout_root)

    all_nodes = collect_nodes(layout_root)

    # Scale to fit nicely
    xs = [n.x for n in all_nodes]
    ys = [n.y for n in all_nodes]
    if len(xs) == 1:
        xs = [0, 1]
    xmin, xmax = min(xs), max(xs)
    ymin, ymax = min(ys), max(ys)
    margin_x = max(1.5, (xmax - xmin) * 0.08)
    margin_y = 1.2

    ax.set_xlim(xmin - margin_x, xmax + margin_x)
    ax.set_ylim(ymin - margin_y, ymax + margin_y)
    ax.axis('off')

    # Draw edges first
    for n in all_nodes:
        for child in n.children:
            ax.plot([n.x, child.x], [n.y, child.y],
                    color='#4a90d9', linewidth=1.5, alpha=0.7, zorder=1)

    # Draw nodes
    box_w = 0.75
    box_h = 0.55
    fs = max(5, min(9, 120 / max(len(all_nodes), 1)))

    for n in all_nodes:
        label = n.node.node_label()
        color = get_node_color(label)

        # Glow effect
        fancy = FancyBboxPatch(
            (n.x - box_w / 2, n.y - box_h / 2), box_w, box_h,
            boxstyle="round,pad=0.05",
            linewidth=2,
            edgecolor='#4a90d9',
            facecolor=color,
            zorder=2,
            alpha=0.95,
        )
        ax.add_patch(fancy)

        ax.text(n.x, n.y, label,
                ha='center', va='center',
                fontsize=fs,
                color='white',
                fontfamily='monospace',
                fontweight='bold',
                zorder=3,
                wrap=True)

    ax.set_title(title, color='white', fontsize=13, fontweight='bold', pad=10)