from __future__ import annotations

from typing import Dict, Iterable, Mapping, Optional, Tuple, Union
import math

import matplotlib.pyplot as plt
import matplotlib as mpl


Coord = Tuple[float, float]
ColorMapLike = Mapping[int, str]


def index_to_xy(
    j: int,
    rows: int,
    cols: int,
    origin_top_left: bool = True,
) -> Tuple[int, int]:
    """
    Convert 1-based cell index j to integer grid coordinates (x, y).

    Numbering is row-major. With origin_top_left=True, index 1 is top-left.
    Otherwise, index 1 is bottom-left.

    Returns (x, y) with x, y in {1, 2, ..., cols/rows}.
    """
    if not (1 <= j <= rows * cols):
        raise ValueError(f"Index {j} is out of range 1..{rows * cols}.")

    x = (j - 1) % cols + 1
    row_idx = (j - 1) // cols  # 0-based row
    if origin_top_left:
        y = rows - row_idx
    else:
        y = row_idx + 1
    return x, y


def xy_to_index(
    x: int,
    y: int,
    rows: int,
    cols: int,
    origin_top_left: bool = True,
) -> int:
    """
    Convert integer grid coordinates (x, y) to 1-based cell index j.

    Inverse of index_to_xy for valid inputs.
    """
    if not (1 <= x <= cols and 1 <= y <= rows):
        raise ValueError(
            f"(x, y)=({x}, {y}) out of bounds 1..{cols}, 1..{rows}."
        )

    if origin_top_left:
        row_idx = rows - y
    else:
        row_idx = y - 1
    j = row_idx * cols + x
    return j


