placement.py

"""
VLSI Cell Placement Optimization Challenge
==========================================

CHALLENGE OVERVIEW:
You are tasked with implementing a critical component of a chip placement optimizer.
Given a set of cells (circuit components) with fixed sizes and connectivity requirements,
you need to find positions for these cells that:
1. Minimize total wirelength (wiring cost between connected pins)
2. Eliminate all overlaps between cells

YOUR TASK:
Implement the `overlap_repulsion_loss()` function to prevent cells from overlapping.
The function must:
- Be differentiable (uses PyTorch operations for gradient descent)
- Detect when cells overlap in 2D space
- Apply increasing penalties for larger overlaps
- Work efficiently with vectorized operations

SUCCESS CRITERIA:
After running the optimizer with your implementation:
- overlap_count should be 0 (no overlapping cell pairs)
- total_overlap_area should be 0.0 (no overlap)
- wirelength should be minimized
- Visualization should show clean, non-overlapping placement

GETTING STARTED:
1. Read through the existing code to understand the data structures
2. Look at wirelength_attraction_loss() as a reference implementation
3. Implement overlap_repulsion_loss() following the TODO instructions
4. Run main() and check the overlap metrics in the output
5. Tune hyperparameters (lambda_overlap, lambda_wirelength) if needed
6. Generate visualization to verify your solution

BONUS CHALLENGES:
- Improve convergence speed by tuning learning rate or adding momentum
- Implement better initial placement strategy
- Add visualization of optimization progress over time
"""

import os
from enum import IntEnum
from datetime import datetime

import torch
import torch.optim as optim

from arg_parse_util import parse_args
from benchmark_test_cases import TEST_CASES_BY_ID
from hyperparameter_search import run_optuna_search
from learning_rate_scheduler_util import (
    build_scheduler_kwargs_from_args,
    create_lr_scheduler,
)
from loss_tracking_utils import (
    create_loss_history,
    create_loss_tracking_db,
    save_loss_history_sqlite,
)
from torch_profiler_util import (
    build_torch_profiler_config_from_args,
    create_torch_profiler_session,
    run_with_optional_profile,
)


