from inire.geometry.collision import CollisionEngine
from inire.geometry.primitives import Port
from inire.router.astar import AStarContext, route_astar
from inire.router.cost import CostEvaluator
from inire.router.danger_map import DangerMap
from inire.router.pathfinder import PathFinder
from inire.utils.visualization import plot_routing_results


def main() -> None:
    print("Running Example 04: SBends and Radii Strategy...")

    # 1. Setup Environment
    bounds = (0, 0, 100, 100)
    engine = CollisionEngine(clearance=2.0)
    danger_map = DangerMap(bounds=bounds)
    danger_map.precompute([])

    evaluator = CostEvaluator(engine, danger_map, bend_penalty=200.0, sbend_penalty=400.0)
    
    # Define a custom router with multiple SBend radii and specific offsets
    context = AStarContext(
        evaluator,
        snap_size=1.0,
        bend_radii=[20.0, 50.0],
        sbend_radii=[5.0, 10.0, 50.0],
        sbend_offsets=[2.0, 5.0, 10.0, 20.0, 50.0]
    )
    pf = PathFinder(context)

    # 2. Define Netlist
    # High-density parallel nets with varying offsets
    netlist = {}
    for i in range(10):
        # Starts at x=50, y=50+i*10. Targets at x=450, y=60+i*10.
        # This forces small vertical jogs (SBends)
        netlist[f"net_{i}"] = (Port(50, 50 + i * 10, 0), Port(450, 55 + i * 10, 0))
    
    net_widths = {nid: 2.0 for nid in netlist}

    # 3. Route
    print(f"Routing {len(netlist)} nets with custom SBend strategy...")
    results = pf.route_all(netlist, net_widths, shuffle_nets=True)

    # 4. Visualize
    fig, ax = plot_routing_results(results, [], bounds, netlist=netlist)
    fig.savefig("examples/04_sbends_and_radii.png")
    print("Saved plot to examples/04_sbends_and_radii.png")


if __name__ == "__main__":
    main()