"""
Defines strategies for selecting and assigning colors to chart elements based on data attributes.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import ClassVar

import polars as pl
from polychromos.palette import HSLColorScale, HSLColorSequence, Palette


class ColorSelectorStrategy:
    """
    Base strategy for selecting data points to apply colors to.
    """

    type_id: ClassVar[str] = "All"

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        """
        Selects which entries in the DataFrame should be colored.

        :param df: The DataFrame containing the data.
        :type df: pl.DataFrame
        :param category_var: The name of the category variable.
        :type category_var: str
        :param value_var: The name of the value variable.
        :type value_var: str
        :return: A list of booleans indicating which rows are selected.
        :rtype: list[bool]
        """
        return [True] * df.height

    @classmethod
    def get_name(cls) -> str:
        """
        Gets the name of the color selector strategy.

        :return: Human-readable name.
        :rtype: str
        """
        return "All entries"

    @classmethod
    def get_description(cls) -> str:
        """
        Gets the description of the color selector strategy.

        :return: Human-readable description.
        :rtype: str
        """
        return (
            "Selects all entries. This should be used only in the first color rule, as it will "
            "override all previous rules' selection criteria."
        )

    @classmethod
    def get_icon(cls) -> str | None:
        """
        Gets an optional SVG icon for the color selector strategy.

        :return: An SVG string, or ``None`` if not set.
        :rtype: str | None
        """
        return None

    def explain(self) -> str:
        """
        Returns a human-readable participial phrase describing the selection criteria.

        :return: A phrase such as ``"selecting all entries"``.
        :rtype: str
        """
        return "selecting all entries"


class ColorSelectByCategoryName(ColorSelectorStrategy):
    """
    Selects data points based on category names.
    """

    type_id: ClassVar[str] = "ByCategoryName"

    def __init__(
        self,
        names: list[str],
    ) -> None:
        self._names = names

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        return [name in self._names for name in df[category_var]]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "By category names"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries whose category names are in a given list."

    def explain(self) -> str:
        return f"selecting entries in categories {self._names}"


class ColorSelectValuesBelow(ColorSelectorStrategy):
    """
    Selects data points with values below a certain threshold.
    """

    type_id: ClassVar[str] = "ValuesBelow"

    def __init__(
        self,
        threshold: float,
    ) -> None:
        self._threshold = threshold

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        return [value <= self._threshold for value in df[value_var]]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "Values below threshold"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries with values below a specified threshold."

    def explain(self) -> str:
        return f"selecting entries with a value below {self._threshold}"


class ColorSelectValuesAbove(ColorSelectorStrategy):
    """
    Selects data points with values above a certain threshold.
    """

    type_id: ClassVar[str] = "ValuesAbove"

    def __init__(
        self,
        threshold: float,
    ) -> None:
        self._threshold = threshold

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        return [value >= self._threshold for value in df[value_var]]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "Values above threshold"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries with values above a specified threshold."

    def explain(self) -> str:
        return f"selecting entries with a value above {self._threshold}"


class ColorSelectMaxValue(ColorSelectorStrategy):
    """
    Selects data points with the maximum value.
    """

    type_id: ClassVar[str] = "MaxValue"

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        if df.height == 0:
            return []
        max_value = df[value_var].max()
        if max_value is None:
            return [False] * df.height
        return [value == max_value for value in df[value_var]]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "Maximum value"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries with the maximum value."

    def explain(self) -> str:
        return "selecting the entry with the highest value"


class ColorSelectMinValue(ColorSelectorStrategy):
    """
    Selects data points with the minimum value.
    """

    type_id: ClassVar[str] = "MinValue"

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        if df.height == 0:
            return []
        min_value = df[value_var].min()
        if min_value is None:
            return [False] * df.height
        return [value == min_value for value in df[value_var]]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "Minimum value"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries with the minimum value."

    def explain(self) -> str:
        return "selecting the entry with the lowest value"


class ColorSelectValuesBelowMean(ColorSelectorStrategy):
    """
    Selects data points with values below the total weighted mean.
    """

    type_id: ClassVar[str] = "ValuesBelowMean"

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        return [value <= df["__total_mean__"][i] for i, value in enumerate(df[value_var])]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "Values below mean"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries with values below the total weighted mean."

    def explain(self) -> str:
        return "selecting entries with a value below the mean"


class ColorSelectValuesAboveMean(ColorSelectorStrategy):
    """
    Selects data points with values above the total weighted mean.
    """

    type_id: ClassVar[str] = "ValuesAboveMean"

    def selector(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
    ) -> list[bool]:
        return [value >= df["__total_mean__"][i] for i, value in enumerate(df[value_var])]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this selector strategy."""
        return "Values above mean"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this selector strategy."""
        return "Selects entries with values above the total weighted mean."

    def explain(self) -> str:
        return "selecting entries with a value above the mean"


class ColorAssignerStrategy(ABC):
    """
    Base strategy for assigning colors to data points.
    """

    type_id: ClassVar[str]

    def __init__(
        self,
        anchor_colors: HSLColorSequence,
    ) -> None:
        """
        Constructor method.

        :param anchor_colors: A list of colors used to generate the palette to assign colors from.
        :type anchor_colors: HSLColorSequence
        :raises ValueError: When no anchor colors are provided.
        """
        if len(anchor_colors) == 0:
            raise ValueError("No anchor colors provided")
        self._anchor_colors = anchor_colors

    @property
    def anchor_colors(self) -> HSLColorSequence:
        """
        The anchor colors used to generate the palette.

        :return: The sequence of anchor colors.
        :rtype: HSLColorSequence
        """
        return self._anchor_colors

    @property
    def continuous(self) -> bool:
        """
        Whether the color assigner is continuous (``True``) or discrete (``False``).
        """
        return False

    @abstractmethod
    def assign_colors(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
        group_var: str | None = None,
    ) -> HSLColorSequence:
        """
        Generates a palette of colors according to the strategy.

        :param df: The DataFrame containing the data.
        :type df: pl.DataFrame
        :param category_var: The name of the category variable.
        :type category_var: str
        :param value_var: The name of the value variable.
        :type value_var: str
        :param group_var: The name of the groups variable, if any. Defaults to ``None``.
        :type group_var: str | None, optional
        :return: A sequence of colors assigned to the data points.
        :rtype: HSLColorSequence
        """

    @classmethod
    @abstractmethod
    def get_name(cls) -> str:
        """
        Gets the name of the color assigner strategy.
        """

    @classmethod
    @abstractmethod
    def get_description(cls) -> str:
        """
        Gets the description of the color assigner strategy.
        """

    @classmethod
    def get_icon(cls) -> str | None:
        """
        Gets an optional SVG icon for the color assigner strategy.

        :return: An SVG string, or ``None`` if not set.
        :rtype: str | None
        """
        return None

    def _hex_anchors(self) -> list[str]:
        return [c.to_css_hex() for c in self._anchor_colors]

    @abstractmethod
    def explain(self) -> str:
        """
        Returns a human-readable imperative phrase describing the color assignment.

        :return: A phrase such as ``"colorize them each in a distinct color from ['#ff0000']"``.
        :rtype: str
        """


class DiscreteColorAssignerStrategy(ColorAssignerStrategy):
    """
    Strategy for assigning discrete colors to data points.
    """

    def __init__(
        self,
        anchor_colors: HSLColorSequence,
        shuffle: bool = False,
    ) -> None:
        """
        Constructor method.

        :param anchor_colors: A list of colors used to generate the palette to assign colors from.
        :type anchor_colors: HSLColorSequence
        :param shuffle: Whether to alternate colors or not. Defaults to ``False``.
        :type shuffle: bool, optional
        """
        super().__init__(anchor_colors=anchor_colors)
        self._shuffle = shuffle

    def _calculate_interpolated_sequence(
        self,
        n_colors: int,
    ) -> HSLColorSequence:
        if n_colors <= 1:
            return [self._anchor_colors[0]]

        sequence: HSLColorSequence
        if len(self._anchor_colors) == 1:
            sequence = self._anchor_colors * n_colors
        elif len(self._anchor_colors) == 2 and n_colors > 2:
            sequence = Palette.sequence_from_cylindrical_interpolation(
                self._anchor_colors[0],
                self._anchor_colors[1],
                n_colors,
                path_strategy=Palette.CylindricalInterpolationPath.SHORTEST,
            )
        else:
            scale: HSLColorScale = Palette.to_color_scale(self._anchor_colors)
            sequence = Palette.color_scale_to_color_sequence(scale, n_colors)

        if self._shuffle:
            sequence = Palette.alternate_colors(sequence)

        return sequence


class ColorAssignerByCategory(DiscreteColorAssignerStrategy):
    """
    Strategy for assigning colors based on category values.
    """

    type_id: ClassVar[str] = "ByCategory"

    def __init__(
        self,
        anchor_colors: HSLColorSequence,
        shuffle: bool = False,
    ) -> None:
        super().__init__(
            anchor_colors=anchor_colors,
            shuffle=shuffle,
        )

    def assign_colors(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
        group_var: str | None = None,
    ) -> HSLColorSequence:
        return self._calculate_interpolated_sequence(len(df[category_var]))

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this assigner strategy."""
        return "Assign by category"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this assigner strategy."""
        return "Assigns colors based on the distinct category values."

    def explain(self) -> str:
        return f"colorize them each in a distinct color from {self._hex_anchors()}"


class ColorAssignerByGroup(DiscreteColorAssignerStrategy):
    """
    Strategy for assigning colors based on group values.
    """

    type_id: ClassVar[str] = "ByGroup"

    def __init__(
        self,
        anchor_colors: HSLColorSequence,
        shuffle: bool = False,
    ) -> None:
        super().__init__(
            anchor_colors=anchor_colors,
            shuffle=shuffle,
        )

    def assign_colors(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
        group_var: str | None = None,
    ) -> HSLColorSequence:
        if group_var is None:
            raise ValueError("Group variable must be provided for ColorAssignerByGroup")
        return self._calculate_interpolated_sequence(len(df[group_var].unique()))

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this assigner strategy."""
        return "Assign by group"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this assigner strategy."""
        return "Assigns colors based on the distinct group values."

    def explain(self) -> str:
        return f"colorize them each in a distinct color by group from {self._hex_anchors()}"


class ColorAssignerByValue(ColorAssignerStrategy):
    """
    Strategy for assigning colors based on entry values.
    """

    type_id: ClassVar[str] = "ByValue"

    def __init__(
        self,
        anchor_colors: HSLColorSequence,
        range_min: float | None = None,
        range_max: float | None = None,
        range_center_at: float | None = None,
    ) -> None:
        super().__init__(anchor_colors=anchor_colors)
        self._range_min = range_min
        self._range_max = range_max
        self._range_center_at = range_center_at

    @property
    def continuous(self) -> bool:
        return True

    @property
    def range_min(self) -> float | None:
        """
        Gets the minimum value of the range for color assignment.
        """
        return self._range_min

    @property
    def range_max(self) -> float | None:
        """
        Gets the maximum value of the range for color assignment.
        """
        return self._range_max

    @property
    def range_center_at(self) -> float | None:
        """
        Gets the center value of the range for color assignment.
        """
        return self._range_center_at

    def assign_colors(
        self,
        df: pl.DataFrame,
        category_var: str,
        value_var: str,
        group_var: str | None = None,
    ) -> HSLColorSequence:
        scale: HSLColorScale = Palette.to_color_scale(self._anchor_colors if len(self._anchor_colors) > 1 else self._anchor_colors * 2)
        temp_min: float = self._range_min if self._range_min is not None else min(df[value_var])
        temp_max: float = self._range_max if self._range_max is not None else max(df[value_var])
        low_end: float
        high_end: float
        if self._range_center_at is not None:
            max_offset: float = max(
                abs(self._range_center_at - temp_min),
                abs(self._range_center_at - temp_max),
            )
            low_end = self._range_center_at - max_offset
            high_end = self._range_center_at + max_offset
        else:
            low_end = temp_min
            high_end = temp_max
        return [Palette.scale_lerp(scale, (low_end, high_end), value) for value in df[value_var]]

    @classmethod
    def get_name(cls) -> str:
        """Returns the name of this assigner strategy."""
        return "Assign by value"

    @classmethod
    def get_description(cls) -> str:
        """Returns the description of this assigner strategy."""
        return "Assigns colors based on the entry values, optionally within a specified range."

    def explain(self) -> str:
        hexes = self._hex_anchors()
        color_part = f"colors from {hexes[0]} to {hexes[-1]}"
        parts = [f"colorize them by their value in {color_part}"]
        if self._range_min is not None and self._range_max is not None:
            parts.append(f"restricted to values {self._range_min}\u2013{self._range_max}")
        elif self._range_min is not None:
            parts.append(f"from {self._range_min}")
        elif self._range_max is not None:
            parts.append(f"up to {self._range_max}")
        if self._range_center_at is not None:
            parts.append(f"centered at {self._range_center_at}")
        return " (".join(parts) + (")" if len(parts) > 1 else "")


class ColorRule:
    """
    Encapsulates a color rule for chart elements.
    """

    def __init__(
        self,
        assigner: ColorAssignerStrategy,
        selector: ColorSelectorStrategy | None = None,
    ) -> None:
        self._assigner = assigner
        self._selector = selector if selector is not None else ColorSelectorStrategy()

    @property
    def selector(self) -> ColorSelectorStrategy:
        """
        Gets the color selector strategy.
        """
        return self._selector

    @property
    def assigner(self) -> ColorAssignerStrategy:
        """
        Gets the color assigner strategy.
        """
        return self._assigner

    def explain(self) -> str:
        """
        Returns a human-readable description of what this color rule does.

        Combines the selector and assigner explanations into a single sentence,
        e.g. ``"selecting entries with a value above 12, colorize them by their
        value in colors from #00a000 to #00c000"``.

        :return: A prose description of the rule.
        :rtype: str
        """
        return f"{self._selector.explain()}, {self._assigner.explain()}"


@dataclass(frozen=True)
class ColorStrategyDescription:
    """
    Value object describing available color strategies.
    """

    supported_selectors: list[type[ColorSelectorStrategy]]
    supported_assigners: list[type[ColorAssignerStrategy]]


__all__ = [
    "ColorSelectorStrategy",
    "ColorSelectByCategoryName",
    "ColorSelectValuesBelow",
    "ColorSelectValuesAbove",
    "ColorSelectMaxValue",
    "ColorSelectMinValue",
    "ColorSelectValuesBelowMean",
    "ColorSelectValuesAboveMean",
    "ColorAssignerStrategy",
    "DiscreteColorAssignerStrategy",
    "ColorAssignerByCategory",
    "ColorAssignerByGroup",
    "ColorAssignerByValue",
    "ColorRule",
    "ColorStrategyDescription",
]
