from __future__ import annotations

import arcade
import numpy as np

from mesarcade.figure import Figure
from mesarcade.utils import parse_color

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Any, Callable

    import mesa

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


def rescale(value, old_min, old_max, new_min, new_max):
    old_range = old_max - old_min
    new_range = new_max - new_min
    if old_range > 0:
        return (value - old_min) / old_range * new_range + new_min
    else:
        return (value - old_min) * new_range + new_min


def rescale_array_column_inplace(
    np_array: np.ndarray,
    col: int,
    old_min: float,
    old_max: float,
    new_min: float,
    new_max: float,
) -> None:
    values = np_array[:, col]

    old_range = old_max - old_min
    new_range = new_max - new_min

    if old_range > 0:
        values[:] = (values - old_min) / old_range * new_range + new_min
    else:
        values[:] = (values - old_min) * new_range + new_min


class _ModelHistoryPlot:
    def __init__(
        self,
        model_attributes: list[str | Callable[[mesa.Model], Any]],
        ylim: list[float | None, float | None] | None = None,
        labels: list[str] | None = None,
        colors: list[str] | None = None,
        rendering_step: int = 5,
        title=None,
        legend: bool = True,
        from_datacollector: bool = False,
    ):
        self.model_attrs = model_attributes
        self.lower_y_lim = ylim[0] if ylim is not None else None
        self.upper_y_lim = ylim[1] if ylim is not None else None
        self.labels = labels
        self.rendering_step = rendering_step
        self.legend = legend
        self.title = title
        self.from_datacollector = from_datacollector
        self.colors = None

        self.padding = 10

        self.validate_input()

        if colors is not None:
            self.colors = [parse_color(color) for color in colors]
        else:
            self.colors = [
                arcade.color.NAVY_BLUE,
                arcade.color.ORANGE,
                arcade.color.GREEN,
                arcade.color.RED,
                arcade.color.PINK,
                arcade.color.PURPLE,
            ]

    def validate_input(self):
        if len(self.model_attrs) > 6:
            raise ValueError("Only 6 lines allowed!")

        if self.labels is not None:
            if len(self.model_attrs) != len(self.labels):
                raise ValueError(
                    "The arguments model_attributes and labels must have the same length."
                )

        if self.colors is not None:
            if len(self.model_attrs) != len(self.colors):
                raise ValueError(
                    "The arguments model_attributes and colors must have the same length."
                )

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

        self.x = self.figure.x
        self.y = self.figure.y

        self.width = self.figure.width
        self.height = self.figure.height

        self.font_size = int(self.height * 0.025)

        self.plot_area_x = self.x + self.width * 0.15
        self.plot_area_y = self.y + self.height * 0.3
        self.plot_area_width = self.width * (0.85 - 0.025)
        self.plot_area_height = self.height * (0.7 - 0.025)

        self.data_dict = {model_attr: [] for model_attr in self.model_attrs}
        self.scaled_data_dict = {model_attr: [] for model_attr in self.model_attrs}
        self.min_y = self.lower_y_lim
        self.max_y = self.upper_y_lim
        self.min_x = 0
        self.max_x = 0

        self.create_plot_area()
        self.create_axis_ticks()

        # add legend
        if self.legend:
            self.create_legend()

    def create_plot_area(self):
        background = arcade.shape_list.create_rectangle_filled(
            center_x=self.plot_area_x + self.plot_area_width / 2,
            center_y=self.plot_area_y + self.plot_area_height / 2,
            width=self.plot_area_width,
            height=self.plot_area_height,
            color=arcade.color.WHITE,
        )
        self.figure.shape_list.append(background)

        outline = arcade.shape_list.create_rectangle_outline(
            center_x=self.plot_area_x + self.plot_area_width / 2,
            center_y=self.plot_area_y + self.plot_area_height / 2,
            width=self.plot_area_width,
            height=self.plot_area_height,
            color=arcade.color.BLACK,
            border_width=2,
        )
        self.figure.shape_list.append(outline)

    def create_axis_ticks(self):
        min_y_tick_y_pos = self.plot_area_y + self.padding
        mid_y_tick_y_pos = (
            self.plot_area_y + self.padding + (self.plot_area_height - self.padding * 2) / 2
        )
        max_y_tick_y_pos = self.plot_area_y + self.plot_area_height - self.padding

        self.min_y_label = arcade.Text(
            text="",
            x=0,
            y=min_y_tick_y_pos - self.font_size / 2,
            color=arcade.color.BLACK,
            font_size=self.font_size,
            batch=self.figure.text_batch,
        )

        self.mid_y_label = arcade.Text(
            text="",
            x=0,
            y=mid_y_tick_y_pos - self.font_size / 2,
            color=arcade.color.BLACK,
            font_size=self.font_size,
            batch=self.figure.text_batch,
        )

        self.max_y_label = arcade.Text(
            text="",
            x=0,
            y=max_y_tick_y_pos - self.font_size / 2,
            color=arcade.color.BLACK,
            font_size=self.font_size,
            batch=self.figure.text_batch,
        )

        for i in range(5):
            self.figure.shape_list.append(
                arcade.shape_list.create_rectangle(
                    center_x=self.plot_area_x - 2,
                    center_y=min_y_tick_y_pos + (self.plot_area_height - self.padding * 2) / 4 * i,
                    width=3,
                    height=3,
                    color=arcade.color.BLACK,
                )
            )

    def create_legend(self):
        label_x = self.figure.x + self.width * 0.1
        label_y = self.plot_area_y - self.height * 0.1

        if self.labels is not None:
            labels = self.labels
        elif self.labels is None and isinstance(self.model_attrs[0], str):
            labels = self.model_attrs
        else:
            labels = ["no label"] * len(self.model_attrs)

        for i, label in enumerate(labels):
            if i % 2 == 0:
                label_y -= self.font_size * 2
            if i % 2 != 0:
                label_x_ = label_x + self.figure.width / 2
            else:
                label_x_ = label_x

            label_element = arcade.Text(
                text=label,
                x=label_x_,
                y=label_y,
                batch=self.figure.text_batch,
                color=self.renderer.font_color,
                font_size=self.font_size,
            )
            self.figure.text_list.append(label_element)

            color_dot = arcade.shape_list.create_ellipse(
                center_x=label_x_ - self.font_size,
                center_y=label_y + self.font_size / 3,
                width=self.font_size,
                height=self.font_size,
                color=self.colors[i],
            )
            self.figure.shape_list.append(color_dot)

    def update(self):
        # get the current time step
        tick = self.renderer.tick

        # check if it is time to update
        if tick % self.rendering_step == 0 or tick <= 1:
            # for each model attribute that has to be collected
            for model_attr in self.model_attrs:
                # model attribute was given as string?
                if isinstance(model_attr, str):
                    # get the value from a model attribute
                    if not self.from_datacollector:
                        y = getattr(self.model, model_attr)

                    # or get the value from the datacollector
                    else:
                        y_data = self.model.datacollector.model_vars[model_attr]
                        y = y_data[-1] if len(y_data) > 0 else None

                # or as lambda?
                else:
                    y = model_attr(self.model)

                # set min and max values initially
                if self.max_y is None:
                    self.max_y = y + 0.000001 * y
                if self.min_y is None:
                    self.min_y = y

                # if there valid new data
                if y is not None and np.isfinite(y):
                    # add new data point
                    self.data_dict[model_attr].append((tick, y))

                    # update min and max values
                    if self.lower_y_lim is None:
                        if y < self.min_y:
                            self.min_y = y

                    if self.upper_y_lim is None:
                        if y > self.max_y:
                            self.max_y = y

                # Rescale the data
                # TODO: Optimize this with numpy
                self.scaled_data_dict[model_attr] = [
                    (
                        rescale(
                            value=x,
                            old_min=0,
                            old_max=tick,
                            new_min=self.plot_area_x + self.padding,
                            new_max=self.plot_area_x + self.plot_area_width - self.padding,
                        ),
                        rescale(
                            value=y,
                            old_min=self.min_y,
                            old_max=self.max_y,
                            new_min=self.plot_area_y + self.padding,
                            new_max=self.plot_area_y + self.plot_area_height - self.padding,
                        ),
                    )
                    for x, y in self.data_dict[model_attr]
                ]

            # update lower y_label
            str_min_y_label = str(round(self.min_y, 3))
            if self.min_y_label.text != str_min_y_label:
                self.min_y_label.text = str_min_y_label
                self.min_y_label.x = (
                    self.plot_area_x - (len(str_min_y_label) + 2) * self.font_size / 1.5
                )

            # update medium y_label
            str_mid_y_label = str(round((self.min_y + self.max_y) / 2, 3))
            if self.mid_y_label.text != str_mid_y_label:
                self.mid_y_label.text = str_mid_y_label
                self.mid_y_label.x = (
                    self.plot_area_x - (len(str_mid_y_label) + 2) * self.font_size / 1.5
                )

            # update upper y_label
            str_max_y_label = str(round(self.max_y, 3))
            if self.max_y_label != str_max_y_label:
                self.max_y_label.text = str_max_y_label
                self.max_y_label.x = (
                    self.plot_area_x - (len(str_max_y_label) + 2) * self.font_size / 1.5
                )

    def draw(self):
        for i, model_attr in enumerate(self.model_attrs):
            arcade.draw_line_strip(
                self.scaled_data_dict[model_attr],
                color=self.colors[i],
                line_width=2,
            )


