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]]

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

class HomogenousIterationLabel:
    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:
    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:
        return self._colormap

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

    @property
    def free_y_axis(self) -> bool:
        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:
        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()))