def get_best_device():
    """Select the fastest available torch device."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def resolve_device(device_name=None):
    """Resolve a requested torch device and validate backend availability."""
    if device_name is None or str(device_name).lower() == "auto":
        return get_best_device()

    device = torch.device(device_name)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA device requested but CUDA is not available.")
    if device.type == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("MPS device requested but MPS is not available.")
    return device


def seed_torch(seed):
    """Seed torch RNGs across supported backends."""
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


seed_torch(66)

# Feature index enums for cleaner code access
class CellFeatureIdx(IntEnum):
    """Indices for cell feature tensor columns."""
    AREA = 0
    NUM_PINS = 1
    X = 2
    Y = 3
    WIDTH = 4
    HEIGHT = 5


class PinFeatureIdx(IntEnum):
    """Indices for pin feature tensor columns."""
    CELL_IDX = 0
    PIN_X = 1  # Relative to cell corner
    PIN_Y = 2  # Relative to cell corner
    X = 3  # Absolute position
    Y = 4  # Absolute position
    WIDTH = 5
    HEIGHT = 6


# Configuration constants
# Macro parameters
MIN_MACRO_AREA = 100.0
MAX_MACRO_AREA = 10000.0

# Standard cell parameters (areas can be 1, 2, or 3)
STANDARD_CELL_AREAS = [1.0, 2.0, 3.0]
STANDARD_CELL_HEIGHT = 1.0

# Pin count parameters
MIN_STANDARD_CELL_PINS = 3
MAX_STANDARD_CELL_PINS = 6

# Output directory
OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))
# ======= SETUP =======

def generate_placement_input(num_macros, num_std_cells, device=None, verbose=True):
    """Generate synthetic placement input data.

    Args:
        num_macros: Number of macros to generate
        num_std_cells: Number of standard cells to generate

    Returns:
        Tuple of (cell_features, pin_features, edge_list):
            - cell_features: torch.Tensor of shape [N, 6] with columns [area, num_pins, x, y, width, height]
            - pin_features: torch.Tensor of shape [total_pins, 7] with columns
              [cell_instance_index, pin_x, pin_y, x, y, pin_width, pin_height]
            - edge_list: torch.Tensor of shape [E, 2] with [src_pin_idx, tgt_pin_idx]
    """
    device = device or get_best_device()
    total_cells = num_macros + num_std_cells

    # Step 1: Generate macro areas (uniformly distributed between min and max)
    macro_areas = (
        torch.rand(num_macros, device=device) * (MAX_MACRO_AREA - MIN_MACRO_AREA)
        + MIN_MACRO_AREA
    )

    # Step 2: Generate standard cell areas (randomly pick from 1, 2, or 3)
    std_cell_areas = torch.tensor(STANDARD_CELL_AREAS, device=device)[
        torch.randint(0, len(STANDARD_CELL_AREAS), (num_std_cells,), device=device)
    ]

    # Combine all areas
    areas = torch.cat([macro_areas, std_cell_areas])

    # Step 3: Calculate cell dimensions
    # Macros are square
    macro_widths = torch.sqrt(macro_areas)
    macro_heights = torch.sqrt(macro_areas)

    # Standard cells have fixed height = 1, width = area
    std_cell_widths = std_cell_areas / STANDARD_CELL_HEIGHT
    std_cell_heights = torch.full(
        (num_std_cells,),
        STANDARD_CELL_HEIGHT,
        device=device,
    )

    # Combine dimensions
    cell_widths = torch.cat([macro_widths, std_cell_widths])
    cell_heights = torch.cat([macro_heights, std_cell_heights])

    # Step 4: Calculate number of pins per cell
    num_pins_per_cell = torch.zeros(total_cells, dtype=torch.int, device=device)

    # Macros: between sqrt(area) and 2*sqrt(area) pins
    for i in range(num_macros):
        sqrt_area = int(torch.sqrt(macro_areas[i]).item())
        num_pins_per_cell[i] = torch.randint(
            sqrt_area,
            2 * sqrt_area + 1,
            (1,),
            device=device,
        ).item()

    # Standard cells: between 3 and 6 pins
    num_pins_per_cell[num_macros:] = torch.randint(
        MIN_STANDARD_CELL_PINS,
        MAX_STANDARD_CELL_PINS + 1,
        (num_std_cells,),
        device=device,
    )

    # Step 5: Create cell features tensor [area, num_pins, x, y, width, height]
    cell_features = torch.zeros(total_cells, 6, device=device)
    cell_features[:, CellFeatureIdx.AREA] = areas
    cell_features[:, CellFeatureIdx.NUM_PINS] = num_pins_per_cell.float()
    cell_features[:, CellFeatureIdx.X] = 0.0  # x position (initialized to 0)
    cell_features[:, CellFeatureIdx.Y] = 0.0  # y position (initialized to 0)
    cell_features[:, CellFeatureIdx.WIDTH] = cell_widths
    cell_features[:, CellFeatureIdx.HEIGHT] = cell_heights

    # Step 6: Generate pins for each cell
    total_pins = num_pins_per_cell.sum().item()
    pin_features = torch.zeros(total_pins, 7, device=device)

    # Fixed pin size for all pins (square pins)
    PIN_SIZE = 0.1  # All pins are 0.1 x 0.1

    pin_idx = 0
    for cell_idx in range(total_cells):
        n_pins = num_pins_per_cell[cell_idx].item()
        cell_width = cell_widths[cell_idx].item()
        cell_height = cell_heights[cell_idx].item()

        # Generate random pin positions within the cell
        # Offset from edges to ensure pins are fully inside
        margin = PIN_SIZE / 2
        if cell_width > 2 * margin and cell_height > 2 * margin:
            pin_x = (
                torch.rand(n_pins, device=device) * (cell_width - 2 * margin)
                + margin
            )
            pin_y = (
                torch.rand(n_pins, device=device) * (cell_height - 2 * margin)
                + margin
            )
        else:
            # For very small cells, just center the pins
            pin_x = torch.full((n_pins,), cell_width / 2, device=device)
            pin_y = torch.full((n_pins,), cell_height / 2, device=device)

        # Fill pin features
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.CELL_IDX] = cell_idx
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.PIN_X] = (
            pin_x  # relative to cell
        )
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.PIN_Y] = (
            pin_y  # relative to cell
        )
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.X] = (
            pin_x  # absolute (same as relative initially)
        )
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.Y] = (
            pin_y  # absolute (same as relative initially)
        )
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.WIDTH] = PIN_SIZE
        pin_features[pin_idx : pin_idx + n_pins, PinFeatureIdx.HEIGHT] = PIN_SIZE

        pin_idx += n_pins

    # Step 7: Generate edges with simple random connectivity
    # Each pin connects to 1-3 random pins (preferring different cells)
    edge_list = []
    avg_edges_per_pin = 2.0

    pin_to_cell = torch.zeros(total_pins, dtype=torch.long, device=device)
    pin_idx = 0
    for cell_idx, n_pins in enumerate(num_pins_per_cell):
        pin_to_cell[pin_idx : pin_idx + n_pins] = cell_idx
        pin_idx += n_pins

    # Create adjacency set to avoid duplicate edges
    adjacency = [set() for _ in range(total_pins)]

    for pin_idx in range(total_pins):
        pin_cell = pin_to_cell[pin_idx].item()
        num_connections = torch.randint(
            1,
            4,
            (1,),
            device=device,
        ).item()  # 1-3 connections per pin

        # Try to connect to pins from different cells
        for _ in range(num_connections):
            # Random candidate
            other_pin = torch.randint(0, total_pins, (1,), device=device).item()

            # Skip self-connections and existing connections
            if other_pin == pin_idx or other_pin in adjacency[pin_idx]:
                continue

            # Add edge (always store smaller index first for consistency)
            if pin_idx < other_pin:
                edge_list.append([pin_idx, other_pin])
            else:
                edge_list.append([other_pin, pin_idx])

            # Update adjacency
            adjacency[pin_idx].add(other_pin)
            adjacency[other_pin].add(pin_idx)

    # Convert to tensor and remove duplicates
    if edge_list:
        edge_list = torch.tensor(edge_list, dtype=torch.long, device=device)
        edge_list = torch.unique(edge_list, dim=0)
    else:
        edge_list = torch.zeros((0, 2), dtype=torch.long, device=device)

    if verbose:
        print(f"\nGenerated placement data:")
        print(f"  Total cells: {total_cells}")
        print(f"  Total pins: {total_pins}")
        print(f"  Total edges: {len(edge_list)}")
        print(f"  Average edges per pin: {2 * len(edge_list) / total_pins:.2f}")

    return cell_features, pin_features, edge_list


def initialize_cell_positions(cell_features, spread_scale=0.6):
    """Initialize cell centers with a random radial spread."""
    total_cells = cell_features.shape[0]
    total_area = cell_features[:, CellFeatureIdx.AREA].sum().item()
    spread_radius = max((total_area ** 0.5) * spread_scale, 1.0)

    angles = torch.rand(total_cells, device=cell_features.device) * 2 * torch.pi
    radii = torch.rand(total_cells, device=cell_features.device) * spread_radius

    cell_features[:, CellFeatureIdx.X] = radii * torch.cos(angles)
    cell_features[:, CellFeatureIdx.Y] = radii * torch.sin(angles)


def total_wire_length(cell_features, pin_features, edge_list):
    # the real goal seem to be to reduce the total wirelength.
    # attraction loss can be a training method.
    return 0
# ======= OPTIMIZATION CODE (edit this part) =======

def wirelength_attraction_loss(cell_features, pin_features, edge_list):
    # Q: Do I change this? 

    # Vraj: change this later once the overlap loss is tuned
    # Ans: No there are down stream calls that need this implementation to help evaluate the result

    # Keep this code for testing and create a loss function when needed.
    """Calculate loss based on total wirelength to minimize routing.

    This is a REFERENCE IMPLEMENTATION showing how to write a differentiable loss function.

    The loss computes the Manhattan distance between connected pins and minimizes
    the total wirelength across all edges.

    Args:
        cell_features: [N, 6] tensor with [area, num_pins, x, y, width, height]
        pin_features: [P, 7] tensor with pin information
        edge_list: [E, 2] tensor with edges

    Returns:
        Scalar loss value
    """
    if edge_list.shape[0] == 0:
        return torch.tensor(
            0.0,
            requires_grad=True,
            device=cell_features.device,
            dtype=cell_features.dtype,
        )

    # Update absolute pin positions based on cell positions
    cell_positions = cell_features[:, 2:4]  # [N, 2]
    cell_indices = pin_features[:, 0].long()

    # Calculate absolute pin positions
    pin_absolute_x = cell_positions[cell_indices, 0] + pin_features[:, 1]
    pin_absolute_y = cell_positions[cell_indices, 1] + pin_features[:, 2]

    # Get source and target pin positions for each edge
    src_pins = edge_list[:, 0].long()
    tgt_pins = edge_list[:, 1].long()

    src_x = pin_absolute_x[src_pins]
    src_y = pin_absolute_y[src_pins]
    tgt_x = pin_absolute_x[tgt_pins]
    tgt_y = pin_absolute_y[tgt_pins]

    # Calculate smooth approximation of Manhattan distance
    # Using log-sum-exp approximation for differentiability
    alpha = 0.1  # Smoothing parameter
    dx = torch.abs(src_x - tgt_x)
    dy = torch.abs(src_y - tgt_y)

    # Smooth L1 distance with numerical stability
    smooth_manhattan = alpha * torch.logsumexp(
        torch.stack([dx / alpha, dy / alpha], dim=0), dim=0
    )

    # Total wirelength
    total_wirelength = torch.sum(smooth_manhattan)
    ret = total_wirelength / edge_list.shape[0]  # Normalize by number of edges
    # print(ret.shape)
    return ret


def compute_pairwise_overlap_areas(cell_features):
    """Return pairwise overlap areas for all cell pairs."""
    num_cells = cell_features.shape[0]
    if num_cells <= 1:
        return torch.zeros(
            (num_cells, num_cells),
            device=cell_features.device,
            dtype=cell_features.dtype,
        )

    x_col = cell_features[:, CellFeatureIdx.X]
    y_col = cell_features[:, CellFeatureIdx.Y]
    widths = cell_features[:, CellFeatureIdx.WIDTH]
    heights = cell_features[:, CellFeatureIdx.HEIGHT]

    x_delta = torch.abs(x_col.unsqueeze(1) - x_col.unsqueeze(0))
    y_delta = torch.abs(y_col.unsqueeze(1) - y_col.unsqueeze(0))

    x_span = (widths.unsqueeze(1) + widths.unsqueeze(0)) / 2
    y_span = (heights.unsqueeze(1) + heights.unsqueeze(0)) / 2

    overlap_x = torch.relu(x_span - x_delta)
    overlap_y = torch.relu(y_span - y_delta)
    return overlap_x * overlap_y


def calculate_overlap_metrics_torch(cell_features):
    """Calculate overlap metrics with vectorized torch operations."""
    num_cells = cell_features.shape[0]
    if num_cells <= 1:
        zero = torch.tensor(0.0, device=cell_features.device, dtype=cell_features.dtype)
        return {
            "overlap_count": 0,
            "total_overlap_area": 0.0,
            "max_overlap_area": 0.0,
            "overlap_percentage": 0.0,
            "cells_with_overlap": 0,
            "has_zero_overlap": True,
            "total_overlap_area_tensor": zero,
            "max_overlap_area_tensor": zero,
        }

    pairwise_overlap_area = compute_pairwise_overlap_areas(cell_features)
    mask = torch.triu(
        torch.ones_like(pairwise_overlap_area, dtype=torch.bool),
        diagonal=1,
    )
    active_overlap_areas = pairwise_overlap_area[mask]
    overlapping_pairs = active_overlap_areas > 0

    overlap_count = int(overlapping_pairs.sum().item())
    total_overlap_area = active_overlap_areas.sum()
    max_overlap_area = (
        active_overlap_areas.max()
        if active_overlap_areas.numel() > 0
        else total_overlap_area.new_zeros(())
    )

    overlap_matrix = (pairwise_overlap_area > 0) & mask
    overlap_matrix = overlap_matrix | overlap_matrix.transpose(0, 1)
    cells_with_overlap = int(overlap_matrix.any(dim=0).sum().item())
    total_area = cell_features[:, CellFeatureIdx.AREA].sum()
    overlap_percentage = (
        overlap_count / num_cells * 100.0 if total_area.item() > 0 else 0.0
    )

    return {
        "overlap_count": overlap_count,
        "total_overlap_area": float(total_overlap_area.item()),
        "max_overlap_area": float(max_overlap_area.item()),
        "overlap_percentage": overlap_percentage,
        "cells_with_overlap": cells_with_overlap,
        "has_zero_overlap": overlap_count == 0,
        "total_overlap_area_tensor": total_overlap_area,
        "max_overlap_area_tensor": max_overlap_area,
    }


def overlap_repulsion_loss(cell_features, pin_features, edge_list):
    """Calculate loss to prevent cell overlaps.

    TODO: IMPLEMENT THIS FUNCTION

    This is the main challenge. You need to implement a differentiable loss function
    that penalizes overlapping cells. The loss should:

    1. Be zero when no cells overlap
    2. Increase as overlap area increases
    3. Use only differentiable PyTorch operations (no if statements on tensors)
    4. Work efficiently with vectorized operations

    HINTS:
    - Two axis-aligned rectangles overlap if they overlap in BOTH x and y dimensions
    - For rectangles centered at (x1, y1) and (x2, y2) with widths (w1, w2) and heights (h1, h2):
      * x-overlap occurs when |x1 - x2| < (w1 + w2) / 2
      * y-overlap occurs when |y1 - y2| < (h1 + h2) / 2
    - Use torch.relu() to compute positive overlaps: overlap_x = relu((w1+w2)/2 - |x1-x2|)
    - Overlap area = overlap_x * overlap_y
    - Consider all pairs of cells: use broadcasting with unsqueeze
    - Use torch.triu() to avoid counting each pair twice (only consider i < j)
    - Normalize the loss appropriately (by number of pairs or total area)

    RECOMMENDED APPROACH:
    1. Extract positions, widths, heights from cell_features
    2. Compute all pairwise distances using broadcasting:
       positions_i = positions.unsqueeze(1)  # [N, 1, 2]
       positions_j = positions.unsqueeze(0)  # [1, N, 2]
       distances = positions_i - positions_j  # [N, N, 2]
    3. Calculate minimum separation distances for each pair
    4. Use relu to get positive overlap amounts
    5. Multiply overlaps in x and y to get overlap areas
    6. Mask to only consider upper triangle (i < j)
    7. Sum and normalize

    Args:
        cell_features: [N, 6] tensor with [area, num_pins, x, y, width, height]
        pin_features: [P, 7] tensor with pin information (not used here)
        edge_list: [E, 2] tensor with edges (not used here)

    Returns:
        Scalar loss value (should be 0 when no overlaps exist)
    """
    N = cell_features.shape[0]
    if N <= 1:
        return torch.tensor(0.0, requires_grad=True, device=cell_features.device)

    pairwise_overlap_area = compute_pairwise_overlap_areas(cell_features)
    mask = torch.triu(torch.ones_like(pairwise_overlap_area), diagonal=1)

    # normalization = torch.sqrt(
    #     torch.tensor(N, device=pairwise_overlap_area.device, dtype=pairwise_overlap_area.dtype)
    # )
    # loss = torch.sum(pairwise_overlap_area * mask) / normalization
    overlap_sclar = 200
    loss = torch.log1p(torch.sum(pairwise_overlap_area * mask)) * overlap_sclar

    return loss

    # TODO: Implement overlap detection and loss calculation here
    #
    # Your implementation should:
    # 1. Extract cell positions, widths, and heights
    # 2. Compute pairwise overlaps using vectorized operations
    # 3. Return a scalar loss that is zero when no overlaps exist
    #
    # Delete this placeholder and add your implementation:

def train_placement(
    cell_features,
    pin_features,
    edge_list,
    num_epochs=1000,
    lr=0.1,
    lambda_wirelength=3.0,
    lambda_overlap=1.0,
    scheduler_name="plateau",
    scheduler_kwargs=None,
    track_loss_history=True,
    verbose=True,
    log_interval=100,
    run_metadata=None,
    torch_profiler_config=None,
    torch_profile_output_dir=None,
    track_overlap_metrics=False,
    early_stop_enabled=True,
    early_stop_patience=75,
    early_stop_min_delta=1e-4,
    early_stop_overlap_threshold=0.0,
    early_stop_zero_overlap_patience=25,
    device=None,
):
    """Train the placement optimization using gradient descent.

    Args:
        cell_features: [N, 6] tensor with cell properties
        pin_features: [P, 7] tensor with pin properties
        edge_list: [E, 2] tensor with edge connectivity
        num_epochs: Number of optimization iterations
        lr: Learning rate for Adam optimizer
        lambda_wirelength: Weight for wirelength loss
        lambda_overlap: Weight for overlap loss
        scheduler_name: Learning-rate scheduler name
        scheduler_kwargs: Scheduler-specific keyword arguments
        track_loss_history: Whether to collect per-epoch loss history
        verbose: Whether to print progress
        log_interval: How often to print progress
        run_metadata: Optional metadata describing the run
        torch_profiler_config: Optional torch profiler configuration
        torch_profile_output_dir: Base directory for torch profiler artifacts
        track_overlap_metrics: Whether to collect per-epoch overlap metrics
        early_stop_enabled: Whether to stop once overlap-first convergence stalls
        early_stop_patience: Plateau patience before zero-overlap is reached
        early_stop_min_delta: Minimum improvement to reset patience
        early_stop_overlap_threshold: Treat overlap below this as effectively zero
        early_stop_zero_overlap_patience: Extra patience after zero-overlap to keep improving wirelength
        device: Optional torch device override. When omitted, CPU inputs are moved
            to the best available device.

    Returns:
        Dictionary with:
            - final_cell_features: Optimized cell positions
            - initial_cell_features: Original cell positions (for comparison)
            - loss_history: Loss values over time
    """
    if device is None:
        device = cell_features.device
        if device.type == "cpu":
            device = get_best_device()
    else:
        device = resolve_device(device)

    # Clone features and create learnable positions
    cell_features = cell_features.clone().to(device)
    pin_features = pin_features.to(device)
    edge_list = edge_list.to(device)
    initial_cell_features = cell_features.clone()

    # Make only cell positions require gradients
    cell_positions = cell_features[:, 2:4].clone().detach()
    cell_positions.requires_grad_(True)

    # Create optimizer
    scheduler_kwargs = dict(scheduler_kwargs or {})
    optimizer = optim.Adam([cell_positions], lr=lr)
    scheduler, scheduler_uses_metric = create_lr_scheduler(
        optimizer,
        scheduler_name=scheduler_name,
        num_epochs=num_epochs,
        scheduler_kwargs=scheduler_kwargs,
    )

    history_run_metadata = {
        "run_label": "train_placement",
        "run_started_at": datetime.now().isoformat(timespec="seconds"),
        "device": str(device),
        "num_epochs": num_epochs,
        "lr": lr,
        "lambda_wirelength": lambda_wirelength,
        "lambda_overlap": lambda_overlap,
        "scheduler_name": scheduler_name,
        "scheduler_kwargs": scheduler_kwargs,
        "track_loss_history": track_loss_history,
        "track_overlap_metrics": track_overlap_metrics,
        "early_stop_enabled": early_stop_enabled,
        "early_stop_patience": early_stop_patience,
        "early_stop_min_delta": early_stop_min_delta,
        "early_stop_overlap_threshold": early_stop_overlap_threshold,
        "early_stop_zero_overlap_patience": early_stop_zero_overlap_patience,
        "log_interval": log_interval,
        "verbose": verbose,
        "total_cells": int(cell_features.shape[0]),
        "total_pins": int(pin_features.shape[0]),
        "total_edges": int(edge_list.shape[0]),
    }
    if run_metadata:
        history_run_metadata.update(run_metadata)

    loss_history = None
    if track_loss_history:
        loss_history = create_loss_history(
            run_metadata=history_run_metadata,
            track_overlap_metrics=track_overlap_metrics,
        )

    best_cell_positions = cell_positions.detach().clone()
    best_overlap_score = float("inf")
    best_zero_overlap_wl = float("inf")
    best_epoch = -1
    epochs_without_improvement = 0
    zero_overlap_epochs_without_improvement = 0
    zero_overlap_reached = False
    stopped_early = False
    stop_reason = ""

    # Training loop
    profiler_output_dir = torch_profile_output_dir or OUTPUT_DIR
    with create_torch_profiler_session(
        config=torch_profiler_config,
        output_dir=profiler_output_dir,
        profile_tag=history_run_metadata.get("profile_tag", ""),
        run_metadata=history_run_metadata,
    ) as profiler_session:
        for epoch in range(num_epochs):
            with torch.profiler.record_function("placement/epoch"):
                overlap_metrics = None
                optimizer.zero_grad()

                with torch.profiler.record_function("placement/forward"):
                    cell_features_current = cell_features.clone()
                    cell_features_current[:, 2:4] = cell_positions

                    wl_loss = wirelength_attraction_loss(
                        cell_features_current, pin_features, edge_list
                    )
                    overlap_loss = overlap_repulsion_loss(
                        cell_features_current, pin_features, edge_list
                    )
                    total_loss = (
                        lambda_wirelength * wl_loss + lambda_overlap * overlap_loss
                    )

                with torch.profiler.record_function("placement/backward"):
                    total_loss.backward()
                    torch.nn.utils.clip_grad_norm_([cell_positions], max_norm=5.0)

                with torch.profiler.record_function("placement/optimizer_step"):
                    optimizer.step()
                    if scheduler is not None:
                        if scheduler_uses_metric:
                            scheduler.step(total_loss.item())
                        else:
                            scheduler.step()

                should_log_epoch = verbose and (
                    epoch % log_interval == 0 or epoch == num_epochs - 1
                )
                should_compute_overlap_metrics = (
                    track_overlap_metrics
                    or early_stop_enabled
                    or should_log_epoch
                )
                updated_cell_features = None
                if should_compute_overlap_metrics:
                    with torch.profiler.record_function("placement/metrics"):
                        updated_cell_features = cell_features.clone()
                        updated_cell_features[:, 2:4] = cell_positions.detach()
                        overlap_metrics = calculate_overlap_metrics_torch(
                            updated_cell_features
                        )

                if early_stop_enabled:
                    overlap_score = overlap_metrics["total_overlap_area"]
                    has_zero_overlap = (
                        overlap_metrics["overlap_count"] == 0
                        or overlap_score <= early_stop_overlap_threshold
                    )
                    if has_zero_overlap:
                        current_wl = wirelength_attraction_loss(
                            updated_cell_features,
                            pin_features,
                            edge_list,
                        ).item()
                        if (
                            not zero_overlap_reached
                            or current_wl < best_zero_overlap_wl - early_stop_min_delta
                        ):
                            zero_overlap_reached = True
                            best_zero_overlap_wl = current_wl
                            best_cell_positions = cell_positions.detach().clone()
                            best_epoch = epoch
                            zero_overlap_epochs_without_improvement = 0
                        else:
                            zero_overlap_epochs_without_improvement += 1

                        if (
                            zero_overlap_reached
                            and zero_overlap_epochs_without_improvement
                            >= early_stop_zero_overlap_patience
                        ):
                            stopped_early = True
                            stop_reason = "zero_overlap_plateau"
                    else:
                        if zero_overlap_reached:
                            zero_overlap_epochs_without_improvement += 1
                            if (
                                zero_overlap_epochs_without_improvement
                                >= early_stop_zero_overlap_patience
                            ):
                                stopped_early = True
                                stop_reason = "zero_overlap_plateau"
                        else:
                            if overlap_score < best_overlap_score - early_stop_min_delta:
                                best_overlap_score = overlap_score
                                best_cell_positions = cell_positions.detach().clone()
                                best_epoch = epoch
                                epochs_without_improvement = 0
                            else:
                                epochs_without_improvement += 1

                            if epochs_without_improvement >= early_stop_patience:
                                stopped_early = True
                                stop_reason = "overlap_plateau"
                should_collect_overlap_metrics = track_overlap_metrics and loss_history is not None

                if loss_history is not None:
                    loss_history["total_loss"].append(total_loss.item())
                    loss_history["wirelength_loss"].append(wl_loss.item())
                    loss_history["overlap_loss"].append(overlap_loss.item())
                    loss_history["learning_rate"].append(optimizer.param_groups[0]["lr"])
                    if should_collect_overlap_metrics:
                        loss_history["overlap_count"].append(
                            overlap_metrics["overlap_count"]
                        )
                        loss_history["total_overlap_area"].append(
                            overlap_metrics["total_overlap_area"]
                        )
                        loss_history["max_overlap_area"].append(
                            overlap_metrics["max_overlap_area"]
                        )

                if should_log_epoch:
                    print(f"Epoch {epoch}/{num_epochs}:")
                    print(f"  Total Loss: {total_loss.item():.6f}")
                    print(f"  Wirelength Loss: {wl_loss.item():.6f}")
                    print(f"  Overlap Loss: {overlap_loss.item():.6f}")
                    print(f"  Learning Rate: {optimizer.param_groups[0]['lr']:.6f}")
                    if overlap_metrics is not None:
                        print(f"  Overlap Count: {overlap_metrics['overlap_count']}")
                        print(
                            f"  Total Overlap Area: {overlap_metrics['total_overlap_area']:.6f}"
                        )
                    if early_stop_enabled:
                        print(f"  Best Epoch: {best_epoch}")

                if stopped_early:
                    if verbose:
                        print(
                            f"Early stopping at epoch {epoch} "
                            f"with reason={stop_reason} best_epoch={best_epoch}"
                        )
                    break

            profiler_session.step()

    # Create final cell features
    final_cell_features = cell_features.clone()
    final_positions = best_cell_positions if early_stop_enabled else cell_positions.detach()
    final_cell_features[:, 2:4] = final_positions

    if loss_history is not None:
        loss_history["run_metadata"]["stopped_early"] = stopped_early
        loss_history["run_metadata"]["stop_reason"] = stop_reason
        loss_history["run_metadata"]["best_epoch"] = best_epoch

    return {
        "final_cell_features": final_cell_features,
        "initial_cell_features": initial_cell_features,
        "loss_history": loss_history,
        "stopped_early": stopped_early,
        "stop_reason": stop_reason,
        "best_epoch": best_epoch,
    }


# ======= FINAL EVALUATION CODE (Don't edit this part) =======

def calculate_overlap_metrics(cell_features):
    """Calculate ground truth overlap statistics (non-differentiable).

    This function provides exact overlap measurements for evaluation and reporting.
    Unlike the loss function, this does NOT need to be differentiable.

    Args:
        cell_features: [N, 6] tensor with [area, num_pins, x, y, width, height]

    Returns:
        Dictionary with:
            - overlap_count: number of overlapping cell pairs (int)
            - total_overlap_area: sum of all overlap areas (float)
            - max_overlap_area: largest single overlap area (float)
            - overlap_percentage: percentage of total area that overlaps (float)
    """
    N = cell_features.shape[0]
    if N <= 1:
        return {
            "overlap_count": 0,
            "total_overlap_area": 0.0,
            "max_overlap_area": 0.0,
            "overlap_percentage": 0.0,
        }

    # Extract cell properties
    positions = cell_features[:, 2:4].detach().cpu().numpy()  # [N, 2]
    widths = cell_features[:, 4].detach().cpu().numpy()  # [N]
    heights = cell_features[:, 5].detach().cpu().numpy()  # [N]
    areas = cell_features[:, 0].detach().cpu().numpy()  # [N]

    overlap_count = 0
    total_overlap_area = 0.0
    max_overlap_area = 0.0
    overlap_areas = []

    # Check all pairs
    for i in range(N):
        for j in range(i + 1, N):
            # Calculate center-to-center distances
            dx = abs(positions[i, 0] - positions[j, 0])
            dy = abs(positions[i, 1] - positions[j, 1])

            # Minimum separation for non-overlap
            min_sep_x = (widths[i] + widths[j]) / 2
            min_sep_y = (heights[i] + heights[j]) / 2

            # Calculate overlap amounts
            overlap_x = max(0, min_sep_x - dx)
            overlap_y = max(0, min_sep_y - dy)

            # Overlap occurs only if both x and y overlap
            if overlap_x > 0 and overlap_y > 0:
                overlap_area = overlap_x * overlap_y
                overlap_count += 1
                total_overlap_area += overlap_area
                max_overlap_area = max(max_overlap_area, overlap_area)
                overlap_areas.append(overlap_area)

    # Calculate percentage of total area
    total_area = sum(areas)
    overlap_percentage = (overlap_count / N * 100) if total_area > 0 else 0.0

    return {
        "overlap_count": overlap_count,
        "total_overlap_area": total_overlap_area,
        "max_overlap_area": max_overlap_area,
        "overlap_percentage": overlap_percentage,
    }


def calculate_cells_with_overlaps(cell_features):
    """Calculate number of cells involved in at least one overlap.

    This metric matches the test suite evaluation criteria.

    Args:
        cell_features: [N, 6] tensor with cell properties

    Returns:
        Set of cell indices that have overlaps with other cells
    """
    N = cell_features.shape[0]
    if N <= 1:
        return set()

    # Extract cell properties
    positions = cell_features[:, 2:4].detach().cpu().numpy()
    widths = cell_features[:, 4].detach().cpu().numpy()
    heights = cell_features[:, 5].detach().cpu().numpy()

    cells_with_overlaps = set()

    # Check all pairs
    for i in range(N):
        for j in range(i + 1, N):
            # Calculate center-to-center distances
            dx = abs(positions[i, 0] - positions[j, 0])
            dy = abs(positions[i, 1] - positions[j, 1])

            # Minimum separation for non-overlap
            min_sep_x = (widths[i] + widths[j]) / 2
            min_sep_y = (heights[i] + heights[j]) / 2

            # Calculate overlap amounts
            overlap_x = max(0, min_sep_x - dx)
            overlap_y = max(0, min_sep_y - dy)

            # Overlap occurs only if both x and y overlap
            if overlap_x > 0 and overlap_y > 0:
                cells_with_overlaps.add(i)
                cells_with_overlaps.add(j)

    return cells_with_overlaps


def calculate_normalized_metrics(cell_features, pin_features, edge_list):
    """Calculate normalized overlap and wirelength metrics for test suite.

    These metrics match the evaluation criteria in the test suite.

    Args:
        cell_features: [N, 6] tensor with cell properties
        pin_features: [P, 7] tensor with pin properties
        edge_list: [E, 2] tensor with edge connectivity

    Returns:
        Dictionary with:
            - overlap_ratio: (num cells with overlaps / total cells)
            - normalized_wl: (wirelength / num nets) / sqrt(total area)
            - num_cells_with_overlaps: number of unique cells involved in overlaps
            - total_cells: total number of cells
            - num_nets: number of nets (edges)
    """
    N = cell_features.shape[0]

    # Calculate overlap metric: num cells with overlaps / total cells
    cells_with_overlaps = calculate_cells_with_overlaps(cell_features)
    num_cells_with_overlaps = len(cells_with_overlaps)
    overlap_ratio = num_cells_with_overlaps / N if N > 0 else 0.0

    # Calculate wirelength metric: (wirelength / num nets) / sqrt(total area)
    if edge_list.shape[0] == 0:
        normalized_wl = 0.0
        num_nets = 0
    else:
        # Calculate total wirelength using the loss function (unnormalized)
        wl_loss = wirelength_attraction_loss(cell_features, pin_features, edge_list)
        total_wirelength = wl_loss.item() * edge_list.shape[0]  # Undo normalization

        # Calculate total area
        total_area = cell_features[:, 0].sum().item()

        num_nets = edge_list.shape[0]

        # Normalize: (wirelength / net) / sqrt(area)
        # This gives a dimensionless quality metric independent of design size
        normalized_wl = (total_wirelength / num_nets) / (total_area ** 0.5) if total_area > 0 else 0.0

    return {
        "overlap_ratio": overlap_ratio,
        "normalized_wl": normalized_wl,
        "num_cells_with_overlaps": num_cells_with_overlaps,
        "total_cells": N,
        "num_nets": num_nets,
    }


def plot_placement(
    initial_cell_features,
    final_cell_features,
    pin_features,
    edge_list,
    filename="placement_result.png",
):
    """Create side-by-side visualization of initial vs final placement.

    Args:
        initial_cell_features: Initial cell positions and properties
        final_cell_features: Optimized cell positions and properties
        pin_features: Pin information
        edge_list: Edge connectivity
        filename: Output filename for the plot
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.patches import Rectangle

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))

        # Plot both initial and final placements
        for ax, cell_features, title in [
            (ax1, initial_cell_features, "Initial Placement"),
            (ax2, final_cell_features, "Final Placement"),
        ]:
            N = cell_features.shape[0]
            positions = cell_features[:, 2:4].detach().cpu().numpy()
            widths = cell_features[:, 4].detach().cpu().numpy()
            heights = cell_features[:, 5].detach().cpu().numpy()

            # Draw cells
            for i in range(N):
                x = positions[i, 0] - widths[i] / 2
                y = positions[i, 1] - heights[i] / 2
                rect = Rectangle(
                    (x, y),
                    widths[i],
                    heights[i],
                    fill=True,
                    facecolor="lightblue",
                    edgecolor="darkblue",
                    linewidth=0.5,
                    alpha=0.7,
                )
                ax.add_patch(rect)

            # Calculate and display overlap metrics
            metrics = calculate_overlap_metrics(cell_features)

            ax.set_aspect("equal")
            ax.grid(True, alpha=0.3)
            ax.set_title(
                f"{title}\n"
                f"Overlaps: {metrics['overlap_count']}, "
                f"Total Overlap Area: {metrics['total_overlap_area']:.2f}",
                fontsize=12,
            )

            # Set axis limits with margin
            all_x = positions[:, 0]
            all_y = positions[:, 1]
            margin = 10
            ax.set_xlim(all_x.min() - margin, all_x.max() + margin)
            ax.set_ylim(all_y.min() - margin, all_y.max() + margin)

        plt.tight_layout()
        output_path = os.path.join(OUTPUT_DIR, filename)
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close()

    except ImportError as e:
        print(f"Could not create visualization: {e}")
        print("Install matplotlib to enable visualization: pip install matplotlib")

