from collections.abc import Iterable, Callable, Sequence
from copy import deepcopy
from dataclasses import asdict
from warnings import warn
from toolz import dissoc

from altbacken.core.state import AnnealingState
from altbacken.optional.pandas import DataFrame


type NamingScheme = Callable[[int], str]

class TemplateNamingScheme:
    """Handles naming schemes based on a template.

    This class is used to generate names dynamically based on a provided template
    and an index. The template can include placeholders that will be replaced
    with indices when the class instance is called.

    Attributes:
        template (str): The naming template used for formatting.
    """
    def __init__(self, template: str):
        self._template: str = template

    def __call__(self, index: int) -> str:
        """
        Formats the given index using a predefined template and returns the formatted string.

        Args:
            index (int): The index to be formatted.

        Returns:
            str: The resulting formatted string.
        """
        return self._template.format(index)


class PredefinedNamingScheme:
    """A naming scheme that uses predefined names and falls back to placeholders.

    This class represents a naming scheme where a list of predefined names is used to generate names based
    on their index. If the provided index is out of the range of predefined names, a placeholder name is
    generated and returned.

    Attributes:
        _names (Sequence[str]): A tuple of predefined names used for the naming scheme.
    """
    def __init__(self, names: Iterable[str]):
        self._names: Sequence[str] = tuple(names)

    def __call__(self, index: int) -> str:
        """
        Calls the object as a function to retrieve a name from a predefined naming scheme
        or a placeholder if the index is out of range.

        Args:
            index (int): The index of the required name in the predefined naming scheme.

        Returns:
            str: The name corresponding to the given index or a placeholder if the index
            is out of range.
        """
        if index < len(self._names):
            return self._names[index]
        else:
            warn(UserWarning(f"Index {index} is out of range for predefined naming scheme. Will use placehoder."))
            return f"_{index}_"



class DataFrameTracer:
    """
    Tracks and logs states over iterations for simulated annealing runs.

    This class is designed to record and transform states during simulated annealing
    processes into a format compatible with a pandas DataFrame. It maintains a buffer
    of selected annealing states based on a frequency parameter and transforms the
    state data into useful representations. The naming scheme for the coordinates
    can be customized via different input formats.

    Attributes:
        naming_scheme (NamingScheme): Defines how coordinate names are generated, either
            through a template, predefined names, or a custom scheme.
        coordinates (list[str]): The names of the current solution coordinates, sorted
            and derived from the naming scheme.
        frame (DataFrame): A pandas DataFrame containing the transformed states with both
            current and best solutions and relevant data.
    """
    def __init__(self, naming_scheme: NamingScheme | str | Iterable | None = None, frequency: int = 1):
        if frequency <= 0:
            raise ValueError("Frequency must be positive")
        self._naming_scheme: NamingScheme = self._parse_naming_scheme(naming_scheme or TemplateNamingScheme("x_{}"))
        self._buffer: list[AnnealingState[Sequence[float | int]]] = []
        self._frequency: int = frequency

    @property
    def naming_scheme(self) -> NamingScheme:
        """
        Gets or sets the naming scheme used in the current context.

        This property provides access to the naming scheme which defines how
        naming conventions or structures are handled in a specific context.
        The naming scheme is based on the `NamingScheme` class.

        Returns:
            NamingScheme: The naming scheme currently in use.
        """
        return self._naming_scheme

    @naming_scheme.setter
    def naming_scheme(self, naming_scheme: NamingScheme | str | Iterable[str]) -> None:
        self._naming_scheme = self._parse_naming_scheme(naming_scheme)


    @property
    def coordinates(self) -> list[str]:
        """
        Gets the list of coordinates based on the current buffer state.

        If the buffer contains data, the function generates a list of coordinate
        strings using the naming scheme provided. The coordinates are sorted
        alphabetically. If the buffer is empty, an empty list is returned.

        Returns:
            list[str]: A sorted list of coordinate strings if the buffer contains data, or
            an empty list if the buffer is not populated.
        """
        if self._buffer:
            return sorted(self._naming_scheme(i) for i in range(len(self._buffer[0].best_solution)))
        else:
            return []

    @property
    def frame(self) -> DataFrame:
        """
        Returns a DataFrame representation of transformed coordinates from the buffer.

        This property retrieves a pandas DataFrame, where each row represents the
        transformed coordinates derived from the buffer.

        Returns:
            DataFrame: The DataFrame containing transformed coordinates.
        """
        return DataFrame(map(self._transform_coordinates, self._buffer))

    def __call__(self, state: AnnealingState[Sequence[float | int]]) -> None:
        if state.iteration == 0:
            self._buffer.clear()
        if state.iteration % self._frequency == 0:
            self._buffer.append(deepcopy(state))

    @classmethod
    def _parse_naming_scheme(cls, naming_scheme: NamingScheme | str | Iterable[str]) -> NamingScheme:
        match naming_scheme:
            case str(template): return TemplateNamingScheme(template)
            case Iterable(): return PredefinedNamingScheme(naming_scheme)
            case other: return other

    def _transform_coordinates(self, state: AnnealingState[Sequence[float | int]]) -> dict[str, float | int]:
        current_transformed: dict[str, float | int] = {
            self._naming_scheme(index): value for index, value in enumerate(state.current_solution)
        }
        best_transformed: dict[str, float | int] = {
            f'best_{self._naming_scheme(index)}': value for index, value in enumerate(state.best_solution)
        }

        return {
            **current_transformed,
            **best_transformed,
            **dissoc(asdict(state), "current_solution", "best_solution")
        }

