from collections.abc import Iterable, Callable, Sequence
from itertools import combinations
from math import ceil

from altbacken.core.state import AnnealingState
from altbacken.internal.tracer.dataframe import NamingScheme, DataFrameTracer
from altbacken.optional.plotnine import (
    ggplot, aes, facet_wrap, geom_point, ggtitle, scale_fill_continuous, geom_raster, coord_fixed, geom_line,
    geom_tile, theme_minimal, geom_density_2d, theme
)
from altbacken.optional.pandas import DataFrame, Series


type PhaseLabeler = Callable[[DataFrame], Series[str]]
#: Deduces the phase label based on the simulation output

_DEFAULT_LABELS: Sequence[str] = ("1. Exploration", "2. Transition", "3. Freezing")

class HomogenousIterationLabel:
    """Assigns labels to iterations in a DataFrame based on a predefined set of labels.

    This class is designed to map iteration numbers to specific labels in a
    consistent and controlled manner. It ensures that a repeated sequence of
    labels is applied to the range of iteration values in a given DataFrame.

    Attributes:
        labels (Iterable[str]): Predefined sequence of string labels used for
            mapping iteration numbers. Defaults to `_DEFAULT_LABELS` if not specified.
    """
    def __init__(self, labels: Iterable[str] =_DEFAULT_LABELS):
        self._labels: Sequence[str] = tuple(labels)


    def __call__(self, df: DataFrame) -> Series[str]:
        width: int = ceil(df.iteration.max() / len(self._labels))
        return df.iteration.map(lambda i: self._labels[min(i // width, len(self._labels) - 1)])

_DEFAULT_LABELER: PhaseLabeler = HomogenousIterationLabel()

class ShowcasePlot:
    """Class to visualize and analyze optimization processes.

    This class provides functionality for visualizing optimization runs via
    dynamic plots, enabling users to monitor optimization progress, phases,
    and variable exploration space interactively. It includes customizable
    visualization options such as colormap and axis scaling.

    Attributes:
        colormap (str): The colormap to use for visualizations.
        free_y_axis (bool): Whether the y-axis is free or constrained across
            facets.
        frame (DataFrame): The DataFrame containing the traced data from the
            optimization process.
    """
    def __init__(self, naming_scheme: NamingScheme | str | Iterable[str] | None = None, frequency: int = 1):
        self._frame_tracer: DataFrameTracer = DataFrameTracer(naming_scheme, frequency)
        self._colormap: str = "coolwarm"
        self._free_y_axis: bool = True

    def __call__(self, state: AnnealingState[Sequence[float | int]]) -> None:
        self._frame_tracer(state)

    @property
    def colormap(self) -> str:
        """
        Gets or sets the name of the colormap used.

        This property retrieves the name of the currently set colormap used
        within the object configuration. The colormap is used for visual
        representations and data coloring purposes.

        Returns:
            str: The name of the colormap.
        """
        return self._colormap

    @colormap.setter
    def colormap(self, colormap: str) -> None:
        self._colormap = colormap

    @property
    def free_y_axis(self) -> bool:
        """
        Indicates whether the Y-axis is free or constrained.

        This property returns a boolean value that signifies whether
        the Y-axis is in a free state (not constrained) or not, based
        on the internal state of the class.

        Returns:
            bool: A boolean value where `True` means the Y-axis is free,
            and `False` means it is constrained.
        """
        return self._free_y_axis

    @free_y_axis.setter
    def free_y_axis(self, free_y_axis: bool) -> None:
        self._free_y_axis = bool(free_y_axis)

    @property
    def frame(self) -> DataFrame:
        return self._frame_tracer.frame

    def plot(self, size: tuple[float, float] = (8.27, 11.69), label: PhaseLabeler = _DEFAULT_LABELER) -> ggplot:
        """
        Plots a ggplot visualization representing multiple phases and exploration data.

        This method generates a comprehensive visualization comprising multiple sub-plots
        depicting different phases and exploration-related metrics. It uses the provided
        data from the internal DataFrame and organizes the plots in a hierarchical manner.

        The sub-plots included in the visualization are:
        - A plot showing the current values across phases.
        - A plot showing the best values across phases.
        - An exploration plot.
        - A temperature-related plot.

        The overall figure size of the visualization can be customized, and an external
        phase labeling system can be supplied for labeling data.

        Args:
            size: The dimensions for the output figure (width, height) in inches. Defaults to A4 portrait size.
            label: A labeler instance used for annotating or categorizing phases within the data.

        Returns:
            ggplot: A ggplot object representing the entire visualization with all sub-plots combined.
        """
        df: DataFrame = self._frame_tracer.frame
        df["phase"] = label(df)
        phases: ggplot = self._create_phase_plot(df, "current_value") / self._create_phase_plot(df, "best_value")
        plot = phases / self._create_exploration_plot(df) / self._create_temperature_plot(df)
        return plot + theme(figure_size=size)

    def _create_phase_plot(self, frame: DataFrame, value: str) -> ggplot:
        scales: str = "free" if self._free_y_axis else "free_x"
        return ggplot(frame, aes(x="iteration", y=value, fill="temperature")) \
        + geom_line() \
        + facet_wrap("phase", ncol=len(frame.phase.unique()), scales=scales) \
        + scale_fill_continuous(cmap_name=self._colormap) \
        + ggtitle(f"Phase Plot for {value}")

    def _create_temperature_plot(self, frame: DataFrame) -> ggplot:
        return ggplot(frame, aes(x="iteration", y="temperature")) + geom_line() + ggtitle("Temperature Plot")

    def _create_exploration_plot(self, frame: DataFrame) -> ggplot:
        surfaces: list[Sequence[str]] = list(combinations(self._frame_tracer.coordinates, 2))
        space: DataFrame = DataFrame(
            {"x": row[surface[0]], "y": row[surface[1]], "value": row["current_value"], "surface": '-'.join(surface)}
            for surface in surfaces
            for _, row in frame.iterrows()
        )
        return ggplot(space, aes(x="x", y="y", fill="value", color="value")) \
        + geom_point() \
        + ggtitle("Explored functions space") \
        + coord_fixed() \
        + geom_density_2d() \
        + facet_wrap("surface", scales="free", ncol=len(frame.phase.unique()))


