import math
from dataclasses import dataclass

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFilter


RESOLUTION_M = 0.10
MAP_W, MAP_H = 920, 520


@dataclass
class WorldBox:
    x_min: float
    x_max: float
    y_min: float
    y_max: float


def build_demo_map() -> np.ndarray:
    """0 free, 1 obstacle, -1 unknown."""
    grid = np.zeros((MAP_H, MAP_W), dtype=np.int8)
    grid[:, :] = 0

    # Unknown bands and off-workspace corners.
    grid[:38, :] = -1
    grid[-35:, :] = -1
    grid[:, :42] = -1
    grid[:, -45:] = -1
    grid[390:480, 650:880] = -1

    # Walls, shelves and typical AMR obstacles.
    obstacles = [
        (115, 95, 225, 145),
        (315, 80, 365, 260),
        (510, 70, 555, 185),
        (675, 120, 805, 170),
        (190, 295, 310, 350),
        (420, 310, 520, 360),
        (610, 265, 720, 315),
        (765, 250, 815, 375),
    ]
    for x0, y0, x1, y1 in obstacles:
        grid[y0:y1, x0:x1] = 1

    # Thin border obstacles.
    grid[70:75, 80:840] = 1
    grid[440:445, 80:620] = 1
    grid[70:445, 80:85] = 1
    grid[70:390, 840:845] = 1
    return grid


def world_to_px(x_m: float, y_m: float) -> tuple[int, int]:
    x = int(round(x_m / RESOLUTION_M)) + 60
    y = MAP_H - (int(round(y_m / RESOLUTION_M)) + 45)
    return x, y


def px_to_world(x: int, y: int) -> tuple[float, float]:
    wx = (x - 60) * RESOLUTION_M
    wy = (MAP_H - y - 45) * RESOLUTION_M
    return wx, wy


def clearance_free_mask(grid: np.ndarray, clearance_px: int) -> np.ndarray:
    unsafe = grid != 0
    unsafe_image = Image.fromarray((unsafe.astype(np.uint8) * 255), "L")
    kernel = max(3, clearance_px * 2 + 1)
    if kernel % 2 == 0:
        kernel += 1
    inflated = unsafe_image.filter(ImageFilter.MaxFilter(kernel))
    return np.asarray(inflated) == 0


def contiguous_segments(row_mask: np.ndarray, min_len: int = 24) -> list[tuple[int, int]]:
    segments = []
    start = None
    for i, ok in enumerate(row_mask):
        if ok and start is None:
            start = i
        if (not ok or i == len(row_mask) - 1) and start is not None:
            end = i if ok and i == len(row_mask) - 1 else i - 1
            if end - start + 1 >= min_len:
                segments.append((start, end))
            start = None
    return segments


def plan_coverage(
    x_min: float,
    x_max: float,
    y_min: float,
    y_max: float,
    lane_spacing_m: float,
    robot_width_m: float,
    robot_length_m: float,
    safety_margin_m: float,
    map_aware: bool,
) -> tuple[Image.Image, str]:
    grid = build_demo_map()
    clearance_m = math.hypot(robot_width_m, robot_length_m) / 2.0 + safety_margin_m
    clearance_px = max(1, int(round(clearance_m / RESOLUTION_M)))
    free_mask = clearance_free_mask(grid, clearance_px) if map_aware else grid != 1

    box = WorldBox(min(x_min, x_max), max(x_min, x_max), min(y_min, y_max), max(y_min, y_max))
    px0, py1 = world_to_px(box.x_min, box.y_min)
    px1, py0 = world_to_px(box.x_max, box.y_max)
    x0, x1 = sorted((max(0, px0), min(MAP_W - 1, px1)))
    y0, y1 = sorted((max(0, py0), min(MAP_H - 1, py1)))

    lane_step = max(10, int(round(lane_spacing_m / RESOLUTION_M)))
    waypoints: list[tuple[int, int]] = []
    path_lines: list[list[tuple[int, int]]] = []
    previous_end: tuple[int, int] | None = None
    direction = 1
    sampled_cells = 0

    for y in range(y0, y1 + 1, lane_step):
        row = free_mask[y, x0 : x1 + 1]
        for seg_start, seg_end in contiguous_segments(row):
            sx = x0 + seg_start
            ex = x0 + seg_end
            sampled_cells += max(0, ex - sx + 1)
            if direction > 0:
                start, end = (sx, y), (ex, y)
            else:
                start, end = (ex, y), (sx, y)
            if previous_end is not None:
                connector = safe_connector(previous_end, start, free_mask)
                if connector:
                    path_lines.append(connector)
            path_lines.append([start, end])
            waypoints.extend([start, end])
            previous_end = end
            direction *= -1

    image = render_map(grid, free_mask, (x0, y0, x1, y1), waypoints, path_lines, clearance_px)
    covered_area = sampled_cells * lane_spacing_m * RESOLUTION_M
    path_len_m = estimate_path_length(path_lines)
    note = (
        f"生成 {len(waypoints)} 个 waypoint，估算覆盖面积 {covered_area:.1f} m²，"
        f"路径长度约 {path_len_m:.1f} m。\n\n"
        f"车体 footprint = {robot_width_m:.2f} m x {robot_length_m:.2f} m，"
        f"安全边距 = {safety_margin_m:.2f} m，中心点 clearance 约 {clearance_m:.2f} m。\n\n"
        "工程含义：先把用户选择区域裁剪到已知 free-space，再生成往返覆盖线；"
        "unknown、地图外区域和按车体膨胀后的障碍区域不会参与规划。真实机器人执行时仍应走 Nav2 "
        "与 collision monitor 等安全链路。"
    )
    return image, note


