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...")

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

    evaluator = CostEvaluator(engine, danger_map, bend_penalty=50.0, sbend_penalty=150.0)
    router = AStarRouter(evaluator, snap_size=1.0, straight_lengths=[1.0, 5.0, 25.0], bend_radii=[10.0])
    pf = PathFinder(router, evaluator)

    # 2. Define Netlist
    netlist = {
        "custom_bend": (Port(20, 20, 0), Port(100, 100, 90)),
    }
    net_widths = {"custom_bend": 2.0}

    # 3. Route with standard arc first
    print("Routing with standard arc...")
    results_std = pf.route_all(netlist, net_widths)

    # 4. Define a custom 'trapezoid' bend model
    # (Just for demonstration - we override the collision model during search)
    custom_poly = Polygon([(0, 0), (20, 0), (20, 20), (0, 20)]) # Oversized box
    
    print("Routing with custom collision model...")
    # Override bend_collision_type with a literal Polygon
    router_custom = AStarRouter(evaluator, snap_size=1.0, straight_lengths=[1.0, 5.0, 25.0], bend_radii=[10.0], bend_collision_type=custom_poly)
    results_custom = PathFinder(router_custom, evaluator, use_tiered_strategy=False).route_all(
        {"custom_model": netlist["custom_bend"]}, {"custom_model": 2.0}
    )

    # 5. Visualize
    all_results = {**results_std, **results_custom}
    fig, ax = plot_routing_results(all_results, [], 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()