app.py

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, FancyBboxPatch
import io
from PIL import Image
from matplotlib.patches import FancyArrowPatch

class TreeNode:
    """Represents a node in the decision tree"""
    def __init__(self, depth=0, bounds=None):
        self.depth = depth
        self.bounds = bounds if bounds else {'x': (0, 10), 'y': (0, 10)}
        self.feature = None  # 'x' or 'y'
        self.threshold = None
        self.left = None
        self.right = None
        self.is_leaf = True
        self.samples = None
        self.class_counts = None
        self.entropy = None
        self.gini = None
        self.majority_class = None
        
class DecisionTreePartitioner:
    def __init__(self):
        self.reset_data()
        self.splits = []  # List of (feature, threshold) tuples
        self.root = None
        
    def reset_data(self):
        """Generate sample data with two classes"""
        np.random.seed(42)
        # Class 0 (blue) - bottom left
        n_samples = 50
        self.X0 = np.random.randn(n_samples, 2) * 1.5 + np.array([3, 3])
        # Class 1 (red) - top right  
        self.X1 = np.random.randn(n_samples, 2) * 1.5 + np.array([7, 7])
        
        self.X = np.vstack([self.X0, self.X1])
        self.y = np.hstack([np.zeros(n_samples), np.ones(n_samples)])
        self.splits = []
        self.root = None
        
    def calculate_entropy(self, y):
        """Calculate entropy for a set of labels"""
        if len(y) == 0:
            return 0
        _, counts = np.unique(y, return_counts=True)
        probabilities = counts / len(y)
        entropy = -np.sum(probabilities * np.log2(probabilities + 1e-10))
        return entropy
    
    def calculate_gini(self, y):
        """Calculate Gini index for a set of labels"""
        if len(y) == 0:
            return 0
        _, counts = np.unique(y, return_counts=True)
        probabilities = counts / len(y)
        gini = 1 - np.sum(probabilities ** 2)
        return gini
    
    def build_tree_from_splits(self):
        """Build tree structure from the list of splits"""
        if not self.splits:
            return None
            
        self.root = TreeNode(depth=0)
        self._build_node(self.root, np.arange(len(self.y)), 0)
        return self.root
    
    def _build_node(self, node, indices, split_idx):
        """Recursively build tree nodes"""
        if len(indices) == 0:
            return
            
        # Calculate node statistics
        node.samples = len(indices)
        y_node = self.y[indices]
        unique, counts = np.unique(y_node, return_counts=True)
        node.class_counts = dict(zip(unique.astype(int), counts))
        node.entropy = self.calculate_entropy(y_node)
        node.gini = self.calculate_gini(y_node)
        node.majority_class = int(unique[np.argmax(counts)])
        
        # Check if we have more splits to apply
        if split_idx >= len(self.splits):
            node.is_leaf = True
            return
            
        # Apply the split
        feature, threshold = self.splits[split_idx]
        feature_idx = 0 if feature == 'x' else 1
        
        X_node = self.X[indices]
        left_mask = X_node[:, feature_idx] <= threshold
        right_mask = ~left_mask
        
        left_indices = indices[left_mask]
        right_indices = indices[right_mask]
        
        # Only create split if both children are non-empty
        if len(left_indices) > 0 and len(right_indices) > 0:
            node.is_leaf = False
            node.feature = feature
            node.threshold = threshold
            
            # Create child nodes with updated bounds
            left_bounds = node.bounds.copy()
            right_bounds = node.bounds.copy()
            
            if feature == 'x':
                left_bounds['x'] = (node.bounds['x'][0], threshold)
                right_bounds['x'] = (threshold, node.bounds['x'][1])
            else:
                left_bounds['y'] = (node.bounds['y'][0], threshold)
                right_bounds['y'] = (threshold, node.bounds['y'][1])
            
            node.left = TreeNode(depth=node.depth + 1, bounds=left_bounds)
            node.right = TreeNode(depth=node.depth + 1, bounds=right_bounds)
            
            # Recursively build children
            self._build_node(node.left, left_indices, split_idx + 1)
            self._build_node(node.right, right_indices, split_idx + 1)
    
    def add_split(self, feature, threshold):
        """Add a new split to the tree"""
        self.splits.append((feature, threshold))
        self.build_tree_from_splits()
        
    def remove_last_split(self):
        """Remove the last split"""
        if self.splits:
            self.splits.pop()
            if self.splits:
                self.build_tree_from_splits()
            else:
                self.root = None
    
    def draw_tree(self, node=None, ax=None, x=0.5, y=1.0, dx=0.25, level=0):
        """Recursively draw the decision tree"""
        if node is None:
            return
            
        # Node styling
        if node.is_leaf:
            box_color = 'lightblue' if node.majority_class == 0 else 'orange'
            alpha = 0.7
        else:
            box_color = 'lightgreen'
            alpha = 0.5
        
        # Create node text
        if node.is_leaf:
            text = f"Leaf\nClass: {node.majority_class}\n"
            text += f"Samples: {node.samples}\n"
            text += f"Entropy: {node.entropy:.3f}\n"
            text += f"Gini: {node.gini:.3f}"
        else:
            feature_name = "Width" if node.feature == 'x' else "Height"
            text = f"{feature_name} ≤ {node.threshold:.2f}\n"
            text += f"Samples: {node.samples}\n"
            text += f"Entropy: {node.entropy:.3f}\n"
            text += f"Gini: {node.gini:.3f}"
        
        # Draw box
        bbox = dict(boxstyle="round,pad=0.3", facecolor=box_color, 
                   edgecolor='black', linewidth=2, alpha=alpha)
        ax.text(x, y, text, ha='center', va='center', fontsize=8,
               bbox=bbox, fontweight='bold')
        
        # Draw connections to children
        if not node.is_leaf and node.left and node.right:
            # Left child
            y_child = y - 0.15
            x_left = x - dx
            x_right = x + dx
            
            # Draw arrows
            arrow_left = FancyArrowPatch((x, y - 0.05), (x_left, y_child + 0.05),
                                        arrowstyle='->', mutation_scale=20, 
                                        linewidth=2, color='blue')
            arrow_right = FancyArrowPatch((x, y - 0.05), (x_right, y_child + 0.05),
                                         arrowstyle='->', mutation_scale=20,
                                         linewidth=2, color='red')
            ax.add_patch(arrow_left)
            ax.add_patch(arrow_right)
            
            # Add Yes/No labels
            ax.text((x + x_left) / 2, (y + y_child) / 2, 'Yes', 
                   fontsize=9, color='blue', fontweight='bold')
            ax.text((x + x_right) / 2, (y + y_child) / 2, 'No',
                   fontsize=9, color='red', fontweight='bold')
            
            # Recursively draw children
            self.draw_tree(node.left, ax, x_left, y_child, dx * 0.5, level + 1)
            self.draw_tree(node.right, ax, x_right, y_child, dx * 0.5, level + 1)
    
    def visualize(self, split_history):
        """Create comprehensive visualization"""
        fig = plt.figure(figsize=(20, 10))
        gs = fig.add_gridspec(2, 2, height_ratios=[1, 1], width_ratios=[1.2, 1])
        
        ax1 = fig.add_subplot(gs[0, 0])  # Partition view
        ax2 = fig.add_subplot(gs[1, 0])  # Decision tree
        ax3 = fig.add_subplot(gs[:, 1])  # Statistics
        
        # Parse split history
        self.splits = []
        if split_history.strip():
            for line in split_history.strip().split('\n'):
                if ',' in line:
                    parts = line.split(',')
                    if len(parts) == 2:
                        feature = parts[0].strip().lower()
                        try:
                            threshold = float(parts[1].strip())
                            self.splits.append((feature, threshold))
                        except ValueError:
                            pass
        
        # Build tree from splits
        if self.splits:
            self.build_tree_from_splits()
        
        # === Plot 1: Partitioned Feature Space ===
        ax1.scatter(self.X[self.y == 0, 0], self.X[self.y == 0, 1], 
                   c='blue', label='Class 0 (Lemon)', s=100, alpha=0.6, edgecolors='k')
        ax1.scatter(self.X[self.y == 1, 0], self.X[self.y == 1, 1], 
                   c='orange', label='Class 1 (Orange)', s=100, alpha=0.6, edgecolors='k')
        
        # Draw all partition lines
        colors = plt.cm.rainbow(np.linspace(0, 1, len(self.splits)))
        for idx, (feature, threshold) in enumerate(self.splits):
            if feature == 'x':
                ax1.axvline(x=threshold, color=colors[idx], linewidth=2.5, 
                           linestyle='--', label=f'Split {idx+1}: x≤{threshold:.1f}', alpha=0.8)
            else:
                ax1.axhline(y=threshold, color=colors[idx], linewidth=2.5,
                           linestyle='--', label=f'Split {idx+1}: y≤{threshold:.1f}', alpha=0.8)
        
        ax1.set_xlabel('Feature 1 (Width)', fontsize=14, fontweight='bold')
        ax1.set_ylabel('Feature 2 (Height)', fontsize=14, fontweight='bold')
        ax1.set_title('Partitioned Feature Space', fontsize=16, fontweight='bold')
        ax1.legend(fontsize=10, loc='upper left')
        ax1.grid(True, alpha=0.3)
        ax1.set_xlim(0, 10)
        ax1.set_ylim(0, 10)
        
        # === Plot 2: Decision Tree ===
        ax2.clear()
        ax2.set_xlim(0, 1)
        ax2.set_ylim(0, 1)
        ax2.axis('off')
        ax2.set_title('Decision Tree Structure', fontsize=16, fontweight='bold', pad=20)
        
        if self.root:
            self.draw_tree(self.root, ax2)
        else:
            ax2.text(0.5, 0.5, 'No splits yet\nAdd splits to build the tree', 
                    ha='center', va='center', fontsize=14,
                    bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        # === Plot 3: Statistics ===
        ax3.clear()
        ax3.axis('off')
        
        # Calculate overall statistics
        entropy_initial = self.calculate_entropy(self.y)
        gini_initial = self.calculate_gini(self.y)
        
        stats_text = "DECISION TREE STATISTICS\n" + "="*50 + "\n\n"
        stats_text += f"Total Samples: {len(self.y)}\n"
        stats_text += f"  • Class 0: {np.sum(self.y == 0)}\n"
        stats_text += f"  • Class 1: {np.sum(self.y == 1)}\n\n"
        stats_text += f"Initial Impurity:\n"
        stats_text += f"  • Entropy: {entropy_initial:.4f}\n"
        stats_text += f"  • Gini: {gini_initial:.4f}\n\n"
        
        if self.splits:
            stats_text += f"Number of Splits: {len(self.splits)}\n\n"
            stats_text += "SPLIT SEQUENCE:\n" + "-"*50 + "\n"
            
            for idx, (feature, threshold) in enumerate(self.splits):
                feature_name = "Width (x)" if feature == 'x' else "Height (y)"
                stats_text += f"\n{idx+1}. {feature_name} ≤ {threshold:.2f}\n"
            
            # Get leaf statistics
            leaves = []
            self._collect_leaves(self.root, leaves)
            
            if leaves:
                stats_text += f"\n\nLEAF NODES: {len(leaves)}\n" + "-"*50 + "\n"
                for idx, leaf in enumerate(leaves):
                    stats_text += f"\nLeaf {idx+1}:\n"
                    stats_text += f"  • Samples: {leaf.samples}\n"
                    stats_text += f"  • Class 0: {leaf.class_counts.get(0, 0)} | "
                    stats_text += f"Class 1: {leaf.class_counts.get(1, 0)}\n"
                    stats_text += f"  • Prediction: Class {leaf.majority_class}\n"
                    stats_text += f"  • Entropy: {leaf.entropy:.4f}\n"
                    stats_text += f"  • Gini: {leaf.gini:.4f}\n"
                
                # Calculate weighted average impurity
                total_samples = sum(leaf.samples for leaf in leaves)
                avg_entropy = sum(leaf.entropy * leaf.samples for leaf in leaves) / total_samples
                avg_gini = sum(leaf.gini * leaf.samples for leaf in leaves) / total_samples
                
                stats_text += f"\n\nWEIGHTED AVERAGE IMPURITY:\n" + "-"*50 + "\n"
                stats_text += f"  • Entropy: {avg_entropy:.4f}\n"
                stats_text += f"  • Gini: {avg_gini:.4f}\n"
                stats_text += f"\nTOTAL INFORMATION GAIN:\n"
                stats_text += f"  • {entropy_initial - avg_entropy:.4f}\n"
                stats_text += f"\nTOTAL GINI REDUCTION:\n"
                stats_text += f"  • {gini_initial - avg_gini:.4f}\n"
        else:
            stats_text += "No splits applied yet.\n"
            stats_text += "\nAdd splits in the format:\n"
            stats_text += "  feature, threshold\n\n"
            stats_text += "Example:\n"
            stats_text += "  x, 5.0\n"
            stats_text += "  y, 6.5\n"
        
        ax3.text(0.05, 0.95, stats_text, transform=ax3.transAxes,
                fontsize=10, verticalalignment='top',
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5),
                family='monospace')
        
        plt.tight_layout()
        
        # Convert to image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=120, bbox_inches='tight')
        buf.seek(0)
        img = Image.open(buf)
        plt.close()
        
        return img
    
    def _collect_leaves(self, node, leaves):
        """Collect all leaf nodes"""
        if node is None:
            return
        if node.is_leaf:
            leaves.append(node)
        else:
            self._collect_leaves(node.left, leaves)
            self._collect_leaves(node.right, leaves)

