"""
Categorical line chart 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 plotly.subplots import make_subplots
from polychromos.palette import HSLColorSequence, Palette
from therismos.grouping import Aggregation as TherismosAggregation
from therismos.grouping import AggregationFunction

from orama.color import (
    ColorAssignerByCategory,
    ColorAssignerStrategy,
    ColorRule,
    ColorSelectByCategoryName,
    ColorSelectorStrategy,
)
from orama.enums import LineInterpolation, LinePlacement, LineSortOrder
from orama.options import ChartOptionsMap, FloatOption, IntOption, LineInterpolationOption, LinePlacementOption, LineSortOrderOption
from orama.query import QueryParams
from orama.strategies.base import FigureStrategy, _hex_to_rgb
from orama.theme import Theme
from orama.variables import (
    Aggregation,
    CategoricalVariable,
    ChartVariablesMap,
    NumericalAggregationVariable,
)


class CategoricalLineChartStrategy(FigureStrategy):
    """
    Categorical line chart figure strategy.

    Plots one or more lines over an ordered categorical X axis (e.g. year, quarter).
    An optional series variable splits the data into separate lines.
    """

    _variables: ChartVariablesMap = {
        "x": CategoricalVariable(
            name="X axis",
            description=(
                "Ordered categorical variable(s) for the X axis (e.g. year, quarter). "
                "Variables are concatenated in order to form the X labels."
            ),
            map_to="x_variables",
            multicategorical=True,
            ordered=True,
        ),
        "series": CategoricalVariable(
            name="Series",
            description=("Categorical variable for splitting data into separate lines. Leave as none for a single line."),
            map_to="series_variable",
            multicategorical=False,
            optional=True,
        ),
        "value": NumericalAggregationVariable(
            name="Value",
            description="Aggregation of a numerical variable for the Y axis.",
            map_to="aggregation",
        ),
    }
    _options: ChartOptionsMap = {
        "fill_opacity": FloatOption(
            name="Fill opacity",
            description="Opacity of the filled area under each line. 0 disables the fill.",
            map_to="fill_opacity",
            default_value=0.0,
            min_value=0.0,
            max_value=1.0,
        ),
        "line_width": IntOption(
            name="Line width",
            description="Width of each line in pixels.",
            map_to="line_width",
            default_value=3,
            min_value=1,
        ),
        "placement": LinePlacementOption(
            name="Placement",
            description="How the series lines/areas are positioned relative to each other.",
            map_to="placement",
            default_value=LinePlacement.OVERLAPPING,
        ),
        "interpolation": LineInterpolationOption(
            name="Interpolation",
            description="Line interpolation mode between data points.",
            map_to="interpolation",
            default_value=LineInterpolation.LINEAR,
        ),
        "sort_order": LineSortOrderOption(
            name="Sort series",
            description="Order in which the category lines are drawn.",
            map_to="sort_order",
            default_value=LineSortOrder.NONE,
        ),
    }

    def __init__(
        self,
        fields_metadata: dict[str, FieldMetadata],
        x_variables: list[str],
        aggregation: Aggregation,
        series_variable: str | None = None,
        fill_opacity: float = 0.0,
        line_width: int = 3,
        placement: LinePlacement = LinePlacement.OVERLAPPING,
        interpolation: LineInterpolation = LineInterpolation.LINEAR,
        sort_order: LineSortOrder = LineSortOrder.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,
        x_label: str | None = None,
        value_label: str | None = None,
        series_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"].validate_fields(fields_metadata, *x_variables)
        if series_variable:
            self._variables["series"].validate_fields(fields_metadata, series_variable)
        if aggregation.field:
            self._variables["value"].validate_fields(fields_metadata, aggregation.field)

        self._options["fill_opacity"].validate(fill_opacity)
        self._options["line_width"].validate(line_width)
        self._options["placement"].validate(placement)
        self._options["interpolation"].validate(interpolation)
        self._options["sort_order"].validate(sort_order)

        self._x_variables = x_variables
        self._series_variable = series_variable
        self._aggregation = aggregation
        self._fill_opacity = fill_opacity
        self._line_width = line_width
        self._placement = placement
        self._interpolation = interpolation
        self._sort_order = sort_order
        self._x_label = x_label
        self._value_label = value_label
        self._series_label = series_label

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

    @classmethod
    def description(cls) -> str:
        return (
            "A categorical line chart presents data as connected lines over an ordered "
            "categorical X axis (e.g. year, quarter). Lines can be split by a series variable."
        )

    def get_query_params(self) -> dict[str, QueryParams]:
        """
        Returns query parameters for the main grouped aggregation.

        :return: Dict with a single ``"main"`` key mapping to its :class:`QueryParams`.
        :rtype: dict[str, QueryParams]
        """
        group_fields = list(self._x_variables)
        if self._series_variable:
            group_fields.append(self._series_variable)
        return {
            "main": QueryParams(
                group_fields=group_fields,
                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]]:
        schema: dict[str, type[pl.DataType]] = {var: self._infer_column_type(self._fields_metadata, var) for var in self._x_variables}
        if self._series_variable:
            schema[self._series_variable] = self._infer_column_type(self._fields_metadata, self._series_variable)
        schema[self._aggregation.name] = pl.Int64 if self._aggregation.function == AggregationFunction.COUNT else pl.Float64
        return schema

    def _prepare_plot_data(
        self,
        dfs: dict[str, pl.DataFrame],
        tr: Translations,
        theme: Theme,
    ) -> dict[str, Any]:
        df: pl.DataFrame = dfs["main"]
        for x_var in reversed(self._x_variables):
            df = self._sort_and_translate(df, x_var, self._fields_metadata[x_var], tr, sort_asc=True)
        df, x_alias = self._concatenate_variables(df, self._x_variables)
        series_values: list[object] = (
            df[self._series_variable].unique(maintain_order=True).to_list() if self._series_variable else [None]
        )
        if self._sort_order != LineSortOrder.NONE and self._series_variable is not None:
            if self._sort_order == LineSortOrder.ALPHABETICAL_ASC:
                series_values = sorted(series_values, key=lambda v: str(v) if v is not None else "")
            elif self._sort_order == LineSortOrder.ALPHABETICAL_DESC:
                series_values = sorted(series_values, key=lambda v: str(v) if v is not None else "", reverse=True)
            elif self._sort_order in (LineSortOrder.FIRST_VALUE_ASC, LineSortOrder.FIRST_VALUE_DESC):

                def _first(sv: object) -> float:
                    sub = df.filter(pl.col(self._series_variable) == sv)  # type: ignore[arg-type]
                    vals = sub[self._aggregation.name].drop_nulls()
                    return float(vals[0]) if vals.len() > 0 else 0.0

                series_values = sorted(series_values, key=_first, reverse=(self._sort_order == LineSortOrder.FIRST_VALUE_DESC))
            elif self._sort_order in (LineSortOrder.LAST_VALUE_ASC, LineSortOrder.LAST_VALUE_DESC):

                def _last(sv: object) -> float:
                    sub = df.filter(pl.col(self._series_variable) == sv)  # type: ignore[arg-type]
                    vals = sub[self._aggregation.name].drop_nulls()
                    return float(vals[-1]) if vals.len() > 0 else 0.0

                series_values = sorted(series_values, key=_last, reverse=(self._sort_order == LineSortOrder.LAST_VALUE_DESC))
        sv_col = self._series_variable or "__series__"
        series_df = (
            pl.DataFrame({self._series_variable: series_values}) if self._series_variable else pl.DataFrame({"__series__": [None]})
        )
        line_colors: HSLColorSequence
        if self._color_rules:
            line_colors = [theme.get_color("primary_color_normal")] * len(series_values)
            for color_rule in self._color_rules:
                selector = color_rule.selector.selector(series_df, sv_col, self._aggregation.name)
                assigned = color_rule.assigner.assign_colors(series_df, sv_col, self._aggregation.name)
                line_colors = Palette.mix_color_sequences(line_colors, assigned, selector, indexing=Palette.MixIndexing.BY_POSITION)
        else:
            line_colors = ColorAssignerByCategory(
                anchor_colors=theme.get_default_color_sequence(),
            ).assign_colors(series_df, sv_col, self._aggregation.name)
        return {
            "x_labels": df[x_alias].unique(maintain_order=True).to_list(),
            "series": [str(sv) if sv is not None else None for sv in series_values],
            "colors": [c.to_css_hex() for c in line_colors],
            "_df": df,
            "_x_alias": x_alias,
            "_series_values": series_values,
            "_line_colors": line_colors,
        }

    def _render_from_payload(self, payload: dict[str, Any]) -> go.Figure:
        df = payload["_df"]
        x_alias = payload["_x_alias"]
        series_values = payload["_series_values"]
        line_colors = payload["_line_colors"]
        effective_title = payload["title"] or (f"Line chart of {self._aggregation.name} by {' - '.join(self._x_variables)}")
        effective_x = self._x_label or " - ".join(self._x_variables)
        effective_value = self._value_label or self._append_unit(self._aggregation.default_label, self._aggregation.field)
        if self._placement == LinePlacement.RIDGELINE:
            fig = make_subplots(rows=len(series_values), cols=1, shared_xaxes=True, vertical_spacing=0.02)
            for i, sv in enumerate(series_values):
                line_df = df.filter(pl.col(self._series_variable) == sv) if sv is not None else df  # type: ignore[arg-type]
                r, g, b = _hex_to_rgb(line_colors[i].to_css_hex())
                fill_kwargs: dict[str, Any] = {}
                if self._fill_opacity > 0:
                    fill_kwargs = dict(fill="tozeroy", fillcolor=f"rgba({r},{g},{b},{self._fill_opacity})")
                fig.add_trace(
                    go.Scatter(
                        x=line_df[x_alias].to_list(),
                        y=line_df[self._aggregation.name].to_list(),
                        mode="lines",
                        line=dict(color=line_colors[i].to_css_hex(), width=self._line_width, shape=self._interpolation.value),
                        name=str(sv) if sv is not None else self._aggregation.name,
                        **fill_kwargs,
                    ),
                    row=i + 1,
                    col=1,
                )
        else:
            traces = []
            for i, sv in enumerate(series_values):
                line_df = df.filter(pl.col(self._series_variable) == sv) if sv is not None else df  # type: ignore[arg-type]
                r, g, b = _hex_to_rgb(line_colors[i].to_css_hex())
                scatter_kwargs: dict[str, Any] = dict(
                    x=line_df[x_alias].to_list(),
                    y=line_df[self._aggregation.name].to_list(),
                    mode="lines",
                    line=dict(color=line_colors[i].to_css_hex(), width=self._line_width, shape=self._interpolation.value),
                    name=str(sv) if sv is not None else self._aggregation.name,
                )
                if self._placement == LinePlacement.OVERLAPPING:
                    if self._fill_opacity > 0:
                        scatter_kwargs["fill"] = "tozeroy"
                        scatter_kwargs["fillcolor"] = f"rgba({r},{g},{b},{self._fill_opacity})"
                elif self._placement == LinePlacement.STACKED:
                    scatter_kwargs["stackgroup"] = "one"
                    scatter_kwargs["fillcolor"] = f"rgba({r},{g},{b},{self._fill_opacity})"
                else:
                    scatter_kwargs["stackgroup"] = "one"
                    scatter_kwargs["groupnorm"] = "percent"
                    scatter_kwargs["fillcolor"] = f"rgba({r},{g},{b},{self._fill_opacity})"
                traces.append(go.Scatter(**scatter_kwargs))
            fig = go.Figure(data=traces)
        if self._placement == LinePlacement.RIDGELINE:
            n_rows = len(series_values)
            fig.update_xaxes(title_text=effective_x, row=n_rows, col=1)
            fig.update_xaxes(showgrid=False)
            fig.update_yaxes(showgrid=False, showticklabels=False)
            fig.update_layout(hovermode="x unified", showlegend=False)
            for i, sv in enumerate(series_values):
                axis_key = "yaxis" if i == 0 else f"yaxis{i + 1}"
                domain = fig.layout[axis_key].domain
                y_center = (domain[0] + domain[1]) / 2
                fig.add_annotation(
                    text=str(sv) if sv is not None else self._aggregation.name,
                    xref="paper",
                    yref="paper",
                    x=1,
                    y=y_center,
                    xanchor="left",
                    yanchor="middle",
                    showarrow=False,
                    font=dict(color=line_colors[i].to_css_hex()),
                )
            top_domain = fig.layout["yaxis"].domain[1]
            last_axis_key = "yaxis" if n_rows == 1 else f"yaxis{n_rows}"
            bottom_domain = fig.layout[last_axis_key].domain[0]
            y_all_center = (top_domain + bottom_domain) / 2
            fig.add_annotation(
                text=effective_value,
                xref="paper",
                yref="paper",
                x=-0.07,
                y=y_all_center,
                xanchor="center",
                yanchor="middle",
                textangle=-90,
                showarrow=False,
            )
        else:
            fig.update_layout(
                xaxis_title=effective_x,
                yaxis_title=effective_value,
                hovermode="x unified",
                xaxis=dict(showgrid=False),
            )
            if self._series_variable is not None:
                fig.update_layout(legend_title_text=self._series_label or self._series_variable)
        self._apply_title(fig, effective_title, payload.get("subtitle"))
        return fig

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

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


__all__ = [
    "CategoricalLineChartStrategy",
]
