from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from importlib.util import find_spec
from typing import TYPE_CHECKING, Any, Final, Iterable, Sequence, TypeVar

if TYPE_CHECKING:
    from ropt.enums import ResultAxis
    from ropt.transforms._transforms import OptModelTransforms


_HAVE_PANDAS: Final = find_spec("pandas") is not None

if TYPE_CHECKING and _HAVE_PANDAS:
    import pandas as pd  # noqa: TC002
if _HAVE_PANDAS:
    from ._pandas import _to_dataframe

TypeResults = TypeVar("TypeResults", bound="Results")


@dataclass(slots=True)
class Results(ABC):
    """Abstract base class for storing optimization results.

    The `Results` class serves as a foundation for storing various types of
    optimization results. It is not intended to be instantiated directly but
    rather serves as a base for derived classes like
    [`FunctionResults`][ropt.results.FunctionResults] and
    [`GradientResults`][ropt.results.GradientResults], which hold the actual
    data.

    This class provides storage for the following generic information:

    *   **Batch ID:** An optional identifier, potentially generated by the
        function evaluator, that uniquely identifies a group of function
        evaluations passed to the evaluator by teh optimizer.
    *   **Metadata:** A dictionary for storing additional information generated
        during optimization. This metadata can include various primitive values
        that are not directly interpreted by the optimization code but are
        useful for reporting and analysis.

    The derived classes, [`FunctionResults`][ropt.results.FunctionResults] and
    [`GradientResults`][ropt.results.GradientResults], extend this base class
    with specific attributes for storing function evaluation results and
    gradient evaluation results, respectively. These derived classes also
    provide methods for exporting the stored data.

    One key method provided by the `Results` class is
    [`to_dataframe`][ropt.results.Results.to_dataframe], which allows exporting
    the contents of a specific field, or a subset of its sub-fields, to a
    `pandas` DataFrame for further data analysis and reporting.

    Attributes:
        batch_id: The ID of the evaluation batch.
        metadata: A dictionary of metadata.
    """

    batch_id: int | None
    metadata: dict[str, Any]

    def to_dataframe(
        self,
        field_name: str,
        select: Iterable[str],
        unstack: Iterable[ResultAxis] | None = None,
        names: dict[str, Sequence[str | int] | None] | None = None,
    ) -> pd.DataFrame:
        """Export a field to a pandas DataFrame.

        Exports the values of a single field to a `pandas` DataFrame.
        The field to export is selected by the `field_name` argument.
        Typically, such a field contains multiple sub-fields. By default,
        all sub-fields are exported as columns in the DataFrame, but a
        subset can be selected using the `select` argument.

        Sub-fields may be multi-dimensional arrays, which are exported in a
        stacked manner. Using the axis types found in the metadata, the exporter
        constructs a multi-index labeled with the corresponding names provided
        via the `names` argument. If `names` is `None`, numerical indices are
        used. These multi-indices can optionally be unstacked into multiple
        columns by providing the axis types to unstack via the `unstack`
        argument.

        Info: The DataFrame Index
            The index of the resulting DataFrame may be a multi-index
            constructed from axis indices or labels. In addition, the `batch_id`
            (if not `None`) is prepended to the index.

        Args:
            field_name: The field to export.
            select:     Select the sub-fields to export. By default, all
                        sub-fields are exported.
            unstack:    Select axes to unstack. By default, no axes are
                        unstacked.
            names:      A dictionary mapping axis types to names.

        Raises:
            NotImplementedError: If the `pandas` module is not installed.
            ValueError:          If the field name is incorrect.

        Returns:
            A `pandas` DataFrame containing the results.

        Warning:
            This function requires the `pandas` module to be installed.
        """
        if not _HAVE_PANDAS:
            msg = "The pandas module must be installed to use to_dataframe"
            raise NotImplementedError(msg)

        result_field = getattr(self, field_name, None)
        if result_field is None:
            msg = f"Invalid result field: {field_name}"
            raise AttributeError(msg)

        return _to_dataframe(
            result_field,
            self.batch_id,
            select,
            unstack,
            names,
        )

    @abstractmethod
    def transform_from_optimizer(self, transforms: OptModelTransforms) -> Results:
        """Transform results from the optimizer domain to the user domain.

        During optimization, variables, objectives, and constraints are often
        transformed to a different domain (the optimizer domain) to enhance
        the performance and stability of the optimization algorithm. The
        [`Results`][ropt.results.Results] objects produced during optimization
        are initially in the optimizer domain. This method reverses these
        transformations, mapping the results back to the user-defined domain.
        The transformations between the user and optimizer domains are defined
        by the classes in the [`ropt.transforms`][ropt.transforms] module.

        For instance, variables might have been scaled and shifted to a range
        more suitable for the optimizer. This method, using the provided
        `OptModelTransforms` object, applies the inverse scaling and shifting to
        restore the variables to their original scale and offset. Similarly,
        objectives and constraints are transformed back to the user domain.

        These transformations are defined and managed by the
        [`OptModelTransforms`][ropt.transforms.OptModelTransforms] object,
        which encapsulates the specific transformations for variables,
        objectives, and nonlinear constraints.

        Args:
            transforms: The `OptModelTransforms` object containing the domain
                transformations to apply.

        Returns:
            A new `FunctionResults` object with all relevant data transformed
            back to the user domain.
        """
