"""
Grouped bar chart figure strategy implementation.
"""

from typing import Any

import plotly.express as px
import plotly.graph_objects as go
import polars as pl
from babel.support import Translations
from fields_metadata import FieldMetadata
from polychromos.palette import HSLColorSequence, Palette

from orama.color import (
    ColorAssignerByGroup,
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectByCategoryName,
    ColorSelectorStrategy,
)
from orama.options import (
    BooleanOption,
    ChartOptionsMap,
)
from orama.query import QueryParams
from orama.strategies.bar import AbstractBarChartStrategy
from orama.theme import Theme
from orama.variables import (
    Aggregation,
    CategoricalVariable,
    ChartVariablesMap,
    NumericalAggregationVariable,
)


class GroupedBarChartStrategy(AbstractBarChartStrategy):
    """
    Grouped bar chart figure strategy.
    """

    _variables: ChartVariablesMap = {
        "category": CategoricalVariable(
            name="Category",
            description=(
                "Categorical variable(s) for the bar chart. "
                "Each category corresponds to a bar. "
                "If several variables are provided, they will be concatenated in a cartesian "
                "product with all of their values to form the final categories."
            ),
            map_to="category_variables",
            multicategorical=True,
        ),
        "value": NumericalAggregationVariable(
            name="Value",
            description=(
                "An aggregation of a numerical variable to produce the final value of each "
                "category. This value is proportional to the length of each bar."
            ),
            map_to="aggregation",
        ),
        "grouping": CategoricalVariable(
            name="Grouping",
            description=(
                "Categorical variable for grouping the bars. "
                "Categories are split in several groups, one per each value of this variable."
            ),
            map_to="group_variable",
            multicategorical=False,
        ),
    }
    _options: ChartOptionsMap = {
        "stacked": BooleanOption(
            name="Stacked bars",
            description=(
                "When enabled, group bars are stacked on top of each other. When disabled, group bars are rendered side by side."
            ),
            map_to="stacked",
            default_value=False,
        ),
        "relative": BooleanOption(
            name="Relative values",
            description=(
                "When enabled, bar values are shown as a percentage of the total value of their "
                "category. Only available when stacked is also enabled."
            ),
            map_to="relative",
            default_value=False,
        ),
        "horizontal": BooleanOption(
            name="Horizontal bars",
            description="When enabled, bars are rendered horizontally instead of vertically.",
            map_to="horizontal",
            default_value=False,
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        category_variables: list[str],
        group_variable: str,
        aggregation: Aggregation,
        stacked: bool = False,
        relative: bool = False,
        horizontal: bool = False,
        color_rules: list[ColorRule] | None = None,
        title: str | None = None,
        subtitle: str | None = None,
        caption: str | None = None,
        description: str | None = None,
        storytelling_context: str | None = None,
        category_label: str | None = None,
        value_label: str | None = None,
        group_label: str | None = None,
    ) -> None:
        super().__init__(
            fields_metadata=fields_metadata,
            category_variables=category_variables,
            aggregation=aggregation,
            color_rules=color_rules,
            title=title,
            subtitle=subtitle,
            caption=caption,
            description=description,
            storytelling_context=storytelling_context,
            category_label=category_label,
            value_label=value_label,
        )

        self._variables["grouping"].validate_fields(fields_metadata, group_variable)

        self._options["stacked"].validate(stacked)
        self._options["relative"].validate(relative)
        if relative and not stacked:
            raise ValueError("Relative option can only be enabled when stacked is also enabled")
        self._options["horizontal"].validate(horizontal)

        self._group_variable = group_variable
        self._stacked = stacked
        self._relative = relative
        self._horizontal = horizontal
        self._group_label = group_label

    @classmethod
    def name(cls) -> str:
        return "Grouped Bar Chart"

    @classmethod
    def description(cls) -> str:
        return (
            "A grouped bar chart presents categorical data with rectangular bars with heights or "
            "lengths proportional to the values that they represent. The bars can be plotted "
            "vertically or horizontally. Categories are split in several groups, one per each "
            "value of a grouping variable. These charts add an extra categorical dimension to "
            "standard bar charts."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        query_params = super().get_query_params()
        query_params["main"].group_fields.append(self._group_variable)
        return query_params

    def _get_schema_fields(self) -> dict[str, type[pl.DataType]]:
        fields = super()._get_schema_fields()
        fields[self._group_variable] = self._infer_column_type(
            self._fields_metadata,
            self._group_variable,
        )
        return fields

    def _prepare_grouped_df(
        self,
        df: pl.DataFrame,
        tr: Translations,
    ) -> tuple[pl.DataFrame, str, str, list[str]]:
        """
        Sort/translate group and category variables, concatenate categories, compute category totals,
        and optionally normalise values for relative display.

        :param df: Raw DataFrame from the query result.
        :type df: pl.DataFrame
        :param tr: Translations for localizing labels.
        :type tr: Translations
        :return: Processed DataFrame, category alias, effective value column name, and sorted group list.
        :rtype: tuple[pl.DataFrame, str, str, list[str]]
        """
        df = self._sort_and_translate(
            df,
            self._group_variable,
            self._fields_metadata[self._group_variable],
            tr,
            sort_asc=True,
        )
        for category_variable in reversed(self._category_variables):
            df = self._sort_and_translate(
                df,
                category_variable,
                self._fields_metadata[category_variable],
                tr,
                sort_asc=True,
            )
        df, alias = self._concatenate_variables(df, self._category_variables)
        value_var = self._aggregation.name
        category_total_var = "category_total"
        total_per_category = df.group_by(alias).agg(pl.sum(value_var).alias(category_total_var))
        df = df.join(total_per_category, on=alias, how="left")
        relative = self._relative and self._stacked
        if relative and df.height > 0:
            new_value_var = f"Relative {value_var}"
            df = df.with_columns((pl.col(value_var) / pl.col(category_total_var)).alias(new_value_var))
            value_var = new_value_var
        sorted_groups = df[self._group_variable].unique(maintain_order=True).to_list()
        return df, alias, value_var, sorted_groups

    def _compute_group_color_map(
        self,
        sorted_groups: list[str],
        theme: Theme,
    ) -> dict[str, str] | None:
        """
        Build a ``color_discrete_map`` dict from color rules, or return ``None`` for the default palette.

        :param sorted_groups: Ordered list of group values.
        :type sorted_groups: list[str]
        :param theme: Theme providing the base color.
        :type theme: Theme
        :return: Hex color map keyed by group value, or ``None``.
        :rtype: dict[str, str] | None
        """
        if not self._color_rules:
            return None
        group_df = pl.DataFrame({self._group_variable: sorted_groups})
        n_groups = len(sorted_groups)
        group_colors: HSLColorSequence = [theme.get_color("primary_color_normal")] * n_groups
        for color_rule in self._color_rules:
            selector = color_rule.selector.selector(
                group_df,
                self._group_variable,
                self._aggregation.name,
            )
            assigned_colors = color_rule.assigner.assign_colors(
                group_df,
                self._group_variable,
                self._aggregation.name,
                self._group_variable,
            )
            group_colors = Palette.mix_color_sequences(
                group_colors,
                assigned_colors,
                selector,
                indexing=Palette.MixIndexing.BY_POSITION,
            )
        return {str(grp): group_colors[i].to_css_hex() for i, grp in enumerate(sorted_groups)}

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df, alias, value_var, sorted_groups = self._prepare_grouped_df(dfs["main"], tr)
        color_map = self._compute_group_color_map(sorted_groups, theme)
        return {
            "categories": df[alias].unique(maintain_order=True).to_list(),
            "groups": sorted_groups,
            "color_map": color_map,
            "_df": df,
            "_alias": alias,
            "_value_var": value_var,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        df = payload["_df"]
        alias = payload["_alias"]
        value_var = payload["_value_var"]
        color_map = payload["color_map"]
        relative = self._relative and self._stacked
        x_var = alias
        y_var = value_var
        if self._horizontal:
            x_var, y_var = y_var, x_var
        extra_hover_data: list[str] = ["category_total"]
        if relative:
            extra_hover_data.append(self._aggregation.name)
        bar_kwargs: dict[str, Any] = dict(
            barmode="stack" if self._stacked else "group",
            orientation="h" if self._horizontal else "v",
            hover_data=extra_hover_data,
        )
        if color_map is not None:
            fig = px.bar(
                df.to_pandas(),
                x=x_var,
                y=y_var,
                color=self._group_variable,
                color_discrete_map=color_map,
                **bar_kwargs,
            )
        else:
            fig = px.bar(
                df.to_pandas(),
                x=x_var,
                y=y_var,
                color=self._group_variable,
                **bar_kwargs,
            )
        effective_title = payload["title"] or (
            f"Grouped bar chart of {self._aggregation.name} by {', '.join(self._category_variables)} and {self._group_variable}"
        )
        effective_category = self._category_label or ", ".join(self._category_variables)
        effective_value = self._value_label or self._append_unit(self._aggregation.default_label, self._aggregation.field)
        effective_group = self._group_label or self._group_variable
        if self._horizontal:
            fig.update_layout(xaxis_title=effective_value, yaxis_title=effective_category)
        else:
            fig.update_layout(xaxis_title=effective_category, yaxis_title=effective_value)
        fig.update_layout(legend_title_text=effective_group)
        if relative:
            if self._horizontal:
                fig.update_xaxes(tickformat=".1%", range=[0, 1])
            else:
                fig.update_yaxes(tickformat=".1%", range=[0, 1])
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

    @classmethod
    def supported_color_selectors(cls) -> list[type[ColorSelectorStrategy]]:
        return [
            ColorSelectorStrategy,
            ColorSelectByCategoryName,
        ]

    @classmethod
    def supported_color_assigners(cls) -> list[type[ColorAssignerStrategy]]:
        return [
            ColorAssignerByGroup,
        ]


__all__ = [
    "GroupedBarChartStrategy",
]
