from __future__ import annotations

from abc import ABC
from dataclasses import dataclass
from importlib.util import find_spec
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, Final, Iterable, Optional, TypeVar, Union

if TYPE_CHECKING:
    from ropt.config.enopt import EnOptConfig
    from ropt.enums import ResultAxisName


_HAVE_PANDAS: Final = find_spec("pandas") is not None
_HAVE_XARRAY: Final = find_spec("xarray") is not None
_HAVE_NETCDF: Final = find_spec("netCDF4") is not None

if TYPE_CHECKING and _HAVE_PANDAS:
    import pandas as pd  # noqa: TCH002
if _HAVE_PANDAS:
    from ._pandas import _to_dataframe
if TYPE_CHECKING and _HAVE_XARRAY:
    import xarray  # noqa: TCH002
if _HAVE_XARRAY:
    from ._xarray import _to_dataset
if _HAVE_XARRAY and _HAVE_NETCDF:
    from ._xarray import _to_netcdf

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


@dataclass
class Results(ABC):
    """The `Results` class serves as an abstract base class for storing results.

    This class contains the following generic information:

    1. The result ID, a unique integer value that is incremented as the
       results become available.
    2. An optional batch ID, which may be generated by the function evaluator.
       The interpretation of this ID depends on the evaluator code. It is
       intended to be used as a unique identifier for the group of function
       evaluations passed to the evaluator by the optimization code.
    3. A dictionary of metadata to be added by optimization steps. This contains
       generic information, the nature of which depends on the steps producing
       them. They are expected to be primitive values not interpreted by the
       optimization code but can be exported and reported.

    The `Results` class is an abstract base class that is not intended to be
    instantiated by itself. Most data of interest will be stored in additional
    fields in one of the derived classes:
    [`FunctionResults`][ropt.results.FunctionResults] or
    [`GradientResults`][ropt.results.GradientResults]. In addition to the
    attributes containing the data, a few methods are provided to export the
    data. These functions will only be useful when used with one of the derived
    classes:

    1. The [`to_dataframe`][ropt.results.Results.to_dataframe] method can be
       used to export the contents, or a sub-set, of a single field to a
       [`pandas`](https://pandas.pydata.org/) data frame.
    2. The [`to_dataset`][ropt.results.Results.to_dataset] method can be used to
       export the contents, or a sub-set, of a single field to an
       [`xarray`](https://xarray.dev/) data set.
    2. The [`to_netcdf`][ropt.results.Results.to_netcdf] method can be used to
       export the entire results object to a
       [`netCDF`](https://www.unidata.ucar.edu/software/netcdf/) file.

    Attributes:
        result_id: The ID of the function/gradient evaluation.
        batch_id:  The ID of the evaluation batch that contains the result.
        metadata:  The metadata.
    """

    result_id: int
    batch_id: Optional[int]
    metadata: Dict[str, Any]

    def to_dataframe(
        self,
        config: EnOptConfig,
        field_name: str,
        select: Optional[Iterable[str]] = None,
        unstack: Optional[Iterable[ResultAxisName]] = None,
    ) -> pd.DataFrame:
        """Export a field to a pandas dataframe.

        The function exports the values of a single field to a pandas data
        frame. The field to export is selected by the `field_name` argument. In
        general such a field is another object with multiple sub-fields. By
        default, these are all exported as columns in the pandas data frame, but
        a sub-set can be selected using the `select` argument.

        Any of the sub-fields in the field that is exported may be a
        multi-dimensional array, which is exported in a stacked manner. Using
        the axis types found in the metadata, the exporter will construct a
        multi-index labeled with the corresponding names found in the optimizer
        configuration (if available, otherwise numerical indices are used). Such
        multi-indices can optionally be unstacked into multiple columns by
        providing the axis types to unstack via the `unstack` argument.

        Info: The data frame index
            As noted above, the index of the resulting data frame may be a
            multi-index constructed from axis indices or labels. In addition,
            the `result_id` field, and the `batch_id` (if not None), are also
            prepended to the index of the resulting frame. This has the
            beneficial effect that the data frames exported from multiple
            results can be concatenated and identified in the frame by result ID
            and/or batch ID.

        Args:
            config:     The ensemble optimizer configuration object
            field_name: The field to export
            select:     Select the sub-fields to export, by default all fields
            unstack:    Select axes to unstack, by default none

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

        Returns:
            A pandas data frame containing the results.

        Warning:
            This function is only available if `pandas` is 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 ValueError(msg)

        return _to_dataframe(
            config,
            result_field,
            self.result_id,
            self.batch_id,
            select,
            unstack,
        )

    def to_dataset(
        self,
        config: EnOptConfig,
        field_name: str,
        select: Optional[Iterable[str]] = None,
        *,
        add_metadata: bool = False,
    ) -> xarray.Dataset:
        """Export a field to an xarray dataset.

        The function exports the values of a single field to an xarray dataset.
        The field to export is selected by the `field_name` argument. In
        general, such a field is another object with multiple sub-fields. By
        default, these are all exported as data arrays in the xarray dataset,
        but a subset can be selected using the `select` argument.

        Any of the sub-fields in the field that is exported may be a
        multi-dimensional array, which is exported directly as an xarray data
        array. Using the axis types found in the metadata, the exporter will add
        coordinate labels constructed from the corresponding names found in the
        optimizer configuration (if available).

        The result ID and the batch ID (if not None) are added to the attributes
        of the dataset, using the `result_id` and `batch_id` keys. If metadata
        is present in the results, and the `add_metadata` flag is set, it is
        also added under the `metadata` key.

        Args:
            config:       The ensemble optimizer configuration object
            field_name:   The field to export
            select:       Select the fields to export; by default, all fields
            add_metadata: If true, add the metadata as a field in the dataset attrs

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

        Returns:
            An xarray dataset containing the results.

        Warning:
            This function is only available if `xarray` is installed.
        """
        if not _HAVE_XARRAY:
            msg = "The xarray module must be installed to use to_dataset"
            raise NotImplementedError(msg)

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

        return _to_dataset(
            config,
            result_field,
            self.result_id,
            self.batch_id,
            self.metadata if add_metadata else {},
            select,
        )

    def to_netcdf(self, config: EnOptConfig, filename: Union[str, Path]) -> None:
        """Write the results to a netCDF4 file.

        The fields of the result are converted to xarray datasets and each
        stored as a group in a netCDF4 file, using the name of the field as the
        group name. In addition, a special `__metadata__` group is added that
        contains the result ID, the batch ID (if not None), and a JSON dump of
        the metadata.

        By default, a file is written with a `.nc` extension. The filename may
        contain a replacement field for a variable named `result_id` or
        `batch_id`, which will be replaced by the corresponding field using the
        `format` string method. For example:
        `filename="results-{result_id:04d}"` will result in a file named
        `results-0001.nc` if the result ID equals 1.

        Info: Reading results from a netCDF4 file
            The results may be read back using the `from_netcdf` class method of
            either [`FunctionResults`][ropt.results.FunctionResults] or
            [`GradientResults`][ropt.results.GradientResults]. Currently, no
            information is written about whether the results were function or
            gradient results. Therefore, to read back the result with the
            correct class, its type must be known.

        Args:
            config:   The configuration used to run the optimization.
            filename: The name of the file to write.

        Raises:
            NotImplementedError:
                If the `xarray` or the `netCDF4` module is not installed.

        Warning:
            Use of this method requires that the `xarray` and `netCDF4` modules
            are installed.
        """
        if not _HAVE_XARRAY or not _HAVE_NETCDF:
            msg = (
                "The xarray and netCDF4 modules must be installed "
                "to use Results.to_netcdf"
            )
            raise NotImplementedError(msg)
        _to_netcdf(self, config, Path(filename))
