modules/infographics_generator/layout_system/sdf/optimizer.py

"""Main optimization function for SDF-based layout optimization."""

import os
import shutil
from typing import List, Tuple, Optional, Dict

import numpy as np
import torch
import torch.nn.functional as F

from modules.infographics_generator.layout_system import parameters as params

from .core import (
    load_binary_mask_from_rgba,
    tight_bbox_ratio,
    binary_to_sdf_norm,
    dilate_mask,
    make_container_grid,
    sdf_to_softmask,
    area_sum,
)
from .bbox import (
    bbox_aspect_from_unconstrained,
    unconstrained_from_bbox,
)
from .losses import (
    compute_alignment_consistency_loss,
    compute_alignment_similarity_loss,
    compute_readability_loss,
    compute_proximity_ratio_loss,
    compute_visual_balance_loss,
    compute_position_size_similarity_loss,
)
from .visualization import (
    visualize_optimization_progress,
    visualize_final_result,
    save_composite_image,
)


def optimize(
    png_list,  # List of image paths for N nodes
    original_png_list=None,  # List of original image paths for saving composite
    Wc=1000, Hc=1000,
    opt_res_list=params.OPT_RES_LIST,  # optimize on these resolutions
    outer_rounds=params.OUTER_ROUNDS,  # augmented-lagrangian outer updates per stage
    inner_steps=params.INNER_STEPS,    # gradient steps per outer round
    tau_schedule=params.TAU_SCHEDULE,
    rho_init=params.RHO_INIT,          # initial penalty parameter
    rho_mult=params.RHO_MULT,          # penalty multiplier
    lr=params.LEARNING_RATE,
    size_min=params.SIZE_MIN,         # Legacy parameter for backward compatibility
    min_sizes=None,                    # List of (min_width, min_height) tuples for each element
    pen_weight=params.PEN_WEIGHT,      # weight for penetration penalty
    pen_eta_px=params.PEN_ETA_PX,
    reference_bboxes=None,             # List of reference element bboxes (x, y, w, h) from Example layout
    reference_parent_bbox=None,        # Reference parent container bbox (x, y, w, h)
    w_similarity=params.W_SIMILARITY,  # weight for position/size similarity loss
    size_rules=None,                   # List of tuples (source_idx, target_idx) for size hierarchy rules
    w_readability=params.W_READABILITY, # weight for readability loss
    w_alignment_consistency=params.W_ALIGNMENT_CONSISTENCY,  # weight for alignment consistency loss
    alignment_constraint=None,          # Dictionary with alignment constraint from JSON (direction, value)
    w_alignment_similarity=params.W_ALIGNMENT_SIMILARITY,  # weight for alignment similarity loss
    proximity_info=None,               # Dictionary with container hierarchy info for proximity ratio calculation
    w_proximity=params.W_PROXIMITY,    # weight for proximity ratio loss
    w_data_ink=params.W_DATA_INK,      # weight for data ink loss (maximize union area)
    w_visual_balance=params.W_VISUAL_BALANCE,  # weight for visual balance loss
    min_gap_px=20.0,                    # minimum gap between elements in pixels
    device=None,
    save_prefix=None,                  # Prefix for saving result images (None = use default names)
    debug=False,                        # Enable debug mode: visualization and detailed logging
):
    """Main SDF-based layout optimization function.
    
    Args:
        png_list: List of image paths for N nodes
        original_png_list: List of original image paths for saving composite
        Wc, Hc: Container width and height
        opt_res_list: Optimize on these resolutions
        outer_rounds: Augmented-lagrangian outer updates per stage
        inner_steps: Gradient steps per outer round
        tau_schedule: Schedule for tau parameter
        rho_init: Initial penalty parameter
        rho_mult: Penalty multiplier
        lr: Learning rate
        size_min: Legacy minimum size parameter
        min_sizes: List of (min_width, min_height) tuples for each element
        pen_weight: Weight for penetration penalty
        pen_eta_px: Eta parameter for penetration penalty
        reference_bboxes: List of reference element bboxes from Example layout
        reference_parent_bbox: Reference parent container bbox
        w_similarity: Weight for position/size similarity loss
        size_rules: List of tuples (source_idx, target_idx) for size hierarchy rules
        w_readability: Weight for readability loss
        w_alignment_consistency: Weight for alignment consistency loss
        alignment_constraint: Dictionary with alignment constraint from JSON
        w_alignment_similarity: Weight for alignment similarity loss
        proximity_info: Dictionary with container hierarchy info
        w_proximity: Weight for proximity ratio loss
        w_data_ink: Weight for data ink loss
        w_visual_balance: Weight for visual balance loss
        min_gap_px: Minimum gap between elements in pixels
        device: PyTorch device
        save_prefix: Prefix for saving result images
        debug: Enable debug mode (visualization and detailed logging)
    
    Returns:
        List of final bounding boxes [(x, y, w, h), ...]
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    print("Device:", device)

    # Validate png_list
    if png_list is None or len(png_list) < 1:
        raise ValueError("png_list must be provided with at least one image path")
    
    num_nodes = len(png_list)
    
    # Handle original_png_list
    if original_png_list is None:
        original_png_list = png_list
    
    # Set dilation_radii to 100 for all elements
    dilation_radii = [20.0] * num_nodes
    
    # print(f"Number of nodes: {num_nodes}")
    # print(f"Dilation radii: {dilation_radii}")

    # Validate reference bboxes
    if reference_bboxes is None:
        reference_bboxes = []
    if reference_parent_bbox is None:
        reference_parent_bbox = (0.0, 0.0, float(Wc), float(Hc))
    
    # Normalize reference bboxes: subtract x_min and y_min to make coordinates start from 0
    if reference_bboxes and len(reference_bboxes) > 0:
        # Find minimum x and y across all reference bboxes
        x_min = min(bbox[0] for bbox in reference_bboxes)
        y_min = min(bbox[1] for bbox in reference_bboxes)
        
        # Subtract x_min and y_min from all bboxes
        normalized_reference_bboxes = []
        for bbox in reference_bboxes:
            if isinstance(bbox, (tuple, list)) and len(bbox) >= 4:
                normalized_bbox = (bbox[0] - x_min, bbox[1] - y_min, bbox[2], bbox[3])
            elif isinstance(bbox, dict):
                normalized_bbox = {
                    'x': bbox.get('x', 0) - x_min,
                    'y': bbox.get('y', 0) - y_min,
                    'width': bbox.get('width', bbox.get('w', 100)),
                    'height': bbox.get('height', bbox.get('h', 100))
                }
            else:
                normalized_bbox = bbox
            normalized_reference_bboxes.append(normalized_bbox)
        
        reference_bboxes = normalized_reference_bboxes
        
        # Also normalize reference_parent_bbox
        if reference_parent_bbox:
            x_p, y_p, w_p, h_p = reference_parent_bbox
            reference_parent_bbox = (x_p - x_min, y_p - y_min, w_p, h_p)
        
        # print(f"Normalized reference bboxes: subtracted x_min={x_min:.1f}, y_min={y_min:.1f}")
    
    # Validate size rules
    if size_rules is None:
        size_rules = []
    
    # Validate and set up min_sizes
    if min_sizes is None:
        # Use default values for all elements
        min_sizes = [(params.MIN_WIDTH_DEFAULT, params.MIN_HEIGHT_DEFAULT)] * num_nodes
    elif len(min_sizes) < num_nodes:
        # Extend with default values
        default_size = (params.MIN_WIDTH_DEFAULT, params.MIN_HEIGHT_DEFAULT)
        min_sizes = min_sizes + [default_size] * (num_nodes - len(min_sizes))
    elif len(min_sizes) > num_nodes:
        # Truncate to num_nodes
        min_sizes = min_sizes[:num_nodes]
    
    # print(f"Reference bboxes: {len(reference_bboxes)} elements")
    # print(f"Reference parent bbox: {reference_parent_bbox}")
    # print(f"Size rules: {len(size_rules)} rules")
    # print(f"Similarity weight: {w_similarity}, Readability weight: {w_readability}")
    # print(f"Min sizes: {min_sizes}")
    # print(f"Debug mode: {debug}")

    # Create visualization folder for this optimization run (only in debug mode)
    viz_folder = None
    if debug:
        if save_prefix:
            viz_folder = f"{save_prefix}_visualization"
        else:
            viz_folder = "optimization_visualization"
        
        # Clear existing folder contents if it exists
        if os.path.exists(viz_folder):
            shutil.rmtree(viz_folder)
        
        os.makedirs(viz_folder, exist_ok=True)
        # print(f"Visualization folder created: {viz_folder}")

    # load masks for all nodes
    masks = []
    ratios = []
    for i, png_path in enumerate(png_list):
        if png_path is None:
            raise ValueError(f"Image path at index {i} is None")
        mask = load_binary_mask_from_rgba(png_path)
        
        # Apply dilation with radius 100 to all elements
        mask = dilate_mask(mask, dilation_radii[i])
        # print(f"Applied dilation to node {i} mask with radius {dilation_radii[i]:.1f} pixels")
        
        masks.append(mask)
        r, _ = tight_bbox_ratio(mask)
        ratios.append(r)
        # print(f"Node {i} aspect ratio (tight alpha bbox): r={r:.4f}")

    # precompute SDF templates for all nodes
    sdf_norms = []
    sdf_tensors = []
    for i, mask in enumerate(masks):
        sdf_norm = binary_to_sdf_norm(mask, pad=16)
        sdf_norms.append(sdf_norm)
        sdf_t = torch.from_numpy(sdf_norm)[None, None].to(device)
        sdf_tensors.append(sdf_t)

    # learnable parameters: (tx,ty,ts) for each object
    # Initialize from reference_bboxes if available, otherwise use zeros
    opt_params = []
    for i in range(num_nodes):
        if reference_bboxes and i < len(reference_bboxes):
            # Initialize from reference bbox
            ref_bbox = reference_bboxes[i]
            if isinstance(ref_bbox, (tuple, list)) and len(ref_bbox) >= 4:
                ref_x, ref_y, ref_w, ref_h = ref_bbox[0], ref_bbox[1], ref_bbox[2], ref_bbox[3]
            elif isinstance(ref_bbox, dict):
                ref_x = ref_bbox.get("x", 0)
                ref_y = ref_bbox.get("y", 0)
                ref_w = ref_bbox.get("width", ref_bbox.get("w", 100))
                ref_h = ref_bbox.get("height", ref_bbox.get("h", 100))
            else:
                ref_x, ref_y, ref_w, ref_h = 0, 0, 100, 100
            
            # Convert reference bbox to unconstrained parameters
            # Use material's aspect ratio (ratios[i]) instead of JSON's w/h
            min_width, min_height = min_sizes[i]
            tx_init, ty_init, ts_init = unconstrained_from_bbox(
                ref_x, ref_y, ref_w, ref_h,
                Wc, Hc, ratios[i],
                min_width=min_width, min_height=min_height,
                size_min=size_min
            )
            # print(f"Node {i}: Initializing from reference bbox ({ref_x:.1f}, {ref_y:.1f}, {ref_w:.1f}, {ref_h:.1f}), "
                #   f"adjusted to aspect ratio {ratios[i]:.4f}")
        else:
            # Initialize with zeros (default)
            tx_init, ty_init, ts_init = 0.0, 0.0, 0.0
        
        tx = torch.nn.Parameter(torch.tensor(tx_init, device=device))
        ty = torch.nn.Parameter(torch.tensor(ty_init, device=device))
        ts = torch.nn.Parameter(torch.tensor(ts_init, device=device))
        opt_params.extend([tx, ty, ts])
    
    opt = torch.optim.Adam(opt_params, lr=lr)

    # augmented lagrangian multipliers for constraint g = A_inter = 0
    lam = torch.tensor(0.0, device=device)
    rho = torch.tensor(rho_init, device=device)

    tau_list = list(tau_schedule)
    if len(tau_list) < outer_rounds:
        tau_list += [tau_list[-1]] * (outer_rounds - len(tau_list))

    # staged optimization over resolutions
    for stage_idx, stage_res in enumerate(opt_res_list):
        Hs = Ws = int(stage_res)
        X, Y = make_container_grid(Hs, Ws, device=device)
        # print(f"\n=== Stage optimize at {Ws}x{Hs} (container {Wc}x{Hc}) ===")

        # Visualize initial state before optimization (only for first stage and only in debug mode)
        if stage_idx == 0 and debug:
            with torch.no_grad():
                # Use full resolution for visualization
                X_init, Y_init = make_container_grid(Hc, Wc, device=device)
                initial_bboxes = []
                initial_softmasks = []
                for i in range(num_nodes):
                    tx_idx = i * 3
                    ty_idx = i * 3 + 1
                    ts_idx = i * 3 + 2
                    tx = opt_params[tx_idx]
                    ty = opt_params[ty_idx]
                    ts = opt_params[ts_idx]
                    
                    min_width, min_height = min_sizes[i]
                    x, y, w, h = bbox_aspect_from_unconstrained(
                        tx, ty, ts, Wc, Hc, ratios[i],
                        min_width=min_width, min_height=min_height,
                        size_min=size_min
                    )
                    # print(f"Node {i} initial bbox: ({x.item():.1f}, {y.item():.1f}, {w.item():.1f}, {h.item():.1f})")
                    initial_bboxes.append((x.item(), y.item(), w.item(), h.item()))
                    
                    # Use smaller tau for initial visualization to show actual mask shape
                    m, _ = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X_init, Y_init, tau_px=0.5)
                    initial_softmasks.append(m)
                
                # Compute initial loss values for display
                initial_union = torch.ones_like(initial_softmasks[0])
                for m in initial_softmasks:
                    initial_union = initial_union * (1.0 - m)
                initial_union = 1.0 - initial_union
                initial_A_union = initial_union.sum().item()  # Full resolution, da=1
                
                initial_inter = torch.zeros_like(initial_softmasks[0])
                for i in range(num_nodes):
                    for j in range(i + 1, num_nodes):
                        initial_inter = initial_inter + initial_softmasks[i] * initial_softmasks[j]
                initial_A_inter = initial_inter.sum().item()  # Full resolution, da=1
                
                # Compute visual balance loss for initial state (using bboxes)
                initial_bboxes_tensors = []
                for i in range(num_nodes):
                    tx_idx = i * 3
                    ty_idx = i * 3 + 1
                    ts_idx = i * 3 + 2
                    tx = opt_params[tx_idx]
                    ty = opt_params[ty_idx]
                    ts = opt_params[ts_idx]
                    
                    min_width, min_height = min_sizes[i]
                    x, y, w, h = bbox_aspect_from_unconstrained(
                        tx, ty, ts, Wc, Hc, ratios[i],
                        min_width=min_width, min_height=min_height,
                        size_min=size_min
                    )
                    initial_bboxes_tensors.append((x, y, w, h))
                
                L_visual_balance_init = compute_visual_balance_loss(
                    initial_bboxes_tensors, Wc, Hc, device=device
                )
                
                initial_loss_info = {
                    'A_union': initial_A_union,
                    'A_inter': initial_A_inter,
                    'visual_balance': w_visual_balance * L_visual_balance_init.item(),
                }
                
                initial_save_path = os.path.join(viz_folder, f"initial_stage{stage_idx}.png")
                
                visualize_optimization_progress(
                    initial_softmasks, initial_bboxes, Wc, Hc,
                    epoch=-1, loss_info=initial_loss_info,
                    save_path=initial_save_path
                )

        for k in range(outer_rounds):
            tau_px = float(tau_list[min(k, len(tau_list)-1)])

            for t in range(inner_steps):
                opt.zero_grad(set_to_none=True)

                # Compute bboxes for all nodes
                bboxes = []
                softmasks = []
                distances = []
                for i in range(num_nodes):
                    tx_idx = i * 3
                    ty_idx = i * 3 + 1
                    ts_idx = i * 3 + 2
                    tx = opt_params[tx_idx]
                    ty = opt_params[ty_idx]
                    ts = opt_params[ts_idx]
                    
                    min_width, min_height = min_sizes[i]
                    x, y, w, h = bbox_aspect_from_unconstrained(
                        tx, ty, ts, Wc, Hc, ratios[i],
                        min_width=min_width, min_height=min_height,
                        size_min=size_min
                    )
                    bboxes.append((x, y, w, h))
                    
                    m, d_px = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X, Y, tau_px=tau_px)
                    softmasks.append(m)
                    distances.append(d_px)

                # Compute union: 1 - product of (1 - mask_i)
                union = torch.ones_like(softmasks[0])
                for m in softmasks:
                    union = union * (1.0 - m)
                union = 1.0 - union

                # Compute intersection: sum of all pairwise intersections
                inter = torch.zeros_like(softmasks[0])
                for i in range(num_nodes):
                    for j in range(i + 1, num_nodes):
                        inter = inter + softmasks[i] * softmasks[j]

                A_union = area_sum(union, Wc, Hc)
                A_inter = area_sum(inter, Wc, Hc)  # must go to 0

                # Visual balance loss (using bboxes directly for differentiability)
                L_visual_balance = compute_visual_balance_loss(
                    bboxes, Wc, Hc, device=device
                )

                # Penetration loss: based on actual bbox gap (not limited by SDF range)
                # Penalize when gap < min_gap_px (ensures minimum gap between elements)
                # Use ReLU for hard cutoff: no penalty when gap >= min_gap_px
                # For "layer" container type, skip this loss (overlapping is allowed)
                L_pen = torch.tensor(0.0, device=device)
                
                # # Check if this is a layer container (overlapping allowed)
                # is_layer = False
                # if proximity_info:
                #     container_types = proximity_info.get("types", [])
                #     if len(container_types) > 0:
                #         is_layer = (container_types[0] == "layer")
                
                # if not is_layer:
                # # Only apply penetration penalty for row/column layouts
                #     for i in range(num_nodes):
                #         for j in range(i + 1, num_nodes):
                #             x_i, y_i, w_i, h_i = bboxes[i]
                #             x_j, y_j, w_j, h_j = bboxes[j]
                            
                #             # Compute gap in each dimension (negative if overlapping)
                #             gap_x = torch.max(x_j - (x_i + w_i), x_i - (x_j + w_j))
                #             gap_y = torch.max(y_j - (y_i + h_i), y_i - (y_j + h_j))
                            
                #             # Combined gap logic:
                #             # - If both separated (both positive): Euclidean distance
                #             # - If both overlapping (both negative): use MAX (smallest overlap, easiest to fix)
                #             # - If one separated, one overlapping: use the separated one (that's the actual gap)
                #             if gap_x >= 0 and gap_y >= 0:
                #                 # Both separated: Euclidean distance
                #                 gap = torch.sqrt(gap_x * gap_x + gap_y * gap_y)
                #             elif gap_x < 0 and gap_y < 0:
                #                 # Both overlapping: use the smaller overlap (easier to separate)
                #                 gap = torch.max(gap_x, gap_y)
                #             else:
                #                 # One separated, one overlapping: they're aligned in overlapping dimension
                #                 # Use the separated dimension's gap
                #                 gap = torch.max(gap_x, gap_y)
                            
                #             # Penalty if gap < min_gap_px
                #             L_pen = L_pen + F.relu(min_gap_px - gap)

                # Position/Size similarity loss
                L_similarity = torch.tensor(0.0, device=device)
                if reference_bboxes and len(reference_bboxes) >= num_nodes:
                    # Current generated bboxes as tensors
                    generated_bboxes = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes]
                    # Generated parent container bbox (current container)
                    generated_parent_bbox = (0.0, 0.0, float(Wc), float(Hc))
                    
                    L_similarity = compute_position_size_similarity_loss(
                        reference_bboxes[:num_nodes],
                        generated_bboxes,
                        reference_parent_bbox,
                        generated_parent_bbox,
                        device=device
                    )

                # Readability loss (size hierarchy consistency)
                L_readability = torch.tensor(0.0, device=device)
                if size_rules and len(size_rules) > 0:
                    generated_bboxes_readability = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes]
                    L_readability = compute_readability_loss(
                        size_rules,
                        generated_bboxes_readability,
                        size_ratio_threshold=params.SIZE_RATIO_THRESHOLD,
                        device=device
                    )

                # Alignment consistency loss (hierarchical alignment)
                L_alignment_consistency = torch.tensor(0.0, device=device)
                if reference_bboxes and len(reference_bboxes) >= num_nodes:
                    generated_bboxes_alignment = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes]
                    generated_parent_bbox_alignment = (0.0, 0.0, float(Wc), float(Hc))
                    
                    L_alignment_consistency = compute_alignment_consistency_loss(
                        reference_bboxes[:num_nodes],
                        generated_bboxes_alignment,
                        reference_parent_bbox,
                        generated_parent_bbox_alignment,
                        device=device
                    )
                
                # Alignment similarity loss (based on JSON constraint)
                L_alignment_similarity = torch.tensor(0.0, device=device)
                if alignment_constraint and w_alignment_similarity > 0:
                    generated_bboxes_alignment_sim = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes]
                    container_bbox_alignment = (0.0, 0.0, float(Wc), float(Hc))
                    
                    L_alignment_similarity = compute_alignment_similarity_loss(
                        generated_bboxes_alignment_sim,
                        container_bbox_alignment,
                        alignment_constraint,
                        device=device
                    )

                # Proximity ratio loss
                L_proximity = torch.tensor(0.0, device=device)
                if proximity_info and w_proximity > 0:
                    # Extract information from proximity_info
                    container_bboxes = proximity_info.get("containers", [])
                    child_bboxes_list = proximity_info.get("children", [])
                    grandchild_bboxes_list = proximity_info.get("grandchildren", [])
                    container_types = proximity_info.get("types", [])
                    container_weights = proximity_info.get("weights", None)
                    
                    # Convert current generated bboxes to tuples for proximity calculation
                    # For N-element case: treat as single container with N children
                    if len(container_bboxes) == 0:
                        # Simplified N-element case: create a container with N children
                        # Note: container type should be provided in proximity_info
                        container_bbox = (0.0, 0.0, float(Wc), float(Hc))
                        # Keep bboxes as tensors for gradient computation
                        child_bboxes = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes]
                        # For N-element case, grandchildren would be empty (children are leaf nodes)
                        grandchild_bboxes = []
                        container_type = container_types[0] if container_types else "row"  # Default to row
                        
                        L_proximity = compute_proximity_ratio_loss(
                            [container_bbox],
                            [child_bboxes],
                            [grandchild_bboxes],
                            [container_type],
                            container_weights=[1.0] if container_weights is None else container_weights,
                            epsilon=params.PROXIMITY_EPSILON,
                            device=device
                        )
                    else:
                        # Use provided proximity_info
                        # Keep bboxes as tensors for gradient computation
                        generated_bboxes_proximity = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes]
                        
                        # Update child_bboxes_list with current generated bboxes if needed
                        # This is a simplified approach - in full implementation, we'd need to map
                        # generated bboxes to the correct containers
                        if len(child_bboxes_list) > 0 and len(child_bboxes_list[0]) == num_nodes:
                            # Update first container's children with generated bboxes (as tensors)
                            updated_child_bboxes_list = [generated_bboxes_proximity] + child_bboxes_list[1:]
                        else:
                            updated_child_bboxes_list = child_bboxes_list
                        
                        L_proximity = compute_proximity_ratio_loss(
                            container_bboxes,
                            updated_child_bboxes_list,
                            grandchild_bboxes_list,
                            container_types,
                            container_weights,
                            epsilon=params.PROXIMITY_EPSILON,
                            device=device
                        )

                g = A_inter
                # Data ink loss: maximize union area (minimize white space)
                # Negative because we want to maximize A_union (minimize -A_union)
                L_data_ink = -A_union
                
                # Scale pen_weight by rho to prevent L_pen from being overwhelmed when rho is large
                # When rho is large, AL constraint dominates, so we need to scale pen_weight accordingly
                pen_weight_scaled = pen_weight * (1.0 + rho.item() / 1e4)
                
                # AL constraint on overlap + penalty term + similarity loss + readability loss + alignment consistency loss + alignment similarity loss + proximity loss + data ink loss + visual balance loss
                loss = (lam * g + 0.5 * rho * g * g + pen_weight_scaled * L_pen + 
                        w_similarity * L_similarity + w_readability * L_readability + 
                        w_alignment_consistency * L_alignment_consistency + w_alignment_similarity * L_alignment_similarity + 
                        w_proximity * L_proximity + w_data_ink * L_data_ink + w_visual_balance * L_visual_balance)

                loss.backward()
                opt.step()

            # outer AL update
            with torch.no_grad():
                # Compute bboxes for logging and AL update
                bboxes_log = []
                for i in range(num_nodes):
                    tx_idx = i * 3
                    ty_idx = i * 3 + 1
                    ts_idx = i * 3 + 2
                    tx = opt_params[tx_idx]
                    ty = opt_params[ty_idx]
                    ts = opt_params[ts_idx]
                    
                    min_width, min_height = min_sizes[i]
                    x, y, w, h = bbox_aspect_from_unconstrained(
                        tx, ty, ts, Wc, Hc, ratios[i],
                        min_width=min_width, min_height=min_height,
                        size_min=size_min
                    )
                    bboxes_log.append((x, y, w, h))
                
                # Compute A_inter for AL update (always needed)
                softmasks_log = []
                for i in range(num_nodes):
                    x, y, w, h = bboxes_log[i]
                    m, d_px = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X, Y, tau_px=tau_px)
                    softmasks_log.append(m)
                
                inter_log = torch.zeros_like(softmasks_log[0])
                for i in range(num_nodes):
                    for j in range(i + 1, num_nodes):
                        inter_log = inter_log + softmasks_log[i] * softmasks_log[j]
                
                A_inter = area_sum(inter_log, Wc, Hc)
                
                # Update Lagrangian multiplier
                lam = lam + rho * A_inter
                rho = rho * rho_mult
                
                # Detailed logging and visualization (only in debug mode)
                if debug:
                    # Compute all loss components for detailed logging
                    distances_log = []
                    for i in range(num_nodes):
                        x, y, w, h = bboxes_log[i]
                        m, d_px = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X, Y, tau_px=tau_px)
                        distances_log.append(d_px)
                    
                    union_log = torch.ones_like(softmasks_log[0])
                    for m in softmasks_log:
                        union_log = union_log * (1.0 - m)
                    union_log = 1.0 - union_log
                    A_union = area_sum(union_log, Wc, Hc)

                    A_union = area_sum(union_log, Wc, Hc)

                    # Recompute penalty for logging (based on actual bbox gap)
                    # Skip for layer containers (overlapping allowed)
                    L_pen_val = torch.tensor(0.0, device=device)
                    
                    # Check if this is a layer container
                    is_layer = False
                    if proximity_info:
                        container_types = proximity_info.get("types", [])
                        if len(container_types) > 0:
                            is_layer = (container_types[0] == "layer")
                    
                    if not is_layer:
                        for i in range(num_nodes):
                            for j in range(i + 1, num_nodes):
                                x_i, y_i, w_i, h_i = bboxes_log[i]
                                x_j, y_j, w_j, h_j = bboxes_log[j]
                                
                                # Compute gap in each dimension
                                gap_x = torch.max(x_j - (x_i + w_i), x_i - (x_j + w_j))
                                gap_y = torch.max(y_j - (y_i + h_i), y_i - (y_j + h_j))
                                
                                # Combined gap logic
                                if gap_x >= 0 and gap_y >= 0:
                                    # Both separated: Euclidean distance
                                    gap = torch.sqrt(gap_x * gap_x + gap_y * gap_y)
                                elif gap_x < 0 and gap_y < 0:
                                    # Both overlapping: use smaller overlap
                                    gap = torch.max(gap_x, gap_y)
                                else:
                                    # One separated, one overlapping
                                    gap = torch.max(gap_x, gap_y)
                                
                                # Penalty if gap < min_gap_px
                                L_pen_val = L_pen_val + F.relu(min_gap_px - gap)

                    # Recompute similarity loss for logging
                    L_similarity_val = torch.tensor(0.0, device=device)
                    if reference_bboxes and len(reference_bboxes) >= num_nodes:
                        generated_bboxes_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log]
                        generated_parent_bbox_log = (0.0, 0.0, float(Wc), float(Hc))
                        
                        L_similarity_val = compute_position_size_similarity_loss(
                            reference_bboxes[:num_nodes],
                            generated_bboxes_log,
                            reference_parent_bbox,
                            generated_parent_bbox_log,
                            device=device
                        )

                    # Recompute readability loss for logging
                    L_readability_val = torch.tensor(0.0, device=device)
                    if size_rules and len(size_rules) > 0:
                        generated_bboxes_readability_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log]
                        L_readability_val = compute_readability_loss(
                            size_rules,
                            generated_bboxes_readability_log,
                            size_ratio_threshold=params.SIZE_RATIO_THRESHOLD,
                            device=device
                        )

                    # Recompute alignment consistency loss for logging
                    L_alignment_consistency_val = torch.tensor(0.0, device=device)
                    if reference_bboxes and len(reference_bboxes) >= num_nodes:
                        generated_bboxes_alignment_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log]
                        generated_parent_bbox_alignment_log = (0.0, 0.0, float(Wc), float(Hc))
                        
                        L_alignment_consistency_val = compute_alignment_consistency_loss(
                            reference_bboxes[:num_nodes],
                            generated_bboxes_alignment_log,
                            reference_parent_bbox,
                            generated_parent_bbox_alignment_log,
                            device=device
                        )
                    
                    # Recompute alignment similarity loss for logging
                    L_alignment_similarity_val = torch.tensor(0.0, device=device)
                    if alignment_constraint and w_alignment_similarity > 0:
                        generated_bboxes_alignment_sim_log = [torch.stack([bbox[0], bbox[1], bbox[2], bbox[3]]) for bbox in bboxes_log]
                        container_bbox_alignment_log = (0.0, 0.0, float(Wc), float(Hc))
                        
                        L_alignment_similarity_val = compute_alignment_similarity_loss(
                            generated_bboxes_alignment_sim_log,
                            container_bbox_alignment_log,
                            alignment_constraint,
                            device=device
                        )

                    # Recompute proximity ratio loss for logging
                    L_proximity_val = torch.tensor(0.0, device=device)
                    if proximity_info and w_proximity > 0:
                        container_bboxes = proximity_info.get("containers", [])
                        child_bboxes_list = proximity_info.get("children", [])
                        grandchild_bboxes_list = proximity_info.get("grandchildren", [])
                        container_types = proximity_info.get("types", [])
                        container_weights = proximity_info.get("weights", None)
                        
                        if len(container_bboxes) == 0:
                            container_bbox = (0.0, 0.0, float(Wc), float(Hc))
                            child_bboxes = [(bbox[0].item(), bbox[1].item(), bbox[2].item(), bbox[3].item()) for bbox in bboxes_log]
                            grandchild_bboxes = []
                            container_type = container_types[0] if container_types else "row"  # Default to row
                            
                            L_proximity_val = compute_proximity_ratio_loss(
                                [container_bbox],
                                [child_bboxes],
                                [grandchild_bboxes],
                                [container_type],
                                container_weights=[1.0] if container_weights is None else container_weights,
                                epsilon=params.PROXIMITY_EPSILON,
                                device=device
                            )
                        else:
                            generated_bboxes_proximity_log = [(bbox[0].item(), bbox[1].item(), bbox[2].item(), bbox[3].item()) for bbox in bboxes_log]
                            if len(child_bboxes_list) > 0 and len(child_bboxes_list[0]) == num_nodes:
                                updated_child_bboxes_list_log = [generated_bboxes_proximity_log] + child_bboxes_list[1:]
                            else:
                                updated_child_bboxes_list_log = child_bboxes_list
                            
                            L_proximity_val = compute_proximity_ratio_loss(
                                container_bboxes,
                                updated_child_bboxes_list_log,
                                grandchild_bboxes_list,
                                container_types,
                                container_weights,
                                epsilon=params.PROXIMITY_EPSILON,
                                device=device
                            )
                    
                    # Compute data ink loss for logging
                    L_data_ink_val = -A_union
                    
                    # Compute visual balance loss for logging (using bboxes_log)
                    L_visual_balance_val = compute_visual_balance_loss(
                        bboxes_log, Wc, Hc, device=device
                    )
                    
                    # Compute total loss for visualization
                    g_val = A_inter
                    total_loss_val = (lam.item() * g_val.item() + 0.5 * rho.item() * g_val.item() * g_val.item() +
                                     pen_weight * L_pen_val.item() +
                                     w_similarity * L_similarity_val.item() +
                                     w_readability * L_readability_val.item() +
                                     w_alignment_consistency * L_alignment_consistency_val.item() +
                                     w_alignment_similarity * L_alignment_similarity_val.item() +
                                     w_proximity * L_proximity_val.item() +
                                     w_data_ink * L_data_ink_val.item() +
                                     w_visual_balance * L_visual_balance_val.item())
                    
                    print(f"[outer {k:02d}] tau={tau_px:.3f}  A_union={A_union.item():.2f}  A_inter={A_inter.item():.6f}  "
                          f"L_pen={pen_weight*L_pen_val.item():.4f}  L_sim={w_similarity*L_similarity_val.item():.4f}  "
                          f"L_read={w_readability*L_readability_val.item():.4f}  "
                          f"L_align_cons={w_alignment_consistency*L_alignment_consistency_val.item():.4f}  "
                          f"L_align_sim={w_alignment_similarity*L_alignment_similarity_val.item():.4f}  "
                          f"L_prox={w_proximity*L_proximity_val.item():.4f}  "
                          f"L_data_ink={w_data_ink*L_data_ink_val.item():.4f}  "
                          f"L_balance={w_visual_balance*L_visual_balance_val.item():.4f}  "
                          f"lam={lam.item():.3e}  rho={rho.item():.3e}")
                    
                    # Visualize optimization progress at end of each outer epoch
                    # Use full resolution for visualization
                    X_viz, Y_viz = make_container_grid(Hc, Wc, device=device)
                    epoch_bboxes = []
                    epoch_softmasks = []
                    for i, bbox_log in enumerate(bboxes_log):
                        x, y, w, h = bbox_log
                        epoch_bboxes.append((x.item(), y.item(), w.item(), h.item()))
                        m_viz, _ = sdf_to_softmask(sdf_tensors[i], x, y, w, h, X_viz, Y_viz, tau_px=tau_px)
                        epoch_softmasks.append(m_viz)
                    
                    epoch_loss_info = {
                        'total': total_loss_val,
                        'A_union': A_union.item(),
                        'A_inter': A_inter.item(),
                        'pen': L_pen_val.item(),
                        'similarity': w_similarity * L_similarity_val.item(),
                        'readability': w_readability * L_readability_val.item(),
                        'alignment': w_alignment_consistency * L_alignment_consistency_val.item() + w_alignment_similarity * L_alignment_similarity_val.item(),
                        'alignment_consistency': w_alignment_consistency * L_alignment_consistency_val.item(),
                        'alignment_similarity': w_alignment_similarity * L_alignment_similarity_val.item(),
                        'proximity': w_proximity * L_proximity_val.item(),
                        'data_ink': w_data_ink * L_data_ink_val.item(),
                    }
                    
                    epoch_save_path = os.path.join(viz_folder, f"stage{stage_idx}_epoch{k:02d}.png")
                    
                    visualize_optimization_progress(
                        epoch_softmasks, epoch_bboxes, Wc, Hc,
                        epoch=k, loss_info=epoch_loss_info,
                        save_path=epoch_save_path
                    )
                else:
                    # Non-debug mode: simple logging
                    print(f"[outer {k:02d}] tau={tau_px:.3f}  A_inter={A_inter.item():.6f}  lam={lam.item():.3e}  rho={rho.item():.3e}")

    # final bbox (continuous)
    final_bboxes = []
    with torch.no_grad():
        for i in range(num_nodes):
            tx_idx = i * 3
            ty_idx = i * 3 + 1
            ts_idx = i * 3 + 2
            tx = opt_params[tx_idx]
            ty = opt_params[ty_idx]
            ts = opt_params[ts_idx]
            
            min_width, min_height = min_sizes[i]
            x, y, w, h = bbox_aspect_from_unconstrained(
                tx, ty, ts, Wc, Hc, ratios[i],
                min_width=min_width, min_height=min_height,
                size_min=size_min
            )
            final_bboxes.append((x.item(), y.item(), w.item(), h.item()))

    # hard evaluation at full 1000x1000: overlap using (SDF<0) AND
    with torch.no_grad():
        Xf, Yf = make_container_grid(Hc, Wc, device=device)
        # use a small tau for union display (not needed for hard overlap)
        softmasks_f = []
        distances_f = []
        for i in range(num_nodes):
            x, y, w, h = final_bboxes[i]
            m, d_px = sdf_to_softmask(sdf_tensors[i], 
                                     torch.tensor(x, device=device), 
                                     torch.tensor(y, device=device),
                                     torch.tensor(w, device=device), 
                                     torch.tensor(h, device=device), 
                                     Xf, Yf, tau_px=0.2)
            softmasks_f.append(m)
            distances_f.append(d_px)

        # Compute union
        union_f = torch.ones_like(softmasks_f[0])
        for m in softmasks_f:
            union_f = union_f * (1.0 - m)
        union_f = 1.0 - union_f
        A_union_f = union_f.sum().item()  # da=1 at full res
        
        # hard inside test: d_px < 0 for all pairs
        hard_overlap = torch.zeros_like(softmasks_f[0])
        for i in range(num_nodes):
            for j in range(i + 1, num_nodes):
                overlap_ij = ((distances_f[i] < 0.0) & (distances_f[j] < 0.0)).float()
                hard_overlap = hard_overlap + overlap_ij
        A_overlap_hard = hard_overlap.sum().item()

    # print("\n=== Final Results ===")
    # for i, bbox in enumerate(final_bboxes):
    #     print(f"bbox{i+1} (x,y,w,h) = {bbox}")
    # print(f"Union area (approx, {Wc}x{Hc}) = {A_union_f:.2f}  -> ratio {A_union_f/(Wc*Hc):.4f}")
    # print(f"Hard overlap area (SDF<0)       = {A_overlap_hard:.0f} pixels")

    # Save visualization and composite images (only in debug mode)
    if debug:
        # Determine save paths based on prefix (save to visualization folder)
        if save_prefix:
            final_result_path = os.path.join(viz_folder, f"{save_prefix}_final_result.png")
            composite_result_path = os.path.join(viz_folder, f"{save_prefix}_composite_result.png")
        else:
            final_result_path = os.path.join(viz_folder, "final_result.png")
            composite_result_path = os.path.join(viz_folder, "composite_result.png")
        
        # Visualize final result (only for 2 nodes, skip for N nodes)
        # TODO: Extend visualize_final_result to support N nodes
        if num_nodes == 2:
            mask1_orig = load_binary_mask_from_rgba(png_list[0])
            mask2_orig = load_binary_mask_from_rgba(png_list[1])
            if mask1_orig is not None and mask2_orig is not None:
                visualize_final_result(softmasks_f[0], softmasks_f[1], distances_f[0], distances_f[1], 
                                      final_bboxes[0], final_bboxes[1], 
                                      mask1_orig, mask2_orig, Wc, Hc, 
                                      save_path=final_result_path)
        
        # Save composite image with original images
        # Use original paths if provided, otherwise use the paths passed to optimize
        save_png_list = []
        for i in range(num_nodes):
            orig_png = original_png_list[i] if i < len(original_png_list) and original_png_list[i] is not None else png_list[i]
            save_png_list.append(orig_png)
        
        if all(png is not None for png in save_png_list):
            save_composite_image(png_list=save_png_list, bbox_list=final_bboxes, Wc=Wc, Hc=Hc, 
                               save_path=composite_result_path)
        else:
            print(f"Warning: Skipping composite image save - some image paths are None")

    return final_bboxes