from shapely.geometry import Polygon

from inire import CongestionOptions, NetSpec, ObjectiveWeights, RoutingOptions, RoutingProblem, RoutingResult, SearchOptions, route
from inire.geometry.primitives import Port
from inire.utils.visualization import plot_routing_results


def _route_scenario(
    bounds: tuple[float, float, float, float],
    obstacles: list[Polygon],
    bend_collision_type: str,
    netlist: dict[str, tuple[Port, Port]],
    widths: dict[str, float],
    *,
    bend_clip_margin: float | None = None,
) -> dict[str, RoutingResult]:
    problem = RoutingProblem(
        bounds=bounds,
        nets=tuple(NetSpec(net_id, start, target, width=widths[net_id]) for net_id, (start, target) in netlist.items()),
        static_obstacles=tuple(obstacles),
    )
    options = RoutingOptions(
        search=SearchOptions(
            bend_radii=(10.0,),
            bend_collision_type=bend_collision_type,
            bend_clip_margin=bend_clip_margin,
        ),
        objective=ObjectiveWeights(
            bend_penalty=50.0,
            sbend_penalty=150.0,
        ),
        congestion=CongestionOptions(use_tiered_strategy=False),
    )
    return route(problem, options=options).results_by_net


def main() -> None:
    print("Running Example 06: Bend Collision Models...")

    bounds = (-20, -20, 170, 170)
    obs_arc = Polygon([(40, 110), (60, 110), (60, 130), (40, 130)])
    obs_bbox = Polygon([(40, 60), (60, 60), (60, 80), (40, 80)])
    obs_clipped = Polygon([(40, 10), (60, 10), (60, 30), (40, 30)])

    obstacles = [obs_arc, obs_bbox, obs_clipped]
    netlist_arc = {"arc_model": (Port(10, 120, 0), Port(90, 140, 90))}
    netlist_bbox = {"bbox_model": (Port(10, 70, 0), Port(90, 90, 90))}
    netlist_clipped = {"clipped_model": (Port(10, 20, 0), Port(90, 40, 90))}

    print("Routing Scenario 1 (Arc)...")
    res_arc = _route_scenario(bounds, obstacles, "arc", netlist_arc, {"arc_model": 2.0})
    print("Routing Scenario 2 (BBox)...")
    res_bbox = _route_scenario(bounds, obstacles, "bbox", netlist_bbox, {"bbox_model": 2.0})
    print("Routing Scenario 3 (Clipped BBox)...")
    res_clipped = _route_scenario(
        bounds,
        obstacles,
        "clipped_bbox",
        netlist_clipped,
        {"clipped_model": 2.0},
        bend_clip_margin=1.0,
    )

    all_results = {**res_arc, **res_bbox, **res_clipped}
    all_netlists = {**netlist_arc, **netlist_bbox, **netlist_clipped}

    fig, _ax = plot_routing_results(all_results, obstacles, bounds, netlist=all_netlists)
    fig.savefig("examples/06_bend_collision_models.png")
    print("Saved plot to examples/06_bend_collision_models.png")


if __name__ == "__main__":
    main()