import numpy as np
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
from shapely.geometry import box

def main() -> None:
    print("Running Example 07: Fan-Out (5 Nets)...")

    # 1. Setup Environment
    # Small area for fast and reliable demonstration
    bounds = (0, 0, 100, 100)
    engine = CollisionEngine(clearance=2.0)
    
    # Wide bottleneck at x=50, 60um gap (from y=20 to y=80)
    obstacles = [
        box(50, 0, 55, 20),
        box(50, 80, 55, 100),
    ]
    for obs in obstacles:
        engine.add_static_obstacle(obs)

    danger_map = DangerMap(bounds=bounds)
    danger_map.precompute(obstacles)

    evaluator = CostEvaluator(engine, danger_map, greedy_h_weight=1.5)
    
    # Increase node_limit for more complex search
    router = AStarRouter(evaluator, node_limit=50000)
    pf = PathFinder(router, evaluator, max_iterations=2)

    # 2. Define Netlist: Fan-Out Configuration
    netlist = {}
    num_nets = 10
    start_x = 10
    # Bundle centered at y=50, 4um pitch
    start_y_base = 50 - (num_nets * 4.0) / 2.0
    
    end_x = 90
    end_y_base = 10
    end_y_pitch = 80.0 / (num_nets - 1)

    for i in range(num_nets):
        sy = start_y_base + i * 4.0
        ey = end_y_base + i * end_y_pitch
        
        net_id = f"net_{i:02d}"
        netlist[net_id] = (Port(start_x, sy, 0), Port(end_x, ey, 0))

    net_widths = {nid: 2.0 for nid in netlist}

    # 3. Route
    print(f"Routing {len(netlist)} nets through 60um bottleneck...")
    results = pf.route_all(netlist, net_widths)

    # 4. Check Results
    success_count = sum(1 for res in results.values() if res.is_valid)
    print(f"Routed {success_count}/{len(netlist)} nets successfully.")

    # 5. Visualize
    fig, ax = plot_routing_results(results, obstacles, bounds, netlist=netlist)
    fig.savefig("examples/07_large_scale_routing.png")
    print("Saved plot to examples/07_large_scale_routing.png")

if __name__ == "__main__":
    main()