def estimate_path_length(path_lines: list[list[tuple[int, int]]]) -> float:
    total_px = 0.0
    for line in path_lines:
        for a, b in zip(line, line[1:]):
            total_px += math.hypot(a[0] - b[0], a[1] - b[1])
    return total_px * RESOLUTION_M


def line_is_free(a: tuple[int, int], b: tuple[int, int], free_mask: np.ndarray) -> bool:
    dx = b[0] - a[0]
    dy = b[1] - a[1]
    steps = max(abs(dx), abs(dy), 1)
    for i in range(steps + 1):
        x = int(round(a[0] + dx * i / steps))
        y = int(round(a[1] + dy * i / steps))
        if x < 0 or x >= MAP_W or y < 0 or y >= MAP_H or not free_mask[y, x]:
            return False
    return True


def safe_connector(
    a: tuple[int, int],
    b: tuple[int, int],
    free_mask: np.ndarray,
) -> list[tuple[int, int]] | None:
    if line_is_free(a, b, free_mask):
        return [a, b]
    for corner in ((a[0], b[1]), (b[0], a[1])):
        if line_is_free(a, corner, free_mask) and line_is_free(corner, b, free_mask):
            return [a, corner, b]
    return None

def render_map(
    grid: np.ndarray,
    free_mask: np.ndarray,
    selected: tuple[int, int, int, int],
    waypoints: list[tuple[int, int]],
    path_lines: list[list[tuple[int, int]]],
    clearance_px: int,
) -> Image.Image:
    rgb = np.zeros((MAP_H, MAP_W, 3), dtype=np.uint8)
    rgb[grid == 0] = [244, 247, 250]
    rgb[grid == -1] = [198, 204, 213]
    rgb[grid == 1] = [31, 41, 55]
    rgb[(grid == 0) & free_mask] = [236, 252, 242]

    image = Image.fromarray(rgb, "RGB")
    draw = ImageDraw.Draw(image, "RGBA")
    x0, y0, x1, y1 = selected
    draw.rectangle((x0, y0, x1, y1), outline=(37, 99, 235, 240), width=4, fill=(37, 99, 235, 28))

    for line in path_lines:
        if len(line) >= 2:
            draw.line(line, fill=(18, 142, 76, 255), width=4, joint="curve")

    if waypoints:
        for idx, (x, y) in enumerate(waypoints):
            r = 4 if idx not in (0, len(waypoints) - 1) else 6
            color = (22, 101, 52, 255) if idx == 0 else (21, 128, 61, 255)
            if idx == len(waypoints) - 1:
                color = (220, 38, 38, 255)
            draw.ellipse((x - r, y - r, x + r, y + r), fill=color)

    legend_x, legend_y = 20, 20
    draw.rounded_rectangle((legend_x, legend_y, legend_x + 360, legend_y + 116), radius=8, fill=(255, 255, 255, 232))
    draw.text((legend_x + 14, legend_y + 12), "RabbitRobot coverage planning demo", fill=(15, 23, 42))
    legend = [
        ("free / safe after clearance", (236, 252, 242)),
        ("unknown or off-map", (198, 204, 213)),
        ("inflated obstacle source", (31, 41, 55)),
        (f"path with clearance radius {clearance_px}px", (18, 142, 76)),
    ]
    for i, (text, color) in enumerate(legend):
        yy = legend_y + 36 + i * 18
        draw.rectangle((legend_x + 16, yy + 3, legend_x + 30, yy + 14), fill=color + (255,))
        draw.text((legend_x + 38, yy), text, fill=(51, 65, 85))
    return image


with gr.Blocks(title="RabbitRobot Coverage Path Planning") as demo:
    gr.Markdown(
        """
        # RabbitRobot Coverage Path Planning

        一个公开可交互的 AMR 覆盖路径规划演示：用合成占据栅格展示 map-aware clipping、
        车体 footprint 安全膨胀和往返式覆盖路径生成。
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            x_min = gr.Slider(-2.0, 65.0, value=2.0, step=0.5, label="区域 x_min (m)")
            x_max = gr.Slider(2.0, 78.0, value=58.0, step=0.5, label="区域 x_max (m)")
            y_min = gr.Slider(-1.0, 42.0, value=5.5, step=0.5, label="区域 y_min (m)")
            y_max = gr.Slider(3.0, 46.0, value=34.0, step=0.5, label="区域 y_max (m)")
            lane_spacing = gr.Slider(0.35, 2.0, value=0.80, step=0.05, label="覆盖行距 / tool width (m)")
            robot_width = gr.Slider(0.25, 1.0, value=0.486, step=0.01, label="车体宽度 (m)")
            robot_length = gr.Slider(0.25, 1.2, value=0.450, step=0.01, label="车体长度 (m)")
            safety_margin = gr.Slider(0.0, 0.8, value=0.08, step=0.01, label="安全边距 (m)")
            map_aware = gr.Checkbox(value=True, label="启用 map-aware free-space clipping")
            run = gr.Button("Generate coverage path", variant="primary")
        with gr.Column(scale=2):
            output_image = gr.Image(label="Coverage path overlay", type="pil")
            output_note = gr.Textbox(label="Planning summary", lines=7)

    inputs = [x_min, x_max, y_min, y_max, lane_spacing, robot_width, robot_length, safety_margin, map_aware]
    run.click(plan_coverage, inputs=inputs, outputs=[output_image, output_note])
    demo.load(plan_coverage, inputs=inputs, outputs=[output_image, output_note])


if __name__ == "__main__":
    demo.launch()