"""Main post-processing interface for Bayesian optimization results.

This module provides the `PPMain` class, which serves as the high-level orchestrator
for post-processing BOSS optimization results. It coordinates the generation of
diagnostic plots and data dumps by delegating to specialized components:

- `boss.pp.plotter.Plotter`: Creates visualization plots (convergence, models,
  acquisition functions, hyperparameters)
- `boss.io.dumper.Dumper`: Writes numerical data to text files for external analysis
- `boss.pp.mesh.Mesh`: Manages coordinate grids for function evaluation

The `PPMain` class automates the entire post-processing workflow, creating a structured
directory (`postprocessing/`) containing:
- Hyperparameter evolution plots and data
- Convergence diagnostics
- Acquisition history visualization
- Per-iteration model and acquisition function outputs
- True function visualizations (when available)

Typical Usage
-------------
Post-processing is typically run after a BO optimization completes:

>>> from boss.bo.bo_main import BOMain
>>> from boss.pp.pp_main import PPMain
>>> bo = BOMain(func, bounds, kernel="rbf", initpts=5, iterpts=10)
>>> res = bo.run()
>>> pp = PPMain(res, pp_models=True, pp_acq_funcs=True)
>>> pp.run()

This creates the `postprocessing/` directory with all diagnostic outputs.

Alternatively, post-processing can be run from saved restart files:

>>> pp = PPMain.from_file("boss.rst", "boss.out")
>>> pp.run()

See Also
--------
boss.pp.plotter.Plotter : Creates diagnostic plots
boss.io.dumper.Dumper : Writes data files
boss.bo.bo_main.BOMain : Main optimization interface
"""
from __future__ import annotations

import shutil
from pathlib import Path
from typing import Any, Literal

import numpy as np

from boss.bo.results import BOResults
from boss.io.dumper import Dumper
from boss.io.main_output import MainOutput
from boss.pp.mesh import Mesh
from boss.pp.plotter import Plotter
from boss.settings import Settings
from copy import deepcopy
import copy


