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:
    def __init__(self, template: str):
        self._template: str = template

    def __call__(self, index: int) -> str:
        return self._template.format(index)


class PredefinedNamingScheme:
    def __init__(self, names: Iterable[str]):
        self._names: Sequence[str] = tuple(names)

    def __call__(self, index: int) -> str:
        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:
    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:
        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]:
        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:
        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")
        }

