from __future__ import annotations

from dataclasses import dataclass, field
from importlib import import_module
from typing import Any, Sequence


@dataclass(frozen=True)
class TopDownSceneElements:
    room_nodes: Sequence[Any]
    object_nodes: Sequence[Any] = field(default_factory=tuple)
    edges: Sequence[tuple[str, str]] = field(default_factory=tuple)


@dataclass(frozen=True)
class TopDownRenderStyle:
    room_color: str = "red"
    room_label_color: str = "yellow"
    room_marker_size: int = 50
    room_label_offset: int = 2
    figure_size: tuple[int, int] = (8, 8)
    map_cmap: str = "gray"
    title: str = "RoomNode Top-Down Map"


def render_topdown_scene_map(
    pathfinder: Any,
    elements: TopDownSceneElements,
    meters_per_pixel: float,
    style: TopDownRenderStyle | None = None,
    maps_module: Any | None = None,
    plt_module: Any | None = None,
    use_matplotlib: bool = True,
) -> Any:
    if not elements.room_nodes:
        raise ValueError("room_nodes must not be empty")

    if meters_per_pixel <= 0:
        raise ValueError("meters_per_pixel must be greater than 0")

    if elements.object_nodes:
        raise NotImplementedError("object_nodes overlay is not implemented yet")

    if elements.edges:
        raise NotImplementedError("edge overlay is not implemented yet")

    if style is None:
        style = TopDownRenderStyle()

    if maps_module is None:
        maps_module = import_module("habitat.utils.visualizations.maps")

    map_height = float(elements.room_nodes[0].center[1])
    top_down_map = maps_module.get_topdown_map(
        pathfinder,
        height=map_height,
        meters_per_pixel=meters_per_pixel,
    )

    if not use_matplotlib:
        return top_down_map

    if plt_module is None:
        try:
            plt_module = import_module("matplotlib.pyplot")
        except ImportError:
            return top_down_map

    plt_module.figure(figsize=style.figure_size)
    plt_module.imshow(top_down_map, cmap=style.map_cmap)

    for room_node in elements.room_nodes:
        grid_y, grid_x = maps_module.to_grid(
            float(room_node.center[2]),
            float(room_node.center[0]),
            top_down_map.shape,
            pathfinder=pathfinder,
        )
        plt_module.scatter(grid_x, grid_y, c=style.room_color, s=style.room_marker_size)
        plt_module.text(
            grid_x + style.room_label_offset,
            grid_y + style.room_label_offset,
            room_node.room_id,
            color=style.room_label_color,
            fontsize=8,
        )

    plt_module.title(style.title)
    plt_module.axis("off")
    plt_module.show()
    return top_down_map