"""
Grouped boxplot chart figure strategy implementation.
"""

from typing import Any

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.boxplot import AbstractBoxplotChartStrategy
from orama.theme import Theme
from orama.variables import (
    CategoricalVariable,
    ChartVariablesMap,
    NonDerivedNumericalVariable,
)


class GroupedBoxplotChartStrategy(AbstractBoxplotChartStrategy):
    """
    Grouped boxplot chart figure strategy.

    Extends the standard boxplot chart with an additional grouping dimension, producing
    side-by-side boxplots per category — one per unique value of the grouping variable.
    """

    _variables: ChartVariablesMap = {
        "category": CategoricalVariable(
            name="Category",
            description=(
                "Categorical variable(s) for the grouped boxplot chart. "
                "Each category corresponds to a set of boxplots. "
                "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": NonDerivedNumericalVariable(
            name="Value",
            description="A numerical variable to extract the boxplot statistics from for each category.",
            map_to="value_variable",
        ),
        "grouping": CategoricalVariable(
            name="Grouping",
            description=(
                "Categorical variable for grouping the boxplots. "
                "Categories are split into several groups, one per each value of this variable."
            ),
            map_to="group_variable",
            multicategorical=False,
        ),
    }
    _options: ChartOptionsMap = {
        "horizontal": BooleanOption(
            name="Horizontal boxplots",
            description="When enabled, boxplots are rendered horizontally instead of vertically.",
            map_to="horizontal",
            default_value=False,
        ),
        "display_mean": BooleanOption(
            name="Display mean",
            description="Whether to display the mean value in each boxplot.",
            map_to="display_mean",
            default_value=False,
        ),
        "display_sd": BooleanOption(
            name="Display standard deviation",
            description=("Whether to display the standard deviation in each boxplot. Only applies if mean is also displayed."),
            map_to="display_sd",
            default_value=False,
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        category_variables: list[str],
        value_variable: str,
        group_variable: str,
        horizontal: bool = False,
        display_mean: bool = False,
        display_sd: 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,
            value_variable=value_variable,
            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["horizontal"].validate(horizontal)
        self._options["display_mean"].validate(display_mean)
        self._options["display_sd"].validate(display_sd)

        self._group_variable = group_variable
        self._horizontal = horizontal
        self._display_mean = display_mean
        self._display_sd = display_sd & self._display_mean
        self._group_label = group_label

    @classmethod
    def name(cls) -> str:
        """Return the human-readable name of this strategy."""
        return "Grouped Boxplot Chart"

    @classmethod
    def description(cls) -> str:
        """Return a description of this strategy."""
        return (
            "A grouped boxplot chart represents statistical distribution of samples segregated by "
            "categories, with an additional grouping dimension. For each category, multiple "
            "side-by-side boxplots are displayed — one per unique value of the grouping variable. "
            "Each boxplot shows the first quartile (Q1), median (Q2), and third quartile (Q3) of "
            "the samples, as well as the computed fences (whiskers) at 1.5 times the "
            "interquartile range (IQR) from the mean."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        """
        Return the query parameters for this strategy.

        :return: Dictionary mapping query names to their parameters.
        :rtype: dict[str, QueryParams]
        """
        params = super().get_query_params()
        params["main"].group_fields.append(self._group_variable)
        return 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_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df: pl.DataFrame = self._compute_fences(dfs["main"])
        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)

        sorted_groups = df[self._group_variable].unique(maintain_order=True).to_list()
        group_df = pl.DataFrame({self._group_variable: sorted_groups})
        n_groups = len(sorted_groups)

        group_colors: HSLColorSequence
        if self._color_rules:
            group_colors = [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, "mean")
                assigned_colors = color_rule.assigner.assign_colors(group_df, self._group_variable, "mean", self._group_variable)
                group_colors = Palette.mix_color_sequences(
                    group_colors,
                    assigned_colors,
                    selector,
                    indexing=Palette.MixIndexing.BY_POSITION,
                )
        else:
            group_colors = ColorAssignerByGroup(
                anchor_colors=theme.get_default_color_sequence(),
            ).assign_colors(group_df, self._group_variable, "mean", self._group_variable)

        return {
            "categories": df[alias].unique(maintain_order=True).to_list(),
            "groups": sorted_groups,
            "color_map": {str(g): group_colors[i].to_css_hex() for i, g in enumerate(sorted_groups)},
            "_df": df,
            "_alias": alias,
            "_group_colors": group_colors,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        df = payload["_df"]
        alias = payload["_alias"]
        group_colors = payload["_group_colors"]
        sorted_groups = payload["groups"]
        data = []
        for i, group in enumerate(sorted_groups):
            group_data = df.filter(pl.col(self._group_variable) == group)
            box_kwargs: dict[str, Any] = dict(
                name=str(group),
                q1=group_data["q1"].to_list(),
                median=group_data["median"].to_list(),
                q3=group_data["q3"].to_list(),
                lowerfence=group_data["lowerfence"].to_list(),
                upperfence=group_data["upperfence"].to_list(),
                boxpoints="outliers",
                marker=dict(color=group_colors[i].to_css_hex()),
                showlegend=True,
            )
            if self._display_mean:
                box_kwargs["mean"] = group_data["mean"].to_list()
                if self._display_sd:
                    box_kwargs["sd"] = group_data["sd"].to_list()
            if self._horizontal:
                data.append(go.Box(y=group_data[alias].to_list(), orientation="h", **box_kwargs))
            else:
                data.append(go.Box(x=group_data[alias].to_list(), **box_kwargs))
        effective_title = payload["title"] or (
            f"Grouped boxplot of {self._value_variable} 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._value_variable, self._value_variable)
        effective_group = self._group_label or self._group_variable
        fig = go.Figure(data=data)
        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(boxmode="group", legend_title_text=effective_group)
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

    @classmethod
    def supported_color_selectors(cls) -> list[type[ColorSelectorStrategy]]:
        """Return the list of supported color selector strategies."""
        return [
            ColorSelectorStrategy,
            ColorSelectByCategoryName,
        ]

    @classmethod
    def supported_color_assigners(cls) -> list[type[ColorAssignerStrategy]]:
        """Return the list of supported color assigner strategies."""
        return [
            ColorAssignerByGroup,
        ]


__all__ = ["GroupedBoxplotChartStrategy"]
