from __future__ import annotations

from typing import TYPE_CHECKING
import matplotlib.pyplot as plt
import numpy
from shapely.geometry import MultiPolygon, Polygon

if TYPE_CHECKING:
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure

    from inire.geometry.primitives import Port
    from inire.router.pathfinder import RoutingResult


def plot_routing_results(
    results: dict[str, RoutingResult],
    static_obstacles: list[Polygon],
    bounds: tuple[float, float, float, float],
    netlist: dict[str, tuple[Port, Port]] | None = None,
    show_actual: bool = True,
) -> tuple[Figure, Axes]:
    """
    Plot obstacles and routed paths using matplotlib.

    Args:
        results: Dictionary of net_id to RoutingResult.
        static_obstacles: List of static obstacle polygons.
        bounds: Plot limits (minx, miny, maxx, maxy).
        netlist: Optional original netlist for port visualization.
        show_actual: If True, overlay high-fidelity geometry if available.

    Returns:
        The matplotlib Figure and Axes objects.
    """
    fig, ax = plt.subplots(figsize=(12, 12))

    # Plot static obstacles (gray)
    for poly in static_obstacles:
        x, y = poly.exterior.xy
        ax.fill(x, y, alpha=0.3, fc="gray", ec="black", zorder=1)

    # Plot paths
    colors = plt.get_cmap("tab20")
    for i, (net_id, res) in enumerate(results.items()):
        color: str | tuple[float, ...] = colors(i % 20)
        if not res.is_valid:
            color = "red"

        label_added = False
        for comp in res.path:
            # 1. Plot Collision Geometry (Translucent fill)
            # This is the geometry used during search (e.g. proxy or arc)
            for poly in comp.geometry:
                if isinstance(poly, MultiPolygon):
                    geoms = list(poly.geoms)
                else:
                    geoms = [poly]

                for g in geoms:
                    if hasattr(g, "exterior"):
                        x, y = g.exterior.xy
                        ax.fill(x, y, alpha=0.15, fc=color, ec=color, linestyle='--', lw=0.5, zorder=2)
                    else:
                        # Fallback for LineString or other geometries
                        x, y = g.xy
                        ax.plot(x, y, color=color, alpha=0.15, linestyle='--', lw=0.5, zorder=2)

            # 2. Plot "Actual" Geometry (The high-fidelity shape used for fabrication)
            # Use comp.actual_geometry if it exists (should be the arc)
            actual_geoms_to_plot = comp.actual_geometry if comp.actual_geometry is not None else comp.geometry
            
            for poly in actual_geoms_to_plot:
                if isinstance(poly, MultiPolygon):
                    geoms = list(poly.geoms)
                else:
                    geoms = [poly]
                for g in geoms:
                    if hasattr(g, "exterior"):
                        x, y = g.exterior.xy
                        ax.plot(x, y, color=color, lw=1.5, alpha=0.9, zorder=3, label=net_id if not label_added else "")
                    else:
                        x, y = g.xy
                        ax.plot(x, y, color=color, lw=1.5, alpha=0.9, zorder=3, label=net_id if not label_added else "")
                    label_added = True

            # 3. Plot subtle port orientation arrow
            p = comp.end_port
            rad = numpy.radians(p.orientation)
            ax.quiver(p.x, p.y, numpy.cos(rad), numpy.sin(rad), color="black", 
                      scale=40, width=0.002, alpha=0.2, pivot="tail", zorder=4)

        if not res.path and not res.is_valid:
            # Best-effort display: If the path is empty but failed, it might be unroutable.
            # We don't have a partial path in RoutingResult currently.
            pass

    # 4. Plot main arrows for netlist ports
    if netlist:
        for net_id, (start_p, target_p) in netlist.items():
            for p in [start_p, target_p]:
                rad = numpy.radians(p.orientation)
                ax.quiver(p.x, p.y, numpy.cos(rad), numpy.sin(rad), color="black", 
                          scale=25, width=0.004, pivot="tail", zorder=6)

    ax.set_xlim(bounds[0], bounds[2])
    ax.set_ylim(bounds[1], bounds[3])
    ax.set_aspect("equal")
    ax.set_title("Inire Routing Results (Dashed: Collision Proxy, Solid: Actual Geometry)")
    
    # Legend handling for many nets
    if len(results) < 25:
        handles, labels = ax.get_legend_handles_labels()
        if labels:
            ax.legend(loc='upper left', bbox_to_anchor=(1, 1), fontsize='small', ncol=2)
            fig.tight_layout()

    plt.grid(True, which='both', linestyle=':', alpha=0.5)
    return fig, ax
            
def plot_danger_map(
    danger_map: DangerMap,
    ax: Axes | None = None,
) -> tuple[Figure, Axes]:
    """
    Plot the pre-computed danger map as a heatmap.
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10))
    else:
        fig = ax.get_figure()

    # Need to transpose because grid is [x, y] and imshow expects [row, col] (y, x)
    # Also origin='lower' to match coordinates
    im = ax.imshow(
        danger_map.grid.T,
        origin='lower',
        extent=[danger_map.minx, danger_map.maxx, danger_map.miny, danger_map.maxy],
        cmap='YlOrRd',
        alpha=0.6
    )
    plt.colorbar(im, ax=ax, label='Danger Cost')
    ax.set_title("Danger Map (Proximity Costs)")
    return fig, ax

def plot_expanded_nodes(
    nodes: list[tuple[float, float, float]],
    ax: Axes | None = None,
    color: str = 'gray',
    alpha: float = 0.3,
) -> tuple[Figure, Axes]:
    """
    Plot A* expanded nodes for debugging.
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 10))
    else:
        fig = ax.get_figure()

    if not nodes:
        return fig, ax

    x, y, _ = zip(*nodes)
    ax.scatter(x, y, s=1, c=color, alpha=alpha, zorder=0)
    return fig, ax

def plot_expansion_density(
    nodes: list[tuple[float, float, float]],
    bounds: tuple[float, float, float, float],
    ax: Axes | None = None,
    bins: int | tuple[int, int] = 50,
) -> tuple[Figure, Axes]:
    """
    Plot a density heatmap (2D histogram) of expanded nodes.

    Args:
        nodes: List of (x, y, orientation) tuples.
        bounds: (minx, miny, maxx, maxy) for the plot range.
        ax: Optional existing axes to plot on.
        bins: Number of bins for the histogram (int or (nx, ny)).

    Returns:
        Figure and Axes objects.
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=(12, 12))
    else:
        fig = ax.get_figure()

    if not nodes:
        ax.text(0.5, 0.5, "No Expansion Data", ha='center', va='center', transform=ax.transAxes)
        return fig, ax

    x, y, _ = zip(*nodes)
    
    # Create 2D histogram
    h, xedges, yedges = numpy.histogram2d(
        x, y, 
        bins=bins, 
        range=[[bounds[0], bounds[2]], [bounds[1], bounds[3]]]
    )
    
    # Plot as image
    im = ax.imshow(
        h.T, 
        origin='lower', 
        extent=[bounds[0], bounds[2], bounds[1], bounds[3]],
        cmap='plasma',
        interpolation='nearest',
        alpha=0.7
    )
    
    plt.colorbar(im, ax=ax, label='Expansion Count')
    ax.set_title("Search Expansion Density")
    ax.set_xlim(bounds[0], bounds[2])
    ax.set_ylim(bounds[1], bounds[3])
    
    return fig, ax
