"""
Bar chart figure strategy implementation.
"""

import copy
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
from therismos.grouping import Aggregation as TherismosAggregation
from therismos.grouping import AggregationFunction

from orama.color import (
    ColorAssignerByCategory,
    ColorAssignerByValue,
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectByCategoryName,
    ColorSelectMaxValue,
    ColorSelectMinValue,
    ColorSelectorStrategy,
    ColorSelectValuesAbove,
    ColorSelectValuesAboveMean,
    ColorSelectValuesBelow,
    ColorSelectValuesBelowMean,
)
from orama.enums import SortOrder
from orama.options import (
    BooleanOption,
    ChartOptionsMap,
    FloatOption,
    IntOption,
    SortOrderOption,
    StringOption,
)
from orama.query import QueryParams
from orama.strategies.base import FigureStrategy
from orama.theme import Theme
from orama.variables import (
    Aggregation,
    CategoricalVariable,
    ChartVariablesMap,
    NumericalAggregationVariable,
)

_OTHERS_SIMPLE_FUNCTIONS = {
    AggregationFunction.COUNT,
    AggregationFunction.SUM,
    AggregationFunction.MIN,
    AggregationFunction.MAX,
}


class AbstractBarChartStrategy(FigureStrategy):
    """
    Abstract base class for bar chart figure strategies.
    """

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        category_variables: list[str],
        aggregation: Aggregation,
        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,
    ) -> None:
        super().__init__(
            fields_metadata=fields_metadata,
            color_rules=color_rules,
            title=title,
            subtitle=subtitle,
            caption=caption,
            description=description,
            storytelling_context=storytelling_context,
        )

        self._variables["category"].validate_fields(fields_metadata, *category_variables)
        if aggregation.field:
            self._variables["value"].validate_fields(fields_metadata, aggregation.field)

        self._category_variables = category_variables
        self._aggregation = aggregation
        self._category_label = category_label
        self._value_label = value_label

    def get_query_params(self) -> dict[str, QueryParams]:
        return {
            "main": QueryParams(
                group_fields=copy.copy(self._category_variables),
                aggregations=[
                    TherismosAggregation(
                        id=self._aggregation.name,
                        function=AggregationFunction(self._aggregation.function),
                        field=self._aggregation.field,
                    ),
                ],
                dataframe_schema=pl.Schema(
                    self._get_schema_fields(),
                ),
            )
        }

    def _get_schema_fields(self) -> dict[str, type[pl.DataType]]:
        return {
            **{var: self._infer_column_type(self._fields_metadata, var) for var in self._category_variables},
            self._aggregation.name: (pl.Int64 if self._aggregation.function == AggregationFunction.COUNT else pl.Float64),
        }


