from shapely.geometry import Polygon
from inire.geometry.collision import CollisionEngine
from inire.geometry.primitives import Port
from inire.router.astar import AStarRouter
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 08: Custom Bend Geometry Models...")

    bounds = (0, 0, 150, 150)
    engine = CollisionEngine(clearance=2.0)
    
    # Static obstacle to force specific bend paths
    obstacle = Polygon([(60, 40), (90, 40), (90, 110), (60, 110)])
    engine.add_static_obstacle(obstacle)
    
    danger_map = DangerMap(bounds=bounds)
    danger_map.precompute([obstacle])
    evaluator = CostEvaluator(engine, danger_map)

    # We will route three nets, each with a DIFFERENT collision model
    # To do this cleanly with the current architecture, we'll use one router 
    # but change its config per route call (or use tiered escalation in PathFinder).
    # Since AStarRouter.route now accepts bend_collision_type, we can do it directly.
    
    router = AStarRouter(evaluator)
    pf = PathFinder(router, evaluator)

    netlist = {
        "model_arc": (Port(10, 130, 0), Port(130, 100, -90)),
        "model_bbox": (Port(10, 80, 0), Port(130, 50, -90)),
        "model_clipped": (Port(10, 30, 0), Port(130, 10, -90)),
    }
    net_widths = {nid: 2.0 for nid in netlist}

    # Manual routing to specify different models per net
    results = {}
    
    print("Routing with 'arc' model...")
    results["model_arc"] = pf.router.route(netlist["model_arc"][0], netlist["model_arc"][1], 2.0, 
                                          net_id="model_arc", bend_collision_type="arc")
    
    print("Routing with 'bbox' model...")
    results["model_bbox"] = pf.router.route(netlist["model_bbox"][0], netlist["model_bbox"][1], 2.0, 
                                           net_id="model_bbox", bend_collision_type="bbox")
    
    print("Routing with 'clipped_bbox' model...")
    results["model_clipped"] = pf.router.route(netlist["model_clipped"][0], netlist["model_clipped"][1], 2.0, 
                                              net_id="model_clipped", bend_collision_type="clipped_bbox")

    # Wrap in RoutingResult for visualization
    from inire.router.pathfinder import RoutingResult
    final_results = {
        nid: RoutingResult(nid, path if path else [], path is not None, 0)
        for nid, path in results.items()
    }

    fig, ax = plot_routing_results(final_results, [obstacle], bounds, netlist=netlist)
    fig.savefig("examples/08_custom_bend_geometry.png")
    print("Saved plot to examples/08_custom_bend_geometry.png")

if __name__ == "__main__":
    main()