"""
Pie chart figure strategy implementation.
"""

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

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


class PieChartStrategy(FigureStrategy):
    """
    Pie/Doughnut chart figure strategy.
    """

    _variables: ChartVariablesMap = {
        "category": CategoricalVariable(
            name="Category",
            description=(
                "Categorical variable(s) for the pie chart. "
                "Each category corresponds to a pie sector. "
                "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 angle of each sector."
            ),
            map_to="aggregation",
        ),
    }
    _options: ChartOptionsMap = {
        "sort_by_value": SortOrderOption(
            name="Sort by value",
            description=(
                "Sorting order of the pie sectors according to their value. "
                "If none, sectors are sorted according to the category variable(s)."
            ),
            map_to="sort_by_value",
            default_value=SortOrder.NONE,
        ),
        "hole": FloatOption(
            name="Hole size",
            description=(
                "Size of the hole in the center of the pie chart, as a fraction of the pie "
                "radius. A value of 0 produces a standard pie chart, while a value close to 1 "
                "produces a thin ring chart."
            ),
            map_to="hole",
            default_value=0.0,
            min_value=0.0,
            max_value=0.9,
        ),
        "inner_overall": FloatOption(
            name="Show overall value inside",
            description=(
                "Whether to show the overall value in the center of the pie chart when there is a "
                "hole, and the font size to use in that case. The overall is calculated by using "
                "the aggregation function over all the documents that apply."
            ),
            map_to="inner_overall",
            default_value=0.0,
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        category_variables: list[str],
        aggregation: Aggregation,
        sort_by_value: SortOrder = SortOrder.NONE,
        color_rules: list[ColorRule] | None = None,
        hole: float = 0.0,
        inner_overall: float = 0.0,
        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._options["sort_by_value"].validate(sort_by_value)
        self._options["hole"].validate(hole)
        self._options["inner_overall"].validate(inner_overall)

        self._category_variables = category_variables
        self._aggregation = aggregation
        self._sort_by_value = sort_by_value
        self._hole = hole
        self._inner_overall = inner_overall
        self._category_label = category_label
        self._value_label = value_label

    @classmethod
    def name(cls) -> str:
        return "Pie/Doughnut Chart"

    @classmethod
    def description(cls) -> str:
        return (
            "A pie chart is a circular statistical graphic, which is divided into slices to "
            "illustrate numerical proportions. In a pie chart, the arc length of each slice is "
            "proportional to the quantity it represents. A doughnut chart is a variant of the pie "
            "chart with a hole in the center."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        query_params: dict[str, QueryParams] = {
            "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(),
                ),
            )
        }
        if self._inner_overall > 0.0 and self._hole > 0.0:
            query_params["overall"] = QueryParams(
                group_fields=[],
                aggregations=[
                    TherismosAggregation(
                        id="overall",
                        function=AggregationFunction(self._aggregation.function),
                        field=self._aggregation.field,
                    ),
                ],
                dataframe_schema=pl.Schema(
                    {
                        "overall": (pl.Int64 if self._aggregation.function == AggregationFunction.COUNT else pl.Float64),
                    }
                ),
            )
        return query_params

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

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df: pl.DataFrame = dfs["main"]
        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)
        if self._sort_by_value != SortOrder.NONE:
            df = df.sort(
                self._aggregation.name,
                descending=self._sort_by_value == SortOrder.DESCENDING,
            )
        sector_colors: HSLColorSequence
        if self._color_rules:
            sector_colors = [theme.get_color("primary_color_normal")] * df.height
            for color_rule in self._color_rules:
                selector = color_rule.selector.selector(df, alias, self._aggregation.name)
                assigned_colors = color_rule.assigner.assign_colors(df, alias, self._aggregation.name)
                sector_colors = Palette.mix_color_sequences(
                    sector_colors, assigned_colors, selector, indexing=Palette.MixIndexing.BY_POSITION
                )
        else:
            sector_colors = ColorAssignerByCategory(anchor_colors=theme.get_default_color_sequence(), shuffle=True).assign_colors(
                df, alias, self._aggregation.name
            )
        return {
            "categories": df[alias].to_list(),
            "values": df[self._aggregation.name].to_list(),
            "colors": [c.to_css_hex() for c in sector_colors],
            "_alias": alias,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        categories = payload["categories"]
        values = payload["values"]
        colors = payload["colors"]
        fig = go.Figure(
            data=[
                go.Pie(
                    labels=categories,
                    values=values,
                    marker=dict(colors=colors),
                    sort=False,
                    rotation=0,
                    direction="counterclockwise",
                    hole=self._hole,
                )
            ],
        )
        effective_title = payload["title"] or (f"Pie chart of {self._aggregation.name} by {', '.join(self._category_variables)}")
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        if self._inner_overall > 0.0 and self._hole > 0.0:
            dfs = payload["dfs"]
            fig.add_annotation(
                text=(
                    f"{dfs['overall']['overall'][0]:.2f}"
                    if self._aggregation.function != AggregationFunction.COUNT
                    else f"{int(dfs['overall']['overall'][0])}"
                ),
                x=0.5,
                y=0.5,
                font_size=self._inner_overall,
                showarrow=False,
            )
        return fig

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

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


__all__ = [
    "PieChartStrategy",
]