# Create the partitioner
partitioner = DecisionTreePartitioner()

# Create Gradio interface
with gr.Blocks(title="Multi-Split Decision Tree Visualizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🌳 Interactive Multi-Split Decision Tree Visualizer
    
    Build a decision tree step-by-step and visualize the partitioning process!
    
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            split_input = gr.Textbox(
                label="📝 Split Sequence (one per line: feature, threshold)",
                placeholder="x, 5.0\ny, 6.5\nx, 3.0",
                lines=10,
                value="x, 5.0"
            )
            
            update_btn = gr.Button("🔄 Update Visualization", variant="primary", size="lg")
            
            gr.Markdown("""
            ### Example Splits:
            **Simple 2-split tree:**
            ```
            x, 5.0
            y, 6.5
            ```
            
            **Complex 4-split tree:**
            ```
            x, 5.0
            y, 6.5
            x, 3.0
            y, 8.0
            ```
            """)
            
        with gr.Column(scale=2):
            output_image = gr.Image(label="Visualization", height=800)
    
    # Update visualization
    update_btn.click(
        fn=partitioner.visualize,
        inputs=[split_input],
        outputs=output_image
    )
    
    # Initial visualization
    demo.load(
        fn=partitioner.visualize,
        inputs=[split_input],
        outputs=output_image
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()