# ======= MAIN FUNCTION =======

def main(args):
    """Main function demonstrating the placement optimization challenge."""
    torch_profiler_config = build_torch_profiler_config_from_args(args)
    device = resolve_device(args.device)
    if args.optuna:
        run_optuna_search(
            args,
            get_best_device=lambda: device,
            seed_torch=seed_torch,
            generate_placement_input=generate_placement_input,
            initialize_cell_positions=initialize_cell_positions,
            train_placement=train_placement,
            calculate_normalized_metrics=calculate_normalized_metrics,
        )
        return

    print("=" * 70)
    print("VLSI CELL PLACEMENT OPTIMIZATION CHALLENGE")
    print("=" * 70)
    print("\nObjective: Implement overlap_repulsion_loss() to eliminate cell overlaps")
    print("while minimizing wirelength.\n")

    test_case = None
    if args.test_case_id is not None:
        test_case = TEST_CASES_BY_ID[args.test_case_id]

    # Set random seed for reproducibility
    seed = test_case["seed"] if test_case is not None else args.seed
    seed_torch(seed)

    # Generate placement problem
    num_macros = (
        test_case["num_macros"] if test_case is not None else args.num_macros
    )
    num_std_cells = (
        test_case["num_std_cells"] if test_case is not None else args.num_std_cells
    )

    print(f"Generating placement problem:")
    if test_case is not None:
        print(f"  - benchmark test case: {test_case['test_id']}")
    print(f"  - {num_macros} macros")
    print(f"  - {num_std_cells} standard cells")
    print(f"  - seed: {seed}")
    print(f"  - device: {device}")

    cell_features, pin_features, edge_list = generate_placement_input(
        num_macros,
        num_std_cells,
        device=device,
    )

    # Initialize positions with random spread to reduce initial overlaps
    initialize_cell_positions(cell_features)

    # Calculate initial metrics
    print("\n" + "=" * 70)
    print("INITIAL STATE")
    print("=" * 70)
    initial_metrics = calculate_overlap_metrics(cell_features)
    print(f"Overlap count: {initial_metrics['overlap_count']}")
    print(f"Total overlap area: {initial_metrics['total_overlap_area']:.2f}")
    print(f"Max overlap area: {initial_metrics['max_overlap_area']:.2f}")
    print(f"Overlap percentage: {initial_metrics['overlap_percentage']:.2f}%")

    # Run optimization
    print("\n" + "=" * 70)
    print("RUNNING OPTIMIZATION")
    print("=" * 70)

    loss_tracking_db_path = None
    if args.track_loss_history:
        loss_tracking_db_path = create_loss_tracking_db(OUTPUT_DIR)

    result = train_placement(
        cell_features,
        pin_features,
        edge_list,
        num_epochs=args.num_epochs,
        lr=args.lr,
        lambda_wirelength=args.lambda_wirelength,
        lambda_overlap=args.lambda_overlap,
        scheduler_name=args.scheduler,
        scheduler_kwargs=build_scheduler_kwargs_from_args(args),
        track_loss_history=args.track_loss_history,
        verbose=True,
        log_interval=200,
        run_metadata={
            "runner": "placement.main",
            "profile_tag": args.profile_tag,
            "seed": seed,
            "num_macros": num_macros,
            "num_std_cells": num_std_cells,
            "test_id": None if test_case is None else test_case["test_id"],
        },
        torch_profiler_config=torch_profiler_config,
        torch_profile_output_dir=OUTPUT_DIR,
        track_overlap_metrics=args.track_overlap_metrics,
        early_stop_enabled=args.early_stop,
        early_stop_patience=args.early_stop_patience,
        early_stop_min_delta=args.early_stop_min_delta,
        early_stop_overlap_threshold=args.early_stop_overlap_threshold,
        early_stop_zero_overlap_patience=args.early_stop_zero_overlap_patience,
        device=device,
    )
    if args.track_loss_history:
        loss_history_path = save_loss_history_sqlite(
            result["loss_history"],
            loss_tracking_db_path,
        )
        print(f"Loss history saved to: {loss_history_path}")
    else:
        print("Loss history tracking disabled.")

    # Calculate final metrics (both detailed and normalized)
    print("\n" + "=" * 70)
    print("FINAL RESULTS")
    print("=" * 70)

    final_cell_features = result["final_cell_features"]

    # Detailed metrics
    final_metrics = calculate_overlap_metrics(final_cell_features)
    print(f"Overlap count (pairs): {final_metrics['overlap_count']}")
    print(f"Total overlap area: {final_metrics['total_overlap_area']:.2f}")
    print(f"Max overlap area: {final_metrics['max_overlap_area']:.2f}")

    # Normalized metrics (matching test suite)
    print("\n" + "-" * 70)
    print("TEST SUITE METRICS (for leaderboard)")
    print("-" * 70)
    normalized_metrics = calculate_normalized_metrics(
        final_cell_features, pin_features, edge_list
    )
    print(f"Overlap Ratio: {normalized_metrics['overlap_ratio']:.4f} "
          f"({normalized_metrics['num_cells_with_overlaps']}/{normalized_metrics['total_cells']} cells)")
    print(f"Normalized Wirelength: {normalized_metrics['normalized_wl']:.4f}")

    # Success check
    print("\n" + "=" * 70)
    print("SUCCESS CRITERIA")
    print("=" * 70)
    if normalized_metrics["num_cells_with_overlaps"] == 0:
        print("✓ PASS: No overlapping cells!")
        print("✓ PASS: Overlap ratio is 0.0")
        print("\nCongratulations! Your implementation successfully eliminated all overlaps.")
        print(f"Your normalized wirelength: {normalized_metrics['normalized_wl']:.4f}")
    else:
        print("✗ FAIL: Overlaps still exist")
        print(f"  Need to eliminate overlaps in {normalized_metrics['num_cells_with_overlaps']} cells")
        print("\nSuggestions:")
        print("  1. Check your overlap_repulsion_loss() implementation")
        print("  2. Change lambdas (try increasing lambda_overlap)")
        print("  3. Change learning rate or number of epochs")

    # Generate visualization
    plot_placement(
        result["initial_cell_features"],
        result["final_cell_features"],
        pin_features,
        edge_list,
        filename="placement_result.png",
    )

if __name__ == "__main__":
    args = parse_args()
    run_with_optional_profile(lambda: main(args), args, OUTPUT_DIR)