"""Data output interface for Bayesian optimization results.

This module provides the `Dumper` class, which serves as the data export interface
for post-processing BOSS optimization results. It is typically used by the `PPMain`
object (defined in `boss/pp/pp_main.py`) to write numerical data to text files for
external analysis and archival purposes.

The module is part of the post-processing infrastructure alongside:
- `boss.pp.pp_main.PPMain`: Main post-processing orchestrator that coordinates
  data dumping and plotting
- `boss.pp.plotter.Plotter`: Creates visualization plots from the same results
- `boss.pp.mesh.Mesh`: Manages coordinate grids for function evaluation

The `Dumper` class exports various types of optimization data to human-readable
text files with descriptive headers:
- Hyperparameter evolution over iterations
- Gaussian process model predictions (mean and variance) on grids
- Acquisition function values on grids
- True objective function values (when available)
- Convergence diagnostics
- Predicted minimum locations and values

Typical Usage
-------------
The `Dumper` is usually created and used via `PPMain`:

>>> from boss.pp.pp_main import PPMain
>>> pp = PPMain(results, pp_models=True, pp_acq_funcs=True)
>>> pp.run()

This creates the `postprocessing/` directory with data files alongside plots.

Direct usage is also supported:

>>> from boss.io.dumper import Dumper
>>> dumper = Dumper.from_grid_slice(results)
>>> dumper.dump_model("model_output.dat", itr=-1)
>>> dumper.dump_hyperparams("hyperparams.dat")

Output Format
-------------
All output files are text files with:
- Header lines starting with '#' describing the data format
- Column-oriented data with scientific notation (default precision: 9 decimal places)
- Special formatting for iteration numbers (integer) and ensemble sizes (integer)

See Also
--------
boss.pp.pp_main.PPMain : Main post-processing interface
boss.pp.plotter.Plotter : Creates diagnostic plots
boss.bo.results.BOResults : Container for optimization results
"""
from pathlib import Path

import numpy as np

from boss.bo.results import BOResults
from boss.pp.mesh import Mesh