def plot_warehouse_layout(
    rows: int,
    cols: int,
    io_positions: Mapping[int, Coord],
    labels: Optional[Mapping[int, Union[str, int]]] = None,
    colors: Optional[Union[str, ColorMapLike]] = None,
    *,
    origin_top_left: bool = True,
    figsize: Tuple[float, float] = (6.5, 6.5),
    cell_edgecolor: str = "black",
    cell_linewidth: float = 1.0,
    annotate_cell_index: bool = True,
    annotate_user_label: bool = True,
    index_style: Dict[str, Union[str, float, int]] = None,
    label_style: Dict[str, Union[str, float, int]] = None,
    io_style: Dict[str, Union[str, float, int]] = None,
    show_axes: bool = True,
    show_grid: bool = True,
    minor_grid: bool = True,
    show: bool = True,
    title: Optional[str] = None
):
    """
    Plot a warehouse grid and IO points.

    Args:
      rows: Number of grid rows R.
      cols: Number of grid columns C.
      io_positions: Dict mapping IO id -> (x_float, y_float) coordinates.
      labels: Optional dict mapping cell index -> user label (e.g., product code).
      colors: Either a single color for all cells, or a dict {cell index -> color}.
      origin_top_left: If True, index 1 is top-left. Else index 1 is bottom-left.
      figsize: Figure size in inches.
      cell_edgecolor: Edge color for each cell rectangle.
      cell_linewidth: Edge line width for each cell rectangle.
      annotate_cell_index: If True, draws the cell index j near the top of cell.
      annotate_user_label: If True and labels provided, draws the label centered.
      index_style: Matplotlib text kwargs for cell index annotations.
      label_style: Matplotlib text kwargs for user label annotations.
      io_style: Matplotlib kwargs for IO markers and text style:
                keys: radius, facecolor, edgecolor, text_color, text_size.
      show_axes: If False, hides axes spines and ticks.
      show_grid: If True, shows major grid at integer coordinates.
      minor_grid: If True, shows minor grid at 0.5 spacing.
      show: If True, calls plt.show().
      title: Optional title for the plot.

    Returns:
      (fig, ax) so you can further customize outside the function.
    """
    # ------------------------
    # 0) Basic input checking
    # ------------------------
    if rows <= 0 or cols <= 0:
        raise ValueError("rows and cols must be positive integers.")

    max_index = rows * cols

    if labels is not None:
        bad = [j for j in labels.keys() if not (1 <= j <= max_index)]
        if bad:
            raise ValueError(f"labels has out-of-range cell indices: {bad}")

    color_map: Optional[ColorMapLike] = None
    if isinstance(colors, dict):
        bad = [j for j in colors.keys() if not (1 <= j <= max_index)]
        if bad:
            raise ValueError(f"colors has out-of-range cell indices: {bad}")
        color_map = colors
    elif isinstance(colors, str) or colors is None:
        # Single color or no color
        pass
    else:
        raise TypeError("colors must be a str, a dict, or None.")

    # Defaults for text styles
    index_style = index_style or dict(
        ha="center", va="top", fontsize=8, color="black"
    )
    label_style = label_style or dict(
        ha="center", va="center", fontsize=12, color="black", weight="bold"
    )
    io_style = io_style or {}
    io_radius = float(io_style.get("radius", 0.18))
    io_face = str(io_style.get("facecolor", "black"))
    io_edge = str(io_style.get("edgecolor", "none"))
    io_text_color = str(io_style.get("text_color", "white"))
    io_text_size = float(io_style.get("text_size", 9))

    # -----------------------------
    # 1) Prepare figure and axes
    # -----------------------------
    fig, ax = plt.subplots(figsize=figsize)
    ax.set_aspect("equal")

    # Set bounds with 0.5 margin so cells centered at integers fit nicely
    ax.set_xlim(0.0, cols + 1.0)
    ax.set_ylim(0.0, rows + 1.0)

    # Ticks at integer coordinates
    ax.set_xticks(range(1, cols + 1))
    ax.set_yticks(range(1, rows + 1))

    # Optional minor grid at 0.5 increments
    if minor_grid:
        ax.set_xticks([i + 0.5 for i in range(0, cols)], minor=True)
        ax.set_yticks([i + 0.5 for i in range(0, rows)], minor=True)

    if show_grid:
        ax.grid(which="major", linewidth=0.8, alpha=0.35)
        if minor_grid:
            ax.grid(which="minor", linewidth=0.5, alpha=0.15)

    # -----------------------------
    # 2) Draw all cells
    # -----------------------------
    for j in range(1, max_index + 1):
        x, y = index_to_xy(j, rows, cols, origin_top_left=origin_top_left)

        # Resolve facecolor
        if color_map is not None:
            facecolor = color_map.get(j, "none")
        else:
            facecolor = colors if isinstance(colors, str) else "none"

        rect = mpl.patches.Rectangle(
            (x - 0.5, y - 0.5),
            width=1.0,
            height=1.0,
            fill=True,
            facecolor=facecolor,
            edgecolor=cell_edgecolor,
            linewidth=cell_linewidth,
        )
        ax.add_patch(rect)

        # Cell index near the top edge
        if annotate_cell_index:
            ax.text(x, y + 0.40, str(j), **index_style)

        # User label in the center
        if annotate_user_label and labels and j in labels:
            ax.text(x, y, str(labels[j]), **label_style)

    # -----------------------------
    # 3) Plot IO points
    # -----------------------------
    for io_id, (xk, yk) in io_positions.items():
        # Marker
        circ = mpl.patches.Circle(
            (xk, yk), radius=io_radius, facecolor=io_face, edgecolor=io_edge
        )
        ax.add_patch(circ)
        # Text overlay
        ax.text(
            xk,
            yk,
            f"IO{io_id}",
            ha="center",
            va="center",
            fontsize=io_text_size,
            color=io_text_color,
        )

    # -----------------------------
    # 4) Axes cosmetics
    # -----------------------------
    ax.set_xlabel("x")
    ax.set_ylabel("y")

    if not show_axes:
        for side in ("top", "right", "left", "bottom"):
            ax.spines[side].set_visible(False)
        ax.set_xticks([])
        ax.set_yticks([])

    # Set title
    if title:
        ax.set_title(title)

    # Compact layout
    fig.tight_layout()

    if show:
        plt.show()

    return fig, ax

