"""Grid A*, conservative footprint inflation and ideal differential-drive tracking.

Standard library only. This is an original numerical experiment, not Nav2 or SLAM.
"""
import heapq
import json
import math

WIDTH, HEIGHT, RESOLUTION, DT = 26, 20, 0.25, 0.02


def world():
    walls = {(x, y) for x in range(WIDTH) for y in range(HEIGHT)
             if x in (0, WIDTH - 1) or y in (0, HEIGHT - 1)}
    walls |= {(12, y) for y in range(3, 16) if y not in (8, 9, 10, 11)}
    return walls


def inflated(obstacles):
    # A full one-cell Chebyshev margin is conservative for a 0.12 m disk.
    return {(x + dx, y + dy) for x, y in obstacles
            for dx in (-1, 0, 1) for dy in (-1, 0, 1)}


def astar(start, goal, occupied):
    queue, previous, distance = [(0, start)], {}, {start: 0}
    while queue:
        _, current = heapq.heappop(queue)
        if current == goal:
            path = [current]
            while current in previous:
                current = previous[current]
                path.append(current)
            return path[::-1]
        for dx, dy in ((1, 0), (0, 1), (-1, 0), (0, -1)):
            nxt = current[0] + dx, current[1] + dy
            if not (0 <= nxt[0] < WIDTH and 0 <= nxt[1] < HEIGHT) or nxt in occupied:
                continue
            candidate = distance[current] + 1
            if candidate < distance.get(nxt, float("inf")):
                distance[nxt], previous[nxt] = candidate, current
                heuristic = abs(nxt[0] - goal[0]) + abs(nxt[1] - goal[1])
                heapq.heappush(queue, (candidate + heuristic, nxt))
    raise ValueError("no route")


def center(cell):
    return ((cell[0] + 0.5) * RESOLUTION, (cell[1] + 0.5) * RESOLUTION)


def track(path, occupied):
    x, y = center(path[0])
    heading, target_index, traveled = 0.0, 1, 0.0
    for step in range(50000):
        tx, ty = center(path[target_index])
        distance = math.hypot(tx - x, ty - y)
        if distance < 0.025:
            target_index += 1
            if target_index == len(path):
                gx, gy = center(path[-1])
                return {"steps": step, "simulated_seconds": round(step * DT, 2),
                        "distance_m": traveled, "goal_error_m": math.hypot(gx - x, gy - y),
                        "collision_checks_passed": True}
            continue
        desired = math.atan2(ty - y, tx - x)
        error = math.atan2(math.sin(desired - heading), math.cos(desired - heading))
        omega = max(-2.0, min(2.0, 4.0 * error))
        speed = min(0.45, 2.5 * distance) * max(0.0, math.cos(error))
        nx = x + speed * math.cos(heading + 0.5 * omega * DT) * DT
        ny = y + speed * math.sin(heading + 0.5 * omega * DT) * DT
        cell = math.floor(nx / RESOLUTION), math.floor(ny / RESOLUTION)
        if cell in occupied:
            raise AssertionError(f"footprint margin entered at {cell}")
        traveled += math.hypot(nx - x, ny - y)
        x, y, heading = nx, ny, heading + omega * DT
    raise AssertionError("tracking timeout")


def main():
    start, goal = (3, 3), (22, 16)
    walls = world()
    occupancy = inflated(walls)
    path = astar(start, goal, occupancy)
    report = {"grid_resolution_m": RESOLUTION, "robot_radius_m": 0.12,
              "static_path_edges": len(path) - 1, "static": track(path, occupancy)}
    # An obstacle becomes known after reaching this discrete state. Replanning
    # uses an updated map; there is no simulated obstacle detector.
    restart = path[5]
    obstacles = walls | {(12, 9), (12, 10)}
    changed = inflated(obstacles)
    alternative = astar(restart, goal, changed)
    report["replanned_path_edges"] = len(alternative) - 1
    report["replanned"] = track(alternative, changed)
    report["new_obstacle_cells"] = [[12, 9], [12, 10]]
    assert set(alternative).isdisjoint(changed)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()