class ModelHistoryPlot(Figure):
    """A line plot that displays the history of model attributes over time.

    Tracks one or more model attributes and visualizes their values as lines
    over the course of the simulation. Supports up to 6 lines.

    Args:
        model_attributes: List of model attributes to track. Each element can be
            a string (attribute name) or a callable that takes a mesa.Model and
            returns a numeric value.
        ylim: Optional y-axis limits as [min, max]. If None, limits are
            determined automatically from the data. Individual elements can be
            None to auto-scale only that bound.
        labels: Optional labels for the legend. If None, attribute names are
            used for string attributes, or "no label" for callables.
        colors: Optional colors for the lines. Accepts color names, RGB tuples,
            or RGBA tuples. Defaults to a predefined color palette.
        legend: Whether to display a legend. Defaults to True.
        title: Optional title for the plot.
        from_datacollector: If True, reads values from the model's datacollector
            instead of directly from model attributes. Only works with string
            attributes. Defaults to False.
        rendering_step: Simulation steps between plot updates. Defaults to 3.
    """

    def __init__(
        self,
        model_attributes: list[str],
        ylim: list[float, float] | None = None,
        labels: list[str] | None = None,
        colors: list[Color] | None = None,
        legend: bool = True,
        title: str | None = None,
        from_datacollector: bool = False,
        rendering_step: int = 3,
    ) -> None:
        plot = _ModelHistoryPlot(
            model_attributes=model_attributes,
            ylim=ylim,
            labels=labels,
            colors=colors,
            legend=legend,
            title=title,
            rendering_step=rendering_step,
            from_datacollector=from_datacollector,
        )
        super().__init__(components=[plot], title=title, get_space=None)