class BarChartStrategy(AbstractBarChartStrategy):
    """
    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",
        ),
    }
    _options: ChartOptionsMap = {
        "sort_by_value": SortOrderOption(
            name="Sort by value",
            description=(
                "Sorting order of the bars according to their value. If none, bars are sorted according to the category variable(s)."
            ),
            map_to="sort_by_value",
            default_value=SortOrder.NONE,
        ),
        "horizontal": BooleanOption(
            name="Horizontal bars",
            description="When enabled, bars are rendered horizontally instead of vertically.",
            map_to="horizontal",
            default_value=False,
        ),
        "min_value": FloatOption(
            name="Minimum value",
            description="Minimum value for the bar chart value axis.",
            map_to="min_value",
            default_value=None,
        ),
        "max_value": FloatOption(
            name="Maximum value",
            description="Maximum value for the bar chart value axis.",
            map_to="max_value",
            default_value=None,
        ),
        "mean_category_name": StringOption(
            name="Mean category name",
            description=(
                "The name to use for the category that represents the mean of every category. "
                "If not provided, the mean category is not shown. When the value aggregation "
                "function is the count, the mean category represents the average count per "
                "category. Otherwise, it represents the mean value weighted by the count of each "
                "category."
            ),
            map_to="mean_category_name",
            default_value="",
        ),
        "mean_annotation": BooleanOption(
            name="Mean annotation",
            description=(
                "When enabled, shows an annotation with the mean value on the bar chart. "
                "Only applicable if a mean category name is provided. The mean value is "
                "calculated as the overall average when the value aggregation function is the "
                "count; otherwise, it is the mean value weighted by the count of each category."
            ),
            map_to="mean_annotation",
            default_value=False,
        ),
        "others_threshold": FloatOption(
            name="Others threshold",
            description=(
                "When set, all bars whose value is strictly below this threshold are grouped into a "
                "single bar. Only applicable when sorting is active. "
                "Mutually exclusive with 'others_top_n'."
            ),
            map_to="others_threshold",
            default_value=None,
        ),
        "others_top_n": IntOption(
            name="Others top N",
            description=(
                "When set, only the top N bars by value are shown individually; all remaining bars "
                "are grouped into a single bar. Only applicable when sorting is active. "
                "Mutually exclusive with 'others_threshold'."
            ),
            map_to="others_top_n",
            min_value=1,
            default_value=None,
        ),
        "others_category_name": StringOption(
            name="Others category name",
            description="Label for the grouped bar. Defaults to 'Others'.",
            map_to="others_category_name",
            default_value="Others",
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        category_variables: list[str],
        aggregation: Aggregation,
        sort_by_value: SortOrder = SortOrder.NONE,
        horizontal: bool = False,
        min_value: float | None = None,
        max_value: float | None = None,
        mean_category_name: str = "",
        mean_annotation: bool = False,
        others_threshold: float | None = None,
        others_top_n: int | None = None,
        others_category_name: str = "Others",
        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,
    ) -> 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._options["sort_by_value"].validate(sort_by_value)
        self._options["horizontal"].validate(horizontal)
        if min_value is not None:
            self._options["min_value"].validate(min_value)
        if max_value is not None:
            self._options["max_value"].validate(max_value)
        if mean_category_name is not None:
            self._options["mean_category_name"].validate(mean_category_name)
        self._options["mean_annotation"].validate(mean_annotation)
        if others_threshold is not None and others_top_n is not None:
            raise ValueError("'others_threshold' and 'others_top_n' are mutually exclusive")
        if others_threshold is not None and sort_by_value == SortOrder.NONE:
            raise ValueError("'others_threshold' requires sorting to be active (sort_by_value != SortOrder.NONE)")
        if others_top_n is not None and sort_by_value == SortOrder.NONE:
            raise ValueError("'others_top_n' requires sorting to be active (sort_by_value != SortOrder.NONE)")
        if others_category_name and mean_category_name and others_category_name == mean_category_name:
            raise ValueError(f"'others_category_name' and 'mean_category_name' must not be the same: '{others_category_name}'")
        if others_threshold is not None:
            self._options["others_threshold"].validate(others_threshold)
        if others_top_n is not None:
            self._options["others_top_n"].validate(others_top_n)
        self._options["others_category_name"].validate(others_category_name)

        self._sort_by_value = sort_by_value
        self._horizontal = horizontal
        self._min_value = min_value
        self._max_value = max_value
        self._mean_category_name = mean_category_name
        self._mean_annotation = mean_annotation
        self._others_threshold = others_threshold
        self._others_top_n = others_top_n
        self._others_category_name = others_category_name

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

    @classmethod
    def description(cls) -> str:
        return (
            "A bar chart represents categorical data with rectangular bars with heights or "
            "lengths proportional to the values that they represent. The bars can be plotted "
            "vertically or horizontally."
        )

    @property
    def _requires_mean(self) -> bool:
        return any(
            [
                self._mean_category_name,
                self._mean_annotation,
                *[
                    color_rule.selector.__class__
                    in {
                        ColorSelectValuesBelowMean,
                        ColorSelectValuesAboveMean,
                    }
                    for color_rule in (self._color_rules or [])
                ],
            ]
        )

    @property
    def _requires_total_count(self) -> bool:
        others_need_count = (
            self._others_threshold is not None or self._others_top_n is not None
        ) and self._aggregation.function not in _OTHERS_SIMPLE_FUNCTIONS
        return self._requires_mean or others_need_count

    def get_query_params(self) -> dict[str, QueryParams]:
        query_params = super().get_query_params()
        if self._requires_total_count:
            query_params["main"].aggregations.append(
                TherismosAggregation(
                    id="__total_count__",
                    function=AggregationFunction.COUNT,
                ),
            )
        return query_params

    def _get_schema_fields(self) -> dict[str, type[pl.DataType]]:
        fields = super()._get_schema_fields()
        if self._requires_total_count:
            fields["__total_count__"] = pl.Int64
        return fields

    def _apply_others_grouping(self, df: pl.DataFrame, alias: str) -> pl.DataFrame:
        """
        Collapses low-signal bars into a single "Others" bar.

        :param df: The DataFrame to process.
        :type df: pl.DataFrame
        :param alias: The column name for the category axis.
        :type alias: str
        :return: The processed DataFrame with an "Others" bar appended if applicable.
        :rtype: pl.DataFrame
        """
        agg_name = self._aggregation.name
        if self._mean_category_name:
            mean_rows = df.filter(pl.col(alias) == self._mean_category_name)
            data_rows = df.filter(pl.col(alias) != self._mean_category_name)
        else:
            mean_rows = df.head(0)
            data_rows = df

        if self._others_threshold is not None:
            others_rows = data_rows.filter(pl.col(agg_name) < self._others_threshold)
            kept_rows = data_rows.filter(pl.col(agg_name) >= self._others_threshold)
        else:
            assert self._others_top_n is not None
            sorted_desc = data_rows.sort(agg_name, descending=True)
            kept_rows_unordered = sorted_desc.head(self._others_top_n)
            others_rows = sorted_desc.slice(self._others_top_n)
            kept_indices = kept_rows_unordered[alias].to_list()
            kept_rows = data_rows.filter(pl.col(alias).is_in(kept_indices))

        if others_rows.is_empty():
            return pl.concat([data_rows, mean_rows])

        fn = self._aggregation.function
        if fn in (AggregationFunction.COUNT, AggregationFunction.SUM):
            others_value = others_rows[agg_name].sum()
        elif fn == AggregationFunction.MIN:
            others_value = others_rows[agg_name].min()  # type: ignore[assignment]
        elif fn == AggregationFunction.MAX:
            others_value = others_rows[agg_name].max()  # type: ignore[assignment]
        else:
            weighted_sum = (others_rows[agg_name] * others_rows["__total_count__"]).sum()
            total_weight = others_rows["__total_count__"].sum()
            others_value = weighted_sum / total_weight  # type: ignore[operator]

        others_row_data: dict[str, list[Any]] = {
            alias: [self._others_category_name],
            agg_name: [others_value],
        }
        if "__total_count__" in df.columns:
            others_row_data["__total_count__"] = [others_rows["__total_count__"].sum()]
        if "__total_mean__" in df.columns:
            others_row_data["__total_mean__"] = [df["__total_mean__"][0]]

        others_row = pl.DataFrame(others_row_data).select(df.columns)
        return pl.concat([kept_rows, others_row, mean_rows], how="vertical_relaxed")

    def _prepare_df(
        self,
        df: pl.DataFrame,
        tr: Translations,
    ) -> tuple[pl.DataFrame, str, float | None]:
        """
        Sort, translate, compute the mean row (if needed), sort by value, and apply others grouping.

        :param df: Raw DataFrame from the query result.
        :type df: pl.DataFrame
        :param tr: Translations for localizing category labels.
        :type tr: Translations
        :return: Processed DataFrame, category alias column name, and optional mean value.
        :rtype: tuple[pl.DataFrame, str, float | None]
        """
        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)
        mean_value: float | None = None
        if self._requires_mean:
            df = df.with_columns(pl.col(self._aggregation.name).cast(pl.Float64))
            total_count_value = df["__total_count__"].sum()
            if self._aggregation.function == AggregationFunction.COUNT:
                mean_value = float(df[self._aggregation.name].mean())  # type: ignore[arg-type]
            else:
                weighted_sum = float((df[self._aggregation.name] * df["__total_count__"]).sum())
                mean_value = float(weighted_sum / total_count_value)  # type: ignore[operator]
            if self._mean_category_name:
                df = pl.concat(
                    [
                        df,
                        pl.DataFrame(
                            {
                                alias: [self._mean_category_name],
                                self._aggregation.name: [mean_value],
                                "__total_count__": [total_count_value],
                            }
                        ).select(df.columns),
                    ],
                    how="vertical_relaxed",
                )
            df = df.with_columns(pl.lit(mean_value).alias("__total_mean__"))
        if self._sort_by_value != SortOrder.NONE:
            df = df.sort(
                self._aggregation.name,
                descending=self._sort_by_value == SortOrder.DESCENDING,
            )
        if self._others_threshold is not None or self._others_top_n is not None:
            df = self._apply_others_grouping(df, alias)
        return df, alias, mean_value

    def _compute_bar_colors(
        self,
        df: pl.DataFrame,
        alias: str,
        theme: Theme,
    ) -> HSLColorSequence:
        """
        Produce the per-bar color sequence via color rules or the default category assigner.

        :param df: Prepared DataFrame (post :meth:`_prepare_df`).
        :type df: pl.DataFrame
        :param alias: Category axis column name.
        :type alias: str
        :param theme: Theme providing color sequences.
        :type theme: Theme
        :return: A color sequence with one entry per bar.
        :rtype: HSLColorSequence
        """
        if self._color_rules:
            return self._apply_color_rules(df, alias, self._aggregation.name, theme)
        return ColorAssignerByCategory(
            anchor_colors=theme.get_default_color_sequence(),
        ).assign_colors(df, alias, self._aggregation.name)

    def _configure_bar_layout(
        self,
        fig: go.Figure,
        mean_value: float | None,
        effective_category: str,
        effective_value: str,
        annotation_color: str | None,
    ) -> None:
        """
        Apply axis titles, value range, legend visibility, and optional mean annotation.

        :param fig: Plotly figure to update.
        :type fig: go.Figure
        :param mean_value: Overall mean used for the annotation line, or ``None``.
        :type mean_value: float | None
        :param effective_category: Resolved category axis label.
        :type effective_category: str
        :param effective_value: Resolved value axis label.
        :type effective_value: str
        :param annotation_color: Hex color for the mean annotation line, or ``None``.
        :type annotation_color: str | None
        """
        fig.update_layout(showlegend=False)
        value_range = (self._min_value, self._max_value)
        if self._horizontal:
            fig.update_layout(xaxis_title=effective_value, yaxis_title=effective_category)
            fig.update_xaxes(range=value_range)
            if self._mean_annotation and mean_value is not None:
                fig.add_vline(
                    x=mean_value,
                    line_color=annotation_color,
                    line_width=3,
                    opacity=0.7,
                    layer="below",
                )
        else:
            fig.update_layout(xaxis_title=effective_category, yaxis_title=effective_value)
            fig.update_yaxes(range=value_range)
            if self._mean_annotation and mean_value is not None:
                fig.add_hline(
                    y=mean_value,
                    line_color=annotation_color,
                    line_width=3,
                    opacity=0.7,
                    layer="below",
                )

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df, alias, mean_value = self._prepare_df(dfs["main"], tr)
        bar_colors = self._compute_bar_colors(df, alias, theme)
        return {
            "categories": df[alias].to_list(),
            "values": df[self._aggregation.name].to_list(),
            "colors": [c.to_css_hex() for c in bar_colors],
            "mean_value": mean_value,
            "_alias": alias,
            "_annotation_color": theme.get_color("secondary_color_softer").to_css_hex(),
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        alias = payload["_alias"]
        categories = payload["categories"]
        values = payload["values"]
        colors = payload["colors"]
        mean_value = payload["mean_value"]
        df = pl.DataFrame({alias: categories, self._aggregation.name: values})
        x_var = alias
        y_var = self._aggregation.name
        if self._horizontal:
            x_var, y_var = y_var, x_var
        color_map = dict(zip(categories, colors))
        fig = px.bar(
            df.to_pandas(),
            x=x_var,
            y=y_var,
            color=alias,
            color_discrete_map=color_map,
            orientation="h" if self._horizontal else "v",
        )
        effective_title = payload["title"] or (f"Bar chart of {self._aggregation.name} by {', '.join(self._category_variables)}")
        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)
        self._configure_bar_layout(fig, mean_value, effective_category, effective_value, payload["_annotation_color"])
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

    @classmethod
    def supported_color_selectors(cls) -> list[type[ColorSelectorStrategy]]:
        return [
            ColorSelectorStrategy,
            ColorSelectByCategoryName,
            ColorSelectValuesBelow,
            ColorSelectValuesAbove,
            ColorSelectMaxValue,
            ColorSelectMinValue,
            ColorSelectValuesBelowMean,
            ColorSelectValuesAboveMean,
        ]

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


__all__ = [
    "AbstractBarChartStrategy",
    "BarChartStrategy",
]
