"""
Boxplot 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,
    ColorSelectValuesAboveMean,
    ColorSelectValuesBelow,
    ColorSelectValuesBelowMean,
)
from orama.options import (
    BooleanOption,
    ChartOptionsMap,
    StringOption,
)
from orama.query import QueryParams
from orama.strategies.base import FigureStrategy
from orama.theme import Theme
from orama.variables import (
    CategoricalVariable,
    ChartVariablesMap,
    NonDerivedNumericalVariable,
)


class AbstractBoxplotChartStrategy(FigureStrategy):
    """
    Abstract base class for boxplot chart figure strategies.
    """

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        category_variables: list[str],
        value_variable: str,
        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)
        self._variables["value"].validate_fields(fields_metadata, value_variable)

        self._category_variables = category_variables
        self._value_variable = value_variable
        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="q1",
                        function=AggregationFunction.Q1,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="median",
                        function=AggregationFunction.MEDIAN,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="q3",
                        function=AggregationFunction.Q3,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="mean",
                        function=AggregationFunction.AVERAGE,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="sd",
                        function=AggregationFunction.STDDEV,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="min",
                        function=AggregationFunction.MIN,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="max",
                        function=AggregationFunction.MAX,
                        field=self._value_variable,
                    ),
                ],
                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},
            "q1": pl.Float64,
            "median": pl.Float64,
            "q3": pl.Float64,
            "mean": pl.Float64,
            "sd": pl.Float64,
            "min": pl.Float64,
            "max": pl.Float64,
        }

    def _compute_fences(
        self,
        df: pl.DataFrame,
    ) -> pl.DataFrame:
        iqr = df["q3"] - df["q1"]
        lower_fence: pl.Series = pl.Series(
            name="lowerfence",
            values=[max(df["min"][i], df["q1"][i] - 1.5 * iqr[i]) for i in range(df.height)],
            dtype=pl.Float64,
        )
        upper_fence: pl.Series = pl.Series(
            name="upperfence",
            values=[min(df["max"][i], df["q3"][i] + 1.5 * iqr[i]) for i in range(df.height)],
            dtype=pl.Float64,
        )
        return df.with_columns(
            lower_fence,
            upper_fence,
        )


class BoxplotChartStrategy(AbstractBoxplotChartStrategy):
    """
    Boxplot chart figure strategy.
    """

    _variables: ChartVariablesMap = {
        "category": CategoricalVariable(
            name="Category",
            description=(
                "Categorical variable(s) for the boxplot chart. "
                "Each category corresponds to a boxplot. "
                "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",
        ),
    }
    _options: ChartOptionsMap = {
        "horizontal": BooleanOption(
            name="Horizontal bars",
            description="When enabled, bars are rendered horizontally instead of vertically.",
            map_to="horizontal",
            default_value=False,
        ),
        "overall_category_name": StringOption(
            name="Overall category name",
            description=(
                "When set, an additional boxplot representing the overall distribution "
                "of the values (regardless of category) is displayed under the category named here."
            ),
            map_to="overall_category_name",
            default_value="",
        ),
        "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,
        horizontal: bool = False,
        overall_category_name: str = "",
        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,
    ) -> 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._options["horizontal"].validate(horizontal)
        self._options["overall_category_name"].validate(overall_category_name)
        self._options["display_mean"].validate(display_mean)
        self._options["display_sd"].validate(display_sd)

        self._horizontal = horizontal
        self._overall_category_name = overall_category_name
        self._display_mean = display_mean
        self._display_sd = display_sd & self._display_mean

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

    @classmethod
    def description(cls) -> str:
        return (
            "A boxplot chart represents statistical distribution of samples segregated by "
            "categories. 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]:
        params = super().get_query_params()
        if self._overall_category_name:
            params["overall"] = QueryParams(
                group_fields=[],
                aggregations=[
                    TherismosAggregation(
                        id="q1",
                        function=AggregationFunction.Q1,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="median",
                        function=AggregationFunction.MEDIAN,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="q3",
                        function=AggregationFunction.Q3,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="mean",
                        function=AggregationFunction.AVERAGE,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="sd",
                        function=AggregationFunction.STDDEV,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="min",
                        function=AggregationFunction.MIN,
                        field=self._value_variable,
                    ),
                    TherismosAggregation(
                        id="max",
                        function=AggregationFunction.MAX,
                        field=self._value_variable,
                    ),
                ],
                dataframe_schema=pl.Schema(
                    {
                        "q1": pl.Float64,
                        "median": pl.Float64,
                        "q3": pl.Float64,
                        "mean": pl.Float64,
                        "sd": pl.Float64,
                        "min": pl.Float64,
                        "max": pl.Float64,
                    }
                ),
            )
        return params

    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"])
        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._overall_category_name:
            overall_df: pl.DataFrame = self._compute_fences(dfs["overall"])
            overall_df = overall_df.with_columns(
                pl.lit(self._overall_category_name).cast(pl.Utf8).alias(alias),
            )
            df = pl.concat(
                [
                    df,
                    overall_df.select(df.columns),
                ],
                how="vertical",
            )

        boxplot_colors: HSLColorSequence
        if self._color_rules:
            boxplot_colors = [theme.get_color("primary_color_normal")] * df.height
            for color_rule in self._color_rules:
                selector = color_rule.selector.selector(df, alias, "mean")
                assigned_colors = color_rule.assigner.assign_colors(df, alias, "mean")
                boxplot_colors = Palette.mix_color_sequences(
                    boxplot_colors,
                    assigned_colors,
                    selector,
                    indexing=Palette.MixIndexing.BY_POSITION,
                )
        else:
            boxplot_colors = ColorAssignerByCategory(
                anchor_colors=theme.get_default_color_sequence(),
            ).assign_colors(df, alias, "mean")

        return {
            "categories": df[alias].to_list(),
            "statistics": df.select(["q1", "median", "q3", "mean", "sd", "min", "max", "lowerfence", "upperfence"]).to_dicts(),
            "colors": [c.to_css_hex() for c in boxplot_colors],
            "_df": df,
            "_alias": alias,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        df = payload["_df"]
        alias = payload["_alias"]
        colors = payload["colors"]
        data = []
        for i, row in enumerate(df.iter_rows(named=True)):
            box_kwargs: dict[str, Any] = dict(
                name=row[alias],
                q1=[row["q1"]],
                median=[row["median"]],
                q3=[row["q3"]],
                lowerfence=[row["lowerfence"]],
                upperfence=[row["upperfence"]],
                boxpoints="outliers",
                marker=dict(color=colors[i]),
            )
            if self._display_mean:
                box_kwargs["mean"] = [row["mean"]]
                if self._display_sd:
                    box_kwargs["sd"] = [row["sd"]]
            if self._horizontal:
                data.append(go.Box(y=[row[alias]], orientation="h", **box_kwargs))
            else:
                data.append(go.Box(x=[row[alias]], **box_kwargs))
        fig = go.Figure(data=data)
        effective_title = payload["title"] or (f"Boxplot of {self._value_variable} 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._value_variable, self._value_variable)
        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)
        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__ = [
    "AbstractBoxplotChartStrategy",
    "BoxplotChartStrategy",
]