class PPMain:
    """Main post-processing orchestrator for Bayesian optimization results.

    The `PPMain` class manages the complete post-processing workflow, coordinating
    the `Plotter` and `Dumper` objects to generate diagnostic visualizations and
    data files. It creates a structured output directory with subdirectories for
    different types of outputs.

    Parameters
    ----------
    results : BOResults
        Bayesian optimization results object containing models, data, and settings.
    **pp_settings : Any
        Optional keyword arguments to override post-processing settings. Common options include:
        - pp_models : bool
            Whether to output model plots and data files (default: False)
        - pp_acq_funcs : bool
            Whether to output acquisition function plots and data (default: False)
        - pp_iters : array_like
            Specific iterations to post-process (default: all iterations)
        - pp_model_slice : array_like
            Grid specification for model/acquisition plots as [dim1, dim2, npts]
        - pp_truef_npts : int
            Number of grid points for true function visualization

    Attributes
    ----------
    results : BOResults
        The BO results object.
    settings : Settings
        Combined settings from results and pp_settings overrides.
    rstfile : str
        Path to the restart file.
    outfile : str
        Path to the output file.
    main_output : MainOutput
        Output handler for logging.
    basepath : Path
        Base directory for post-processing outputs ("postprocessing/").
    mesh : Mesh
        Mesh object for coordinate grid generation.
    plotter : Plotter
        Plotter object for creating visualizations.
    dumper : Dumper
        Dumper object for writing data files.

    See Also
    --------
    boss.pp.plotter.Plotter : Creates diagnostic plots
    boss.io.dumper.Dumper : Writes data files
    boss.bo.results.BOResults : Container for optimization results
    """
    def __init__(
        self,
        results: BOResults,
        **pp_settings: Any,
    ) -> None:
        keywords = results.settings
        keywords.update({k: v for k, v in pp_settings.items() if v is not None})
        settings = Settings(keywords.data, f=results.settings.f)
        settings.correct()

        self.results = results
        self.settings = settings
        self.rstfile = self.settings["rstfile"]
        self.outfile = self.settings["outfile"]
        self.main_output = None
        self._setup()

    @classmethod
    def from_file(
        cls, rstfile: str, outfile: str, main_output: MainOutput | None = None
    ) -> "PPMain":
        """Create a PPMain instance from saved restart and output files.

        This alternative constructor loads BO results and settings from disk,
        allowing post-processing to be performed independently from the optimization run.

        Parameters
        ----------
        rstfile : str
            Path to the restart file (typically "boss.rst") containing settings
            and restart data.
        outfile : str
            Path to the output file (typically "boss.out") containing iteration
            history and results.
        main_output : MainOutput or None, optional
            Optional output handler for logging. If None, a new one is created.

        Returns
        -------
        PPMain
            Initialized PPMain instance ready for post-processing.
        """
        self = cls.__new__(cls)
        self.results = BOResults.from_file(rstfile, outfile)
        self.settings = Settings.from_file(rstfile)
        self.rstfile = rstfile
        self.outfile = outfile
        self.main_output = main_output
        cls._setup(self)
        return self

    def _setup(self) -> None:
        """Initialize post-processing components and default settings.

        This internal method:
        - Sets default iteration range if not specified
        - Creates the MainOutput handler if needed
        - Initializes the Mesh object with grid specification
        - Creates Plotter and Dumper instances for output generation

        The method converts 1-indexed pp_model_slice dimensions to 0-indexed
        for internal use by the Mesh.
        """
        if self.settings["pp_iters"] is None:
            self.settings["pp_iters"] = np.arange(0, self.settings["iterpts"] + 1)

        if not getattr(self, "main_output", None):
            self.main_output = MainOutput(self.settings)
        self.basepath = Path("postprocessing")

        # Create objects to handle mesh, plots and dumps.
        self.zero_indexed_slice = copy.copy(self.settings["pp_model_slice"])
        self.zero_indexed_slice[0] -= 1
        self.zero_indexed_slice[1] -= 1
        self.mesh = Mesh.from_grid_spec(
            self.settings["bounds"], self.zero_indexed_slice
        )
        self.plotter = Plotter(self.results, self.mesh)
        self.dumper = Dumper(self.results, self.mesh)

    def run(self) -> None:
        """Execute the complete post-processing workflow.

        This method orchestrates the entire post-processing pipeline by:

        1. Creating the output directory structure (`postprocessing/` and subdirectories)
        2. Generating global diagnostic outputs:
           - Hyperparameter evolution plots and data files
           - Minimum predictions over iterations
           - Convergence diagnostics
           - Acquisition history visualization
        3. Processing per-iteration outputs (if enabled):
           - Model predictions (mean and variance) as plots and data files
           - Acquisition function values as plots and data files
        4. Generating true function outputs (if available and enabled):
           - True function values at predicted minima
           - True function visualization over the input space
        """
        if self.results is None:
            self.results = BOResults.from_file(self.rstfile, self.outfile)

        # Create the main pp directory.
        if self.basepath.is_dir():
            print("warning: overwriting directory 'postprocessing'")
        shutil.rmtree("postprocessing", ignore_errors=True)
        self.basepath.mkdir(exist_ok=True)
        self.create_subdirs()

        # process hyperparameters
        self.plotter.plot_hyperparameters(
            self.basepath / "hyperparameters.png", show=False
        )
        self.dumper.dump_hyperparams(self.basepath / "hyperparameters.dat")

        # process min preds
        self.dumper.dump_min_predictions(self.basepath / "minimum_predictions.dat")

        # process conv meas
        self.dumper.dump_convergence(self.basepath / "convergence_measures.dat")
        self.plotter.plot_convergence(
            save_as=self.basepath / "convergence_measures.png", show=False
        )

        # process acquisition locations
        self.plotter.plot_acquisitions(
            save_as=self.basepath / "acquisition_locations.png", show=False
        )

        # ensure minimum data is available for all iterations to be processed
        self.results.calc_missing_minima(self.settings["pp_iters"])

        # process per-iteration info
        for itr in self.settings["pp_iters"]:
            # acq. funcs
            if self.settings["pp_acq_funcs"]:
                file_dest = self._get_dest("acq", "data", itr)
                self.dumper.dump_acq_func(file_dest, itr=itr)
                file_dest = self._get_dest("acq", "graph", itr)
                self.plotter.plot_acq_func(itr=itr, save_as=file_dest, show=False)

            # models
            if self.settings["pp_models"]:
                file_dest = self._get_dest("model", "data", itr)
                self.dumper.dump_model(file_dest, itr=itr)
                file_dest = self._get_dest("model", "graph", itr)
                self.plotter.plot_model(itr=itr, save_as=file_dest, show=False)

        if self.settings["pp_truef_at_glmins"]:
            self.dumper.dump_true_func_at_minima(self.basepath / "true_f_at_x_hat.dat")
            self.plotter.plot_true_func_at_minima(
                self.basepath / "true_function_at_xhats.png", show=False
            )

        # process true function
        pp_truef_npts = self.settings["pp_truef_npts"]
        if pp_truef_npts:
            # Hack: temporarily change self.mesh to a new mesh with pp_truef_npts grid points in each dim
            mesh_modified = pp_truef_npts != self.mesh.grid_pts[0]
            if mesh_modified:
                pp_truef_slice = deepcopy(self.settings['pp_model_slice'])
                pp_truef_slice[0] -= 1
                pp_truef_slice[1] -= 1
                pp_truef_slice[2] = pp_truef_npts
                self.mesh = Mesh.from_grid_spec(
                    self.settings["bounds"], pp_truef_slice
                )

            self.plotter.plot_true_func(self.basepath / "true_func.png", show=False)
            self.dumper.dump_true_func(self.basepath / "true_func.dat")
            # revert back to the regular grid
            if mesh_modified:
                self.mesh = Mesh.from_grid_spec(
                    self.settings["bounds"], self.zero_indexed_slice
                )

    def _get_dest(
        self,
        target: Literal["model", "acq"],
        output_type: Literal["data", "graph"],
        itr: int,
    ) -> Path:
        """Generate output file path for per-iteration data or plots.

        Constructs the appropriate file path based on the target type (model or
        acquisition function), output type (data file or plot), and iteration number.
        File names follow the convention `it####_npts####.ext` where the first number
        is the iteration and the second is the ensemble size at that iteration.

        Parameters
        ----------
        target : {"model", "acq"}
            Type of output - "model" for GP model predictions, "acq" for acquisition function.
        output_type : {"data", "graph"}
            Format of output - "data" for text files (.dat), "graph" for plots (.png).
        itr : int
            Iteration number for which to generate the file path.

        Returns
        -------
        Path
            Full path to the output file.
        """
        templ = "it{:04}_npts{:04}"
        dests = {
            ("model", "data"): str(self.basepath) + "/data_models/{}.dat",
            ("model", "graph"): str(self.basepath) + "/graphs_models/{}.png",
            ("acq", "data"): str(self.basepath) + "/data_acqfns/{}.dat",
            ("acq", "graph"): str(self.basepath) + "/graphs_acqfns/{}.png",
        }
        npts = self.results.batch_tracker.ensemble_sizes[itr]
        file_dest = Path(dests[target, output_type].format(templ.format(itr, npts)))
        return file_dest

    def create_subdirs(self) -> None:
        """Create subdirectory structure for post-processing outputs.
        """
        if self.settings["pp_models"]:
            (self.basepath / "data_models").mkdir(exist_ok=True)
            (self.basepath / "graphs_models").mkdir(exist_ok=True)
        if self.settings["pp_acq_funcs"]:
            (self.basepath / "data_acqfns").mkdir(exist_ok=True)
            (self.basepath / "graphs_acqfns").mkdir(exist_ok=True)
        if self.settings["pp_local_minima"] is not None:
            (self.basepath / "data_local_minima").mkdir(exist_ok=True)
