"""
Heat map 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 HSLColorScale, HSLColorSequence, Palette
from therismos.grouping import Aggregation as TherismosAggregation
from therismos.grouping import AggregationFunction

from orama.color import (
    ColorAssignerByValue,
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectorStrategy,
)
from orama.options import (
    BooleanOption,
    ChartOptionsMap,
)
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 HeatMapChartStrategy(FigureStrategy):
    """
    Heat map figure strategy.
    """

    _variables: ChartVariablesMap = {
        "x_category": CategoricalVariable(
            name="X Category",
            description=(
                "Categorical variable(s) for the horizontal axis of the heat map. "
                "Each category corresponds to a heat map column. "
                "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="x_category_variables",
            multicategorical=True,
        ),
        "y_category": CategoricalVariable(
            name="Y Category",
            description=(
                "Categorical variable(s) for the vertical axis of the heat map. "
                "Each category corresponds to a heat map row. "
                "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="y_category_variables",
            multicategorical=True,
        ),
        "value": NumericalAggregationVariable(
            name="Value",
            description=("An aggregation of a numerical variable to produce the final value of each heat map cell."),
            map_to="aggregation",
        ),
    }
    _options: ChartOptionsMap = {
        "display_values": BooleanOption(
            name="Display values",
            description="Whether to display the numerical values in each cell of the heat map.",
            default_value=False,
            map_to="display_values",
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        x_category_variables: list[str],
        y_category_variables: list[str],
        aggregation: Aggregation,
        display_values: 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,
        x_label: str | None = None,
        y_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["x_category"].validate_fields(fields_metadata, *x_category_variables)
        self._variables["y_category"].validate_fields(fields_metadata, *y_category_variables)
        if aggregation.field:
            self._variables["value"].validate_fields(fields_metadata, aggregation.field)

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

        self._x_category_variables = x_category_variables
        self._y_category_variables = y_category_variables
        self._aggregation = aggregation
        self._display_values = display_values
        self._x_label = x_label
        self._y_label = y_label
        self._value_label = value_label

    @classmethod
    def name(cls) -> str:
        return "Heat Map"

    @classmethod
    def description(cls) -> str:
        return (
            "A heat map represents categorical data in a two-dimensional matrix. A color scale is "
            "used to represent the values of each cell in the matrix."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        return {
            "main": QueryParams(
                group_fields=[
                    *self._x_category_variables,
                    *self._y_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._x_category_variables},
            **{var: self._infer_column_type(self._fields_metadata, var) for var in self._y_category_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._x_category_variables):
            df = self._sort_and_translate(df, category_variable, self._fields_metadata[category_variable], tr, sort_asc=True)
        df, x_alias = self._concatenate_variables(df, self._x_category_variables)
        for category_variable in reversed(self._y_category_variables):
            df = self._sort_and_translate(df, category_variable, self._fields_metadata[category_variable], tr, sort_asc=True)
        df, y_alias = self._concatenate_variables(df, self._y_category_variables)
        color_sequence: HSLColorSequence
        if self._color_rules:
            color_sequence = self._color_rules[0].assigner.anchor_colors
        else:
            color_sequence = theme.get_default_color_sequence()
        color_scale: HSLColorScale = Palette.to_color_scale(color_sequence)
        return {
            "x_categories": df[x_alias].to_list(),
            "y_categories": df[y_alias].to_list(),
            "values": df[self._aggregation.name].to_list(),
            "_color_scale": color_scale,
            "_x_alias": x_alias,
            "_y_alias": y_alias,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        color_scale = payload["_color_scale"]
        effective_title = payload["title"] or (
            f"Heat map of {self._aggregation.name} "
            f"by {', '.join(self._x_category_variables)} "
            f"and {', '.join(self._y_category_variables)}"
        )
        effective_x = self._x_label or ", ".join(self._x_category_variables)
        effective_y = self._y_label or ", ".join(self._y_category_variables)
        effective_value = self._value_label or self._append_unit(self._aggregation.default_label, self._aggregation.field)
        fig = go.Figure(
            data=go.Heatmap(
                z=payload["values"],
                x=payload["x_categories"],
                y=payload["y_categories"],
                colorscale=Palette.color_scale_to_plotly(color_scale),
                texttemplate="%{z:.2f}" if self._display_values else None,
            )
        )
        fig.update_traces(colorbar_title_text=effective_value)
        fig.update_layout(xaxis_title=effective_x, yaxis_title=effective_y)
        for color_rule in self._color_rules:
            if isinstance(color_rule.assigner, ColorAssignerByValue) and color_rule.assigner.range_min is not None:
                fig.update_traces(zmin=color_rule.assigner.range_min)
                break
        for color_rule in self._color_rules:
            if isinstance(color_rule.assigner, ColorAssignerByValue) and color_rule.assigner.range_max is not None:
                fig.update_traces(zmax=color_rule.assigner.range_max)
                break
        for color_rule in self._color_rules:
            if isinstance(color_rule.assigner, ColorAssignerByValue) and color_rule.assigner.range_center_at is not None:
                fig.update_traces(zmid=color_rule.assigner.range_center_at)
                break
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

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

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


__all__ = [
    "HeatMapChartStrategy",
]