from typing import Iterable, Literal, Sequence, Tuple
import math

RGB = Tuple[int, int, int]
Point = Tuple[float, float]
Cell = Tuple[int, int]

def tile_color_by_io_distance(
    base_rgb: RGB,
    rows: int,
    cols: int,
    io_points: Sequence[Point],
    cell: Cell,
    metric: Literal["manhattan", "euclidean"] = "manhattan",
    gamma: float = 1.0,
    max_distance_hint: float | None = None,
) -> RGB:
    """
    Compute a tile color that lightens with distance from the nearest I/O point.

    Parameters
    ----------
    base_rgb
        Base color as (R, G, B) integers in [0, 255]. This is the color at
        distance 0. Farther cells are lightened toward white.
    rows, cols
        Warehouse map size (number of grid rows and columns).
    io_points
        Iterable of I/O points as (x, y), where x is column and y is row.
        Coordinates may be integer or float. At least one point is required.
    cell
        Cell position as (row, col).
    metric
        Distance metric: "manhattan" uses $|x_1-x_2|+|y_1-y_2|$,
        "euclidean" uses $\\sqrt{(x_1-x_2)^2+(y_1-y_2)^2}$.
    gamma
        Contrast control. The lightening factor $t$ is raised to the power
        $\\gamma$. Values $\\gamma>1$ keep cells darker near I/O for longer.
        Values $0<\\gamma<1$ lighten more quickly.
    max_distance_hint
        Optional maximum nearest distance to use for normalization. If None,
        it is computed by scanning the whole grid and taking the maximum
        nearest distance to any I/O point.

    Returns
    -------
    RGB
        Lightened color as (R, G, B) in [0, 255].

    Notes
    -----
    The lightening factor is
    $$
      t = \\min\\left(1, \\left(\\frac{d_{\\text{near}}}{d_{\\max}}\\right)^{\\gamma}\\right),
    $$
    and the output color is blended toward white:
    $$
      c_{\\text{out}} = c_{\\text{base}} + t\\,(255 - c_{\\text{base}}).
    $$
    """
    if not io_points:
        raise ValueError("io_points must contain at least one point.")

    r, c = int(cell[0]), int(cell[1])

    def _dist(p: Point, q: Point) -> float:
        if metric == "manhattan":
            return abs(p[0] - q[0]) + abs(p[1] - q[1])
        elif metric == "euclidean":
            return math.hypot(p[0] - q[0], p[1] - q[1])
        else:
            raise ValueError("metric must be 'manhattan' or 'euclidean'.")

    # Distance from this cell to the nearest I/O point
    d_here = min(_dist((c, r), io) for io in io_points)

    # Normalizer: maximum nearest distance across the grid
    if max_distance_hint is None:
        max_near = 0.0
        # Scan every grid cell to find the max nearest distance
        for rr in range(rows):
            for cc in range(cols):
                d = min(_dist((cc, rr), io) for io in io_points)
                if d > max_near:
                    max_near = d
        d_max = max_near if max_near > 0 else 1.0
    else:
        d_max = max(1e-9, float(max_distance_hint))

    # Lightening factor in [0, 1]
    t = min(1.0, (d_here / d_max) ** max(gamma, 1e-9))

    # Blend base color toward white
    r_out = int(round(base_rgb[0] + t * (255 - base_rgb[0])))
    g_out = int(round(base_rgb[1] + t * (255 - base_rgb[1])))
    b_out = int(round(base_rgb[2] + t * (255 - base_rgb[2])))

    return (max(0, min(255, r_out)),
            max(0, min(255, g_out)),
            max(0, min(255, b_out)))


from typing import Dict, Tuple, Sequence

# --- paste your function here or import it ---
# from mymodule import tile_color_by_io_distance

