"""
Sunburst 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, 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,
    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 SunburstChartStrategy(FigureStrategy):
    """
    Sunburst chart figure strategy.
    """

    _variables: ChartVariablesMap = {
        "path": CategoricalVariable(
            name="Path",
            description=(
                "Categorical variable(s) for the concentrical sunburst chart. "
                "Each variable corresponds to a concentrical pie chart (innermost first). "
                "For each of the variables, each category corresponds to a sector of the pie "
                "chart. Each outer pie chart is divided according to the categories of the "
                "inner pie chart it surrounds."
            ),
            map_to="path_variables",
            multicategorical=True,
            ordered=True,
        ),
        "value": NumericalAggregationVariable(
            name="Value",
            description=(
                "An aggregation of a numerical variable to produce the final value of each "
                "categories' path. This value is proportional to the angle of each sector "
                "in the outermost circle."
            ),
            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,
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        path_variables: list[str],
        aggregation: Aggregation,
        sort_by_value: SortOrder = SortOrder.NONE,
        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,
        path_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["path"].validate_fields(fields_metadata, *path_variables)
        if aggregation.field:
            self._variables["value"].validate_fields(fields_metadata, aggregation.field)

        self._options["sort_by_value"].validate(sort_by_value)

        self._path_variables = path_variables
        self._aggregation = aggregation
        self._sort_by_value = sort_by_value
        self._path_label = path_label
        self._value_label = value_label

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

    @classmethod
    def description(cls) -> str:
        return (
            "A sunburst chart is a multi-level pie chart that visualizes hierarchical data. Each "
            "level of the hierarchy is represented by a ring, with the innermost ring being the "
            "top of the hierarchy."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        return {
            "main": QueryParams(
                group_fields=copy.copy(self._path_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._path_variables},
            self._aggregation.name: (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._path_variables):
            df = self._sort_and_translate(
                df,
                category_variable,
                self._fields_metadata[category_variable],
                tr,
                sort_asc=True,
            )
        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, self._path_variables[0], self._aggregation.name)
                assigned_colors = color_rule.assigner.assign_colors(df, self._path_variables[0], 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, self._path_variables[0], self._aggregation.name)

        return {
            "categories": df[self._path_variables[0]].to_list(),
            "values": df[self._aggregation.name].to_list(),
            "colors": [c.to_css_hex() for c in sector_colors],
            "_df": df,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        df = payload["_df"]
        colors = payload["colors"]
        fig = px.sunburst(
            df.to_pandas(),
            path=self._path_variables,
            values=self._aggregation.name,
            color=self._path_variables[0],
            color_discrete_map={cat: c for cat, c in zip(df[self._path_variables[0]], colors)},
        )
        fig.update_traces(sort=False)
        effective_title = payload["title"] or (f"Sunburst chart of {self._aggregation.name} by {', '.join(self._path_variables)}")
        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,
        ]

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


__all__ = [
    "SunburstChartStrategy",
]