class Dumper:
    """Data export interface for writing Bayesian optimization results to text files.

    The `Dumper` class provides methods for exporting various optimization data to
    human-readable text files. It uses `Mesh` objects to define coordinate grids for
    evaluating and dumping GP models, acquisition functions, and true functions over
    the input space.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results object containing models, data, and settings.
    mesh : Mesh
        Mesh object defining the coordinate grid for function evaluation.

    Attributes
    ----------
    results : BOResults
        The BO results object.
    settings : Settings
        Settings object from the BO results.
    bounds : NDArray
        Optimization bounds from settings.
    mesh : Mesh
        The mesh object for coordinate grids.
    grid_slice : ArrayLike or None
        Grid specification from the mesh, if available.
    """
    def __init__(self, results: BOResults, mesh: Mesh) -> None:
        self.results = results
        self.settings = results.settings
        self.bounds = self.settings["bounds"]
        self.mesh = mesh
        try:
            self.grid_slice = mesh.get_grid_spec()
        except RuntimeError:
            self.grid_slice = None

    @classmethod
    def from_grid_slice(cls, results: BOResults, grid_slice=None):
        """Create a Dumper from a grid specification.

        This is a convenience constructor that creates a `Mesh` object from a grid
        specification before initializing the `Dumper`.

        Parameters
        ----------
        results : BOResults
            Bayesian optimization results object.
        grid_slice : ArrayLike or None, optional
            Grid specification for the mesh. If None, uses the grid spec from
            `results.settings["pp_model_slice"]`.

        Returns
        -------
        Dumper
            Initialized dumper instance.
        """
        if grid_slice is None:
            grid_slice = results.settings["pp_model_slice"]
        mesh = Mesh.from_grid_spec(results.settings["bounds"], grid_slice)
        self = cls(results, mesh)
        return self

    def _build_grid_header(self, name: str) -> str:
        shape = self.mesh.grid_pts
        header = f"# {name}, grid of "
        if len(shape) > 1:
            header += f"{shape[0]}x{shape[1]}="
        header += f"{np.prod(shape)} pts\n"
        return header

    def dump_columns(
        self,
        file_path: str | Path,
        col_data,
        header: str = "data",
        formats: str | dict[int, str] = "%18.9e",
    ) -> None:
        """Write column data to a text file with a header.

        This is the core utility method used by all other dump methods to write
        formatted numerical data to text files. It handles column stacking,
        format specification, and header writing.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file.
        col_data : tuple or list of array_like
            Tuple or list of column data arrays to write. Each element can be
            1D (single column) or 2D (multiple columns).
        header : str, default="data"
            Header string to write at the top of the file. Should start with '#'.
        formats : str or dict[int, str], default="%18.9e"
            Output format specification. If str, applies to all columns. If dict,
            maps column indices to format strings. Non-specified columns use
            "%18.8e" as default.
        """
        num_cols = len(col_data)
        if isinstance(formats, str):
            fmt = formats
        elif isinstance(formats, dict):
            fmt = []
            for k in range(num_cols):
                if k in formats.keys():
                    fstr = formats[k]
                else:
                    fstr = "%18.8e"
                col = np.asarray(col_data[k])
                if len(col.shape) == 1:
                    reps = 1
                else:
                    reps = col.shape[1]
                fmt.extend([fstr] * reps)

        data_combined = np.column_stack(col_data)
        with open(file_path, "w") as fd:
            fd.write(header + "\n")
            np.savetxt(fd, data_combined, fmt=fmt)

    def dump_hyperparams(self, file_path: str | Path) -> None:
        """Write Gaussian process hyperparameter evolution to a text file.

        Exports the GP kernel hyperparameters (variance and lengthscales) for each
        iteration to a multi-column text file. Each row contains the iteration number,
        ensemble size, and all hyperparameter values.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically "hyperparameters.dat").

        Notes
        -----
        Output file columns: iteration, ensemble_size, variance, lengthscale_1, ..., lengthscale_d
        where d is the dimensionality of the optimization problem.
        """
        res = self.results
        header = "Model hyperparameter values by iteration (iter npts parameters)"
        iters = np.array(list(res["model_params"].keys()))
        npts = res.batch_tracker.ensemble_sizes[iters]
        model_params = res["model_params"].to_array()
        formats = {0: "%d", 1: "%5d"}
        self.dump_columns(
            file_path, (iters, npts, model_params), header, formats=formats
        )

    def dump_model(self, file_path: str | Path, itr: int=-1) -> None:
        """Write Gaussian process model predictions to a text file.

        Evaluates the GP model (mean and standard deviation) over the mesh grid
        and writes the results to a multi-column text file. For problems with more
        than 2 active dimensions, inactive dimensions are fixed to the predicted
        minimum location.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically in "data_models/" subdirectory).
        itr : int, default=-1
            Iteration number for which to dump the model. Negative indices count
            from the end (e.g., -1 is the last iteration).

        Notes
        -----
        Output file columns: x1, [x2, ...], mu, nu
        where x1, x2, ... are coordinate values, mu is the predicted mean, and
        nu is the predicted standard deviation.

        The header also includes the model's hyperparameter values at the
        specified iteration.
        """
        self.mesh.fix_dim_preset(self.results, preset="min", itr=itr)
        mesher = self.mesh
        model = self.results.reconstruct_model(itr)
        model_params = model.get_all_params()

        header = self._build_grid_header("Model output (x mu nu)")
        header += (
            "# Model parameter values "
            + "({}): ".format(" ".join(model_params.keys()))
            + format(tuple(model_params.values()))
        )
        Y, Y_std = mesher.evaluate_func(model.predict, coord_shaped=True)
        self.dump_columns(file_path, (mesher.coords, Y, Y_std), header)

    def dump_acq_func(self, file_path: str | Path, itr: int=-1) -> None:
        """Write acquisition function values to a text file.

        Evaluates the acquisition function over the mesh grid and writes the results
        to a multi-column text file. For problems with more than 2 active dimensions,
        inactive dimensions are fixed to the next acquisition point location.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically in "data_acqfns/" subdirectory).
        itr : int, default=-1
            Iteration number for which to dump the acquisition function. Negative
            indices count from the end (e.g., -1 is the last iteration).

        Notes
        -----
        Output file columns: x1, [x2, ...], alpha
        where x1, x2, ... are coordinate values and alpha is the acquisition
        function value.
        """
        mesh = self.mesh
        mesh.fix_dim_preset(self.results, preset='next_acq', itr=itr)
        header = self._build_grid_header("# Acquisition function (x af)")
        acqfn = self.results.reconstruct_acq_func(itr=itr)
        acq_vals = mesh.evaluate_func(acqfn, coord_shaped=True)[0]
        self.dump_columns(file_path, (mesh.coords, acq_vals), header)

    def dump_true_func(self, file_path: str | Path) -> None:
        """Write true objective function values to a text file.

        Evaluates the true objective function over the mesh grid and writes the
        results to a multi-column text file. This is useful for comparing GP model
        predictions against ground truth. For problems with more than 2 active
        dimensions, inactive dimensions are fixed to the predicted minimum location
        from the final iteration.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically "true_func.dat").

        Notes
        -----
        Output file columns: x1, [x2, ...], f(x)
        where x1, x2, ... are coordinate values and f(x) is the true objective
        function value.
        """
        mesh = self.mesh
        mesh.fix_dim_preset(self.results, preset='min', itr=-1)
        header = self._build_grid_header("True function output (x f)")
        f_user = lambda X: self.settings.f.evaluate(X).Y
        output = self.mesh.evaluate_func(
            f_user, coord_shaped=True
        )
        self.dump_columns(file_path, (mesh.coords, output[0]), header)

    def dump_convergence(self, file_path: str | Path) -> None:
        """Write convergence diagnostics to a text file.

        Computes and exports convergence metrics that track changes in the predicted
        global minimum over iterations. These metrics help assess whether the
        optimization has converged.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically "convergence_measures.dat").

        Notes
        -----
        Output file columns:
        - iteration: Iteration number
        - npts: Ensemble size (number of training points)
        - dx_glmin: Euclidean norm of change in predicted minimum location
        - dmu_glmin_normalized: Absolute change in predicted minimum value,
          normalized by the estimated y-range

        The convergence metrics start from iteration 1 (not 0) since they measure
        changes between consecutive iterations.
        """
        header = (
            "# Convergence measures by iteration"
            " (iter npts dx_glmin abs(dmu_glmin)/yrange)\n"
        )
        results = self.results
        iters = list(results["x_glmin"].keys())[1:]
        npts = results.batch_tracker.ensemble_sizes[iters]
        x_glmin = results["x_glmin"].to_array()
        mu_glmin = results["mu_glmin"].to_array()
        dmu_glmin = np.abs(np.diff(mu_glmin))
        dy = np.zeros_like(dmu_glmin)
        for ind, itr in enumerate(iters):
            y_min, y_max = results.get_est_yrange(itr)
            dy[ind] = y_max - y_min
        dx_glmin = np.linalg.norm(np.diff(x_glmin, axis=0), axis=1)
        self.dump_columns(
            file_path,
            (iters, npts, dx_glmin, dmu_glmin / dy),
            header,
            formats={0: "%d", 1: "%5d"},
        )

    def dump_min_predictions(self, file_path: str | Path) -> None:
        """Write predicted global minimum locations and values to a text file.

        Exports the predicted global minimum (location, mean, and uncertainty) for
        each iteration. This provides a complete record of how the optimization's
        best estimate evolved over time.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically "minimum_predictions.dat").

        Notes
        -----
        Output file columns:
        - iteration: Iteration number
        - npts: Ensemble size (number of training points)
        - x_glmin: Predicted minimum location (d columns for d-dimensional problems)
        - mu_glmin: GP mean at the predicted minimum
        - nu_glmin: GP standard deviation at the predicted minimum
        """
        results = self.results
        header = (
            "# Global minimum predictions by iteration"
            " (iter npts x_glmin mu_glmin nu_glmin)\n"
        )
        iters = list(results["x_glmin"].keys())
        npts = results.batch_tracker.ensemble_sizes[iters]
        x_glmin = results["x_glmin"].to_array()
        mu_glmin = results["mu_glmin"].to_array()
        nu_glmin = results["nu_glmin"].to_array()
        self.dump_columns(
            file_path,
            (iters, npts, x_glmin, mu_glmin, nu_glmin),
            header,
            formats={0: "%d", 1: "%5d"},
        )

    def dump_true_func_at_minima(self, file_path: str | Path) -> None:
        """Write true function values at predicted minima to a text file.

        Evaluates the true objective function at each iteration's predicted minimum
        location and compares it to the GP model's prediction. This provides insight
        into the quality of the GP model and the accuracy of minimum predictions.

        Parameters
        ----------
        file_path : str or Path
            Path to the output file (typically "true_f_at_x_hat.dat").

        Notes
        -----
        Output file columns:
        - iteration: Iteration number
        - npts: Ensemble size (number of training points)
        - f(x_glmin): True function value at the predicted minimum
        - error: Prediction error (f(x_glmin) - mu_glmin)
        """
        results = self.results
        header = (
            "# True function value at x_glmin locations by iteration"
            " (iter npts f(x_glmin) f(x_glmin)-mu_glmin)\n"
        )
        iters = list(results["x_glmin"].keys())
        npts = results.batch_tracker.ensemble_sizes[iters]
        x_glmin = results["x_glmin"].to_array()
        mu_glmin = results["mu_glmin"].to_array()
        f_true_hats = self.settings.f.evaluate(x_glmin).Y
        diff = f_true_hats - mu_glmin
        self.dump_columns(
            file_path,
            (iters, npts, f_true_hats, diff),
            header,
            formats={0: "%d", 1: "%5d"},
        )