RGB = Tuple[int, int, int]
Coord = Tuple[float, float]

def _rgb_to_hex(rgb: RGB) -> str:
    return "#{:02X}{:02X}{:02X}".format(*rgb)

def build_color_map_from_io(
    rows: int,
    cols: int,
    io_positions: Dict[int, Coord],
    base_rgb: RGB = (40, 40, 40),
    *,
    origin_top_left: bool = True,
    metric: str = "manhattan",
    gamma: float = 1.0,
) -> Dict[int, str]:
    """
    Produce {cell_index -> hex_color} for plot_warehouse_layout, where cells
    get lighter the farther they are from the nearest IO point.
    """
    # Convert IO coords from plotting space (x,y in 1..C, 1..R with possible .5)
    # to the distance function space (col,row) zero based.
    io_points: Sequence[Tuple[float, float]] = []
    for (x, y) in io_positions.values():
        c0 = x - 1.0
        r0 = (rows - y) if origin_top_left else (y - 1.0)
        io_points.append((c0, r0))

    # Precompute a good normalization hint: farthest cell from IOs
    # by scanning the grid once in the same zero based space.
    def nearest_dist_zero_based(r0: float, c0: float) -> float:
        if metric == "manhattan":
            return min(abs(c0 - cx) + abs(r0 - ry) for cx, ry in io_points)
        else:  # euclidean
            from math import hypot
            return min(hypot(c0 - cx, r0 - ry) for cx, ry in io_points)

    max_near = 0.0
    for rr in range(rows):
        for cc in range(cols):
            d = nearest_dist_zero_based(rr, cc)
            if d > max_near:
                max_near = d
    max_hint = max_near if max_near > 0 else 1.0

    # Build the index->color map
    colors: Dict[int, str] = {}
    for j in range(1, rows * cols + 1):
        # Your helper converts 1-based index to plotting coords
        x, y = index_to_xy(j, rows, cols, origin_top_left=origin_top_left)

        # Convert plotting coords to zero based row,col for the distance fn
        c0 = x - 1
        r0 = (rows - y) if origin_top_left else (y - 1)

        rgb = tile_color_by_io_distance(
            base_rgb=base_rgb,
            rows=rows,
            cols=cols,
            io_points=io_points,
            cell=(r0, c0),
            metric=metric,          # "manhattan" or "euclidean"
            gamma=gamma,            # contrast control
            max_distance_hint=max_hint,
        )
        colors[j] = _rgb_to_hex(rgb)

    return colors


# -----------------------------
# Example usage
# -----------------------------
if __name__ == "__main__":
    # io_positions = {
    #     1: (2.5, 0.5),
    #     2: (0.5, 2.5),
    #     3: (2.5, 4.5),
    # }

    # labels = {1: "A", 2: "B", 3: "C", 4: "D"}
    # colors = {1: "red", 2: "blue", 3: "green", 4: "yellow"}

    # plot_warehouse_layout(
    #     rows=4,
    #     cols=4,
    #     io_positions=io_positions,
    #     labels=labels,
    #     colors=colors,
    #     origin_top_left=True,
    #     figsize=(6, 6),
    #     annotate_cell_index=True,
    #     annotate_user_label=True,
    #     show_axes=True,
    #     show_grid=True,
    #     minor_grid=True,
    #     title="Optimal Warehouse Layout\nMin. travel cost: 42,420"
    # )

    io_positions = {
        1: (2.5, 0.5),
        2: (0.5, 2.5),
        3: (2.5, 4.5),
    }

    rows, cols = 10, 10
    colors = build_color_map_from_io(
        rows, cols, io_positions,
        base_rgb=(40, 40, 40),
        origin_top_left=True,
        metric="euclidean",
        gamma=1.2,
    )

    plot_warehouse_layout(
        rows=rows,
        cols=cols,
        io_positions=io_positions,
        colors=colors,
        title="Distance-shaded tiles",
    )