from __future__ import annotations

import math
from typing import TYPE_CHECKING, Any, Callable, Literal

import arcade
import matplotlib
import matplotlib.colors
import networkx as nx

from mesarcade.utils import parse_color

if TYPE_CHECKING:
    import mesa

# Type alias for color values
Color = str | tuple[int, int, int] | tuple[int, int, int, int] | list[int] | None


class Artist:
    def __init__(
        self,
        get_xy_position: Callable[[Any], tuple[float, float]],
        color: Color = "black",
        color_attribute: str | Callable[[Any], Any] | None = None,
        color_map: str | dict[Any, Color] | None = "bwr",
        color_vmin: float | None = None,
        color_vmax: float | None = None,
        shape: Literal["rect", "circle"] = "rect",
        size: float = 1,
        filter_entities: Callable[[Any], bool] = lambda entity: True,
        jitter: bool = False,
        dynamic_color: bool = True,
        dynamic_position: bool = True,
        dynamic_population: bool = True,
        get_population: Callable[[mesa.Model], Any] = lambda model: None,
    ) -> None:
        self.color = parse_color(color=color)
        self.color_attribute = color_attribute
        self.color_map: str | dict | None = color_map
        self.color_vmin = color_vmin
        self.color_vmax = color_vmax
        self.shape = shape
        self.dynamic_color = dynamic_color
        self.dynamic_position = dynamic_position
        self.dynamic_population = dynamic_population
        self.filter_entities = filter_entities
        self.jitter = jitter
        self.size = size
        self.get_population = get_population
        self.get_xy_position = get_xy_position

        self.assert_correct_color_input()
        if self.color_attribute is not None:
            self.color_dict = {}
            self.fill_color_dict()

    def set_size(self):
        self.width = self.figure.cell_width * self.size
        self.height = self.figure.cell_height * self.size

    def setup(self, figure, renderer):
        self.figure = figure
        self.renderer = renderer
        self.model = self.renderer.model
        self.population = self.get_population(self.renderer.model)

        self.setup2()

        self.set_size()
        self.setup_sprites()

    def scale_x(self, x):
        pass

    def scale_y(self, y):
        pass

    def setup2(self):
        pass

    def draw(self):
        self.sprite_list.draw()

    def fill_color_dict(self):
        if isinstance(self.color_map, str):
            norm = matplotlib.colors.Normalize(vmin=self.color_vmin, vmax=self.color_vmax)
            cmap = matplotlib.colormaps[self.color_map]
            for color_value in range(self.color_vmin - 100, self.color_vmax + 101):
                rgb = tuple(int(255 * c) for c in cmap(norm(color_value))[0:3])
                self.color_dict[color_value] = rgb

        elif isinstance(self.color_map, dict):
            self.color_dict = {
                key: parse_color(color=color) for key, color in self.color_map.items()
            }
        else:
            raise ValueError("`color_map` must be a str or a dict.")

    def assert_correct_color_input(self):
        if self.color_attribute is None:
            assert self.color is not None

        elif self.color_attribute is not None:
            assert self.color_map is not None
            if isinstance(self.color_map, str):
                assert self.color_vmin is not None and self.color_vmax is not None

    def set_sprite_position(self, xy_position, sprite):
        x = self.scale_x(xy_position[0])
        y = self.scale_y(xy_position[1])

        if self.jitter:
            x += sprite.mesar_x_jitter
            y += sprite.mesar_y_jitter

        sprite.center_x = x
        sprite.center_y = y

    def set_sprite_color(self, entity, sprite):
        if self.color_attribute is None:
            sprite.color = self.color
        else:
            # get color attribute value with getattr()
            if isinstance(self.color_attribute, str):
                color_attribute_value = getattr(entity, self.color_attribute)
            # get it with a lambda
            else:
                color_attribute_value = self.color_attribute(entity)
            sprite.color = self.color_dict[color_attribute_value]

    def add_sprite(self, entity):
        sprite = self.create_sprite(entity=entity)
        self.sprite_list.append(sprite)
        self.sprite_dict[entity] = sprite

    def setup_sprites(self):
        self.selected_entities = self.select_entities()
        self.sprite_list = arcade.SpriteList(use_spatial_hash=False)
        self.sprite_dict = {}
        for entity in self.selected_entities:
            self.add_sprite(entity=entity)

    def create_sprite(self, entity):
        if self.shape == "rect":
            sprite = arcade.SpriteSolidColor(
                width=max(1, self.width),
                height=max(1, self.height),
            )
        elif self.shape == "circle":
            sprite = arcade.SpriteCircle(
                radius=max(1, min(self.width, self.height) / 2),
                color=arcade.color.BABY_BLUE,  # platzhalter
            )

        if self.jitter:
            sprite.mesar_x_jitter = (self.model.random.random() - 0.5) * self.figure.cell_width
            sprite.mesar_y_jitter = (self.model.random.random() - 0.5) * self.figure.cell_height

        self.set_sprite_position(
            xy_position=self.get_xy_position(entity),
            sprite=sprite,
        )
        self.set_sprite_color(entity=entity, sprite=sprite)
        return sprite

    def select_entities(self) -> list:
        return [entity for entity in self.population if self.filter_entities(entity)]

    def update(self):
        # TODO: use dirty_color and dirty_position to update only specific entity sprites

        # if the set of entities is not fixed during the simulation
        if self.dynamic_population:
            # get all relevant entities
            self.selected_entities = self.select_entities()
            updated_entities = set(self.selected_entities)

            # get all entities that were added previously
            entities_with_sprite = set(self.sprite_dict)

            # get all entities which need a sprite
            add_list = updated_entities - entities_with_sprite

            # add sprites for those entities
            for entity in add_list:
                self.add_sprite(entity=entity)

            # get all entities whose sprites can be removed
            remove_list = entities_with_sprite - updated_entities

            # remove the sprites of those entities
            for entity in remove_list:
                sprite = self.sprite_dict.pop(entity)
                self.sprite_list.remove(sprite)

        # if both colors and positions change during the simulation
        if self.dynamic_color and self.dynamic_position:
            for entity, sprite in self.sprite_dict.items():
                self.set_sprite_color(entity=entity, sprite=sprite)
                self.set_sprite_position(
                    xy_position=self.get_xy_position(entity),
                    sprite=sprite,
                )

        # if only the colors can change
        elif self.dynamic_color:
            for entity, sprite in self.sprite_dict.items():
                self.set_sprite_color(entity=entity, sprite=sprite)

        # if only the positions can change
        elif self.dynamic_position:
            for entity, sprite in self.sprite_dict.items():
                self.set_sprite_position(
                    xy_position=self.get_xy_position(entity),
                    sprite=sprite,
                )


class CellAgentArtists(Artist):
    """Renders agents on a cell-based grid space.

    Visualizes agents that live on cells.
    Each agent is drawn at its cell's coordinate position.

    Args:
        color: Default color for all entities. Used when color_attribute is None.
        color_attribute: Agent attribute name or callable to determine color.
        color_map: Matplotlib colormap name or dict mapping values to colors.
        color_vmin: Minimum value for colormap normalization.
        color_vmax: Maximum value for colormap normalization.
        shape: Entity shape, either "rect" or "circle".
        size: Size multiplier relative to cell size. Defaults to 1.
        filter_entities: Callable to filter which agents are displayed.
        jitter: If True, adds random offset to prevent overlap.
        dynamic_color: If True, updates colors each frame.
        dynamic_position: If True, updates positions each frame.
        dynamic_population: If True, handles agents being added/removed.
        get_population: Callable returning the agent collection from model.
        get_xy_position: Callable returning (x, y) position for an agent.
    """

    def __init__(
        self,
        color: Color = "black",
        color_attribute: str | Callable[[Any], Any] | None = None,
        color_map: str | dict[Any, Color] | None = "bwr",
        color_vmin: float | None = None,
        color_vmax: float | None = None,
        shape: Literal["rect", "circle"] = "circle",
        size: float = 1,
        filter_entities: Callable[[Any], bool] = lambda entity: True,
        jitter: bool = False,
        dynamic_color: bool = True,
        dynamic_position: bool = True,
        dynamic_population: bool = True,
        get_population: Callable[[mesa.Model], Any] = lambda model: model.agents,
        get_xy_position: Callable[[Any], tuple[float, float]] = lambda agent: (
            agent.cell.coordinate
        ),
    ) -> None:
        super().__init__(
            get_xy_position=get_xy_position,
            color=color,
            color_attribute=color_attribute,
            color_map=color_map,
            color_vmin=color_vmin,
            color_vmax=color_vmax,
            shape=shape,
            size=size,
            filter_entities=filter_entities,
            jitter=jitter,
            dynamic_color=dynamic_color,
            dynamic_position=dynamic_position,
            dynamic_population=dynamic_population,
            get_population=get_population,
        )

    def scale_x(self, x):
        return x * self.figure.cell_width + self.figure.x + self.figure.cell_width / 2

    def scale_y(self, y):
        return y * self.figure.cell_height + self.figure.y + self.figure.cell_height / 2


class CellArtists(Artist):
    """Renders cells of a grid space.

    Visualizes the cells themselves (not agents on them). Useful for displaying
    cell properties like terrain type, resource levels, or other cell attributes.

    Args:
        color: Default color for all cells. Used when color_attribute is None.
        color_attribute: Cell attribute name or callable to determine color.
        color_map: Matplotlib colormap name or dict mapping values to colors.
        color_vmin: Minimum value for colormap normalization.
        color_vmax: Maximum value for colormap normalization.
        shape: Cell shape, either "rect" or "circle".
        size: Size multiplier relative to cell size. Defaults to 1.
        filter_entities: Callable to filter which cells are displayed.
        jitter: If True, adds random offset to cell positions.
        dynamic_color: If True, updates colors each frame.
        dynamic_position: If True, updates positions each frame.
        dynamic_population: If True, handles cells being added/removed.
        get_population: Callable returning the cell collection from model.
        get_xy_position: Callable returning (x, y) position for a cell.
    """

    def __init__(
        self,
        color: Color = "grey",
        color_attribute: str | Callable[[Any], Any] | None = None,
        color_map: str | dict[Any, Color] | None = "bwr",
        color_vmin: float | None = None,
        color_vmax: float | None = None,
        shape: Literal["rect", "circle"] = "rect",
        size: float = 1,
        filter_entities: Callable[[Any], bool] = lambda entity: True,
        jitter: bool = False,
        dynamic_color: bool = True,
        dynamic_position: bool = False,
        dynamic_population: bool = True,
        get_population: Callable[[mesa.Model], Any] = lambda model: model.grid,
        get_xy_position: Callable[[Any], tuple[float, float]] = lambda cell: cell.coordinate,
    ) -> None:
        super().__init__(
            get_xy_position=get_xy_position,
            color=color,
            color_attribute=color_attribute,
            color_map=color_map,
            color_vmin=color_vmin,
            color_vmax=color_vmax,
            shape=shape,
            size=size,
            filter_entities=filter_entities,
            jitter=jitter,
            dynamic_color=dynamic_color,
            dynamic_position=dynamic_position,
            dynamic_population=dynamic_population,
            get_population=get_population,
        )

    def scale_x(self, x):
        return x * self.figure.cell_width + self.figure.x + self.figure.cell_width / 2

    def scale_y(self, y):
        return y * self.figure.cell_height + self.figure.y + self.figure.cell_height / 2


class ContinuousSpaceAgentArtists(Artist):
    """Renders agents in a continuous space.

    Visualizes agents that move freely in continuous 2D space rather than
    on a discrete grid. Positions are scaled to fit the plot area.

    Args:
        color: Default color for all agents. Used when color_attribute is None.
        color_attribute: Agent attribute name or callable to determine color.
        color_map: Matplotlib colormap name or dict mapping values to colors.
        color_vmin: Minimum value for colormap normalization.
        color_vmax: Maximum value for colormap normalization.
        shape: Agent shape, either "rect" or "circle".
        size: Size multiplier for agent sprites. Defaults to 2.
        filter_entities: Callable to filter which agents are displayed.
        jitter: If True, adds random offset to prevent overlap.
        dynamic_color: If True, updates colors each frame.
        dynamic_position: If True, updates positions each frame.
        dynamic_population: If True, handles agents being added/removed.
        get_population: Callable returning the agent collection from model.
        get_xy_position: Callable returning (x, y) position for an agent.
    """

    def __init__(
        self,
        color: Color = "black",
        color_attribute: str | Callable[[Any], Any] | None = None,
        color_map: str | dict[Any, Color] | None = "bwr",
        color_vmin: float | None = None,
        color_vmax: float | None = None,
        shape: Literal["rect", "circle"] = "circle",
        size: float = 2,
        filter_entities: Callable[[Any], bool] = lambda entity: True,
        jitter: bool = False,
        dynamic_color: bool = True,
        dynamic_position: bool = True,
        dynamic_population: bool = True,
        get_population: Callable[[mesa.Model], Any] = lambda model: model.agents,
        get_xy_position: Callable[[Any], tuple[float, float]] = lambda agent: agent.position,
    ) -> None:
        super().__init__(
            get_xy_position=get_xy_position,
            color=color,
            color_attribute=color_attribute,
            color_map=color_map,
            color_vmin=color_vmin,
            color_vmax=color_vmax,
            shape=shape,
            size=size,
            filter_entities=filter_entities,
            jitter=jitter,
            dynamic_color=dynamic_color,
            dynamic_position=dynamic_position,
            dynamic_population=dynamic_population,
            get_population=get_population,
        )

    def scale_x(self, x):
        return x * self.figure.cell_width + self.figure.x

    def scale_y(self, y):
        return y * self.figure.cell_height + self.figure.y


class NetworkCellArtists(Artist):
    """Renders nodes and edges of a network graph.

    Visualizes cells as nodes in a network layout.

    Args:
        color: Default color for all nodes. Used when color_attribute is None.
        color_attribute: Cell attribute name or callable to determine color.
        color_map: Matplotlib colormap name or dict mapping values to colors.
        color_vmin: Minimum value for colormap normalization.
        color_vmax: Maximum value for colormap normalization.
        shape: Node shape, either "rect" or "circle".
        size: Size multiplier for node sprites. Defaults to 2.
        filter_entities: Callable to filter which nodes are displayed.
        jitter: If True, adds random offset to node positions.
        dynamic_color: If True, updates colors each frame.
        dynamic_position: If True, updates positions each frame.
        dynamic_population: If True, handles nodes being added/removed.
        networkx_layout: Layout function from networkx (e.g., spring_layout).
        get_population: Callable returning the cell collection from model.
        get_xy_position: Callable returning (x, y) position for a cell.
    """

    def __init__(
        self,
        color: Color = "black",
        color_attribute: str | Callable[[Any], Any] | None = None,
        color_map: str | dict[Any, Color] | None = "bwr",
        color_vmin: float | None = None,
        color_vmax: float | None = None,
        shape: Literal["rect", "circle"] = "circle",
        size: float = 2,
        filter_entities: Callable[[Any], bool] = lambda entity: True,
        jitter: bool = False,
        dynamic_color: bool = True,
        dynamic_position: bool = True,
        dynamic_population: bool = True,
        networkx_layout: Callable[..., dict[Any, Any]] = nx.spring_layout,
        get_population: Callable[[mesa.Model], Any] = lambda model: model.grid,
        get_xy_position: Callable[[Any], tuple[float, float]] = lambda cell: (
            cell._MESARCADE_NETWORK_POSITION
        ),
    ) -> None:
        super().__init__(
            get_xy_position=get_xy_position,
            color=color,
            color_attribute=color_attribute,
            color_map=color_map,
            color_vmin=color_vmin,
            color_vmax=color_vmax,
            shape=shape,
            size=size,
            filter_entities=filter_entities,
            jitter=jitter,
            dynamic_color=dynamic_color,
            dynamic_position=dynamic_position,
            dynamic_population=dynamic_population,
            get_population=get_population,
        )
        self.networkx_layout = networkx_layout

    def _get_node_positions(self):
        self.layout_positions = self.networkx_layout(self.model.grid.G)

        self.max_x_node_position = max(
            [abs(self.layout_positions[i][0]) for i in self.layout_positions]
        )
        self.max_y_node_position = max(
            [abs(self.layout_positions[i][1]) for i in self.layout_positions]
        )

        self.node_size = max(
            10 / math.log10(len(self.layout_positions)) * self.figure.width / 400,
            1,
        )

        self.figure.cell_width = self.node_size
        self.figure.cell_height = self.node_size

        for i, cell in enumerate(self.model.grid):
            cell._MESARCADE_NETWORK_POSITION = self.layout_positions[i]
            cell._MESARCADE_NETWORK_POSITION[0] /= self.max_x_node_position
            cell._MESARCADE_NETWORK_POSITION[1] /= self.max_y_node_position

    def setup2(self):
        self._get_node_positions()

    def scale_x(self, x):
        return x * self.figure.width / 2.15 + self.figure.x + self.figure.width / 2

    def scale_y(self, y):
        return y * self.figure.height / 2.15 + self.figure.y + self.figure.height / 2

    def draw(self):
        edges = self.model.grid.G.edges()
        for u, v in edges:
            node_u_pos = self.layout_positions[u]
            node_v_pos = self.layout_positions[v]

            node_u_pos = [self.scale_x(node_u_pos[0]), self.scale_y(node_u_pos[1])]
            node_v_pos = [self.scale_x(node_v_pos[0]), self.scale_y(node_v_pos[1])]

            arcade.draw_line_strip(
                [node_u_pos, node_v_pos],
                color=arcade.color.BLACK,
                line_width=max(1, self.node_size / 5),
            )

        self.sprite_list.draw()


class NetworkAgentArtists(NetworkCellArtists):
    """Renders agents on a network graph.

    Visualizes agents that live on network nodes.

    Args:
        color: Default color for all agents. Used when color_attribute is None.
        color_attribute: Agent attribute name or callable to determine color.
        color_map: Matplotlib colormap name or dict mapping values to colors.
        color_vmin: Minimum value for colormap normalization.
        color_vmax: Maximum value for colormap normalization.
        shape: Agent shape, either "rect" or "circle".
        size: Size multiplier for agent sprites. Defaults to 2.
        filter_entities: Callable to filter which agents are displayed.
        jitter: If True, adds random offset to prevent overlap.
        dynamic_color: If True, updates colors each frame.
        dynamic_position: If True, updates positions each frame.
        dynamic_population: If True, handles agents being added/removed.
        networkx_layout: Layout function from networkx (e.g., spring_layout).
        get_population: Callable returning the agent collection from model.
        get_xy_position: Callable returning (x, y) position for an agent.
    """

    def __init__(
        self,
        color: Color = "black",
        color_attribute: str | Callable[[Any], Any] | None = None,
        color_map: str | dict[Any, Color] | None = "bwr",
        color_vmin: float | None = None,
        color_vmax: float | None = None,
        shape: Literal["rect", "circle"] = "circle",
        size: float = 2,
        filter_entities: Callable[[Any], bool] = lambda entity: True,
        jitter: bool = False,
        dynamic_color: bool = True,
        dynamic_position: bool = True,
        dynamic_population: bool = True,
        networkx_layout: Callable[..., dict[Any, Any]] = nx.spring_layout,
        get_population: Callable[[mesa.Model], Any] = lambda model: model.agents,
        get_xy_position: Callable[[Any], tuple[float, float]] = lambda agent: (
            agent.cell._MESARCADE_NETWORK_POSITION
        ),
    ) -> None:
        super().__init__(
            get_xy_position=get_xy_position,
            color=color,
            color_attribute=color_attribute,
            color_map=color_map,
            color_vmin=color_vmin,
            color_vmax=color_vmax,
            shape=shape,
            size=size,
            filter_entities=filter_entities,
            jitter=jitter,
            dynamic_color=dynamic_color,
            dynamic_position=dynamic_position,
            dynamic_population=dynamic_population,
            get_population=get_population,
        )
        self.networkx_layout = networkx_layout
