from __future__ import annotations

import copy
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Literal

import numpy as np
from numpy.typing import ArrayLike, NDArray

from boss.bo.results import BOResults


class Mesh:
    def __init__(
        self,
        bounds: ArrayLike,
        active_dims: ArrayLike | None = None,
        grid_pts: int | tuple[int, ...] = 50,
    ) -> None:
        """
        Creates a mesh := coordinate grid over a hypercube given by bounds.

        The idea is to have 2 different representations of the grid:
            (1) As a coordinate array where each grid node is a single row,
               the shape of the array is thus (num_nodes x num_dim). This
               representation is referred to as `coords` in the code.
            (2) As a list of arrays identical to what would be returned by
               np.meshgrid. This representation is referred to as `grid` in the
               code.

        Representation (1) is used for evaluating vectorized functions on the grid
        while (2) is used for contour and surface plots.

        The class furthermore allows us to make grids for arbitrary slices along
        coordinate axes in our input space by specifying which dimensions
        (i.e. coordinates x, y, z...) are free to vary and which dimensions
        are fixed to create a slice in input space.

        Parameters
        ----------
        bounds : ArrayLike[float, 2d]
            Bounds for the entire input space, one row for each dimension
            in the form of [lower, upper].
        active_dims : ArrayLike[int, 1d] | None
            The free dimenensions of the input space. If none are given, all
            dimensions are assumed to be free.
        grid_pts : int | tuple
            The number of grid points per dimension. If only a single value is
            given it will be extended to all dimensions.
        """
        self.bounds = np.atleast_2d(bounds)
        if active_dims is None:
            self.active_dims = np.arange(len(self.bounds))
        else:
            self.active_dims = np.atleast_1d(active_dims)

        # The grid points are the number of points in each dimension.
        if isinstance(grid_pts, int) or len(grid_pts) == 1:
            self.grid_pts = tuple(np.repeat(grid_pts, len(self.active_dims)))
        else:
            self.grid_pts = grid_pts

        # The grid shape is similar to grid points, but
        # with the first and second elements swapped to be
        # consistent with the array shapes returned by np.meshgrid.
        if len(self.grid_pts) >= 2:
            self.grid_shape = (
                self.grid_pts[1],
                self.grid_pts[0],
            )
        else:
            self.grid_shape = self.grid_pts

        self._coords, self._grid = self.build()

    @classmethod
    def from_grid_spec(cls, bounds: ArrayLike, grid_spec: ArrayLike) -> Mesh:
        """
        Constructor that creates a Mesher instance from a grid specification.

        A grid specification is an array with 3 elements [x1, x2, npts] where x1, x2
        are the free dimensions and npts the number of grid points per dimension. This
        is the same format used by the BOSS keyword pp_models_slice, except that here the
        indexing of the dimensions start from 0.
        """
        slc = np.asarray(grid_spec, dtype=int)
        if slc[0] == slc[1]:
            active_dims = [slc[0]]
            shape = (slc[2],)
        else:
            active_dims = [slc[0], slc[1]]
            shape = (slc[2], slc[2])

        self = cls(bounds, active_dims, shape)
        return self

    @property
    def size(self) -> np.int64:
        """The total number of nodes in the grid."""
        return np.prod(self.grid_pts)

    @property
    def coords(self) -> NDArray:
        """
        The grid represented as a coordinate array.

        Each node in the grid occupies a single row in the array
        so the resulting shape is (num_nodes, num_dim).
        """
        return self._coords

    @property
    def grid(self) -> list[NDArray]:
        """
        The grid represented as meshgrid arrays.

        The values for each coordinate are given in separate arrays,
        see np.meshgrid for details.
        """
        return self._grid

    def build(self) -> tuple[NDArray, list[NDArray]]:
        """
        Builds two representations of the grid.

        Returns
        -------
        coords : NDArray
            Grid represented as a coordinate array with shape (num_nodes, num_dim).
        grid : list[NDArray]
            Grid represented as meshgrid arrays, one array per dimension.
        """
        basis = [
            np.linspace(self.bounds[d, 0], self.bounds[d, 1], self.grid_pts[i])
            for i, d in enumerate(self.active_dims)
        ]
        coords = np.nan * np.ones((self.size, len(self.bounds)), dtype=float)
        grid = np.meshgrid(*basis)
        coords[:, self.active_dims] = np.stack([X.flatten() for X in grid], axis=1)
        return coords, grid

    @property
    def fixed_dims(self) -> NDArray:
        """The dimensions that have been fixed to a constant value."""
        return np.setdiff1d(np.arange(len(self.bounds)), self.active_dims)

    def get_grid_spec(self) -> NDArray:
        """
        Returns a grid specification for the current mesh, if one exists.

        A grid specification is an array with 3 elements [x1, x2, npts] where x1, x2
        are the free dimensions and npts is the number of grid points per dimension.

        Returns
        -------
        grid_spec : NDArray
            Array containing [dim1, dim2, num_points] for the current mesh.

        Raises
        ------
        RuntimeError
            If the mesh has 2 active dimensions with different numbers of grid points.
        """
        if len(self.active_dims) == 2 and self.grid_pts[0] != self.grid_pts[1]:
            raise RuntimeError("Grid points in x != y.")
        return np.array([*self.active_dims, self.grid_pts[0]])

    def fix_dims(self, dim_vals: dict[int, float]) -> None:
        """
        Fixes the specified grid dimensions to the given values.

        Parameters
        ----------
        dim_vals : dict[int, float]
            Mapping from non-active dimensions to new fixed values.
        """
        dims = list(dim_vals.keys())
        intersct = np.intersect1d(self.active_dims, dims)
        if len(intersct) > 0:
            raise ValueError(f"Cannot fix free dimension: {intersct}")

        for dim, val in dim_vals.items():
            self._coords[:, dim] = val

    def evaluate_func(
        self, func: Callable, coord_shaped: bool = False
    ) -> tuple[NDArray, ...]:
        """
        Calculates a, possibly multi-valued, function over the grid.

        Parameters
        ----------
        func : Callable
        The function which to calcuate over the grid. The function is assumed
        to be vectorized and accept a 2d array as argument where each row defines
        a separate argument to be evaluated. The return value can be either
        a single array where each column represents

        coord_shaped : bool
        Determines the shape of the array of returned function values.
        If false, values are returned with shape compatible with meshgrid,
        otherwise they are in coordinate format: shape (n_pts, dim).

        Returns
        -------
        outputs : np.ndarray
        The calcuated function values.
        """
        outputs = func(self.coords)
        if coord_shaped:
            if isinstance(outputs, np.ndarray):
                if outputs.shape[1] > 1:
                    return tuple([*outputs.T])
                return (outputs,)
            elif isinstance(outputs, tuple):
                return outputs
        else:  # Reshape outputs to fit the mesh size
            if isinstance(outputs, tuple):
                return tuple([out.reshape(self.grid_shape) for out in outputs])
            elif isinstance(outputs, np.ndarray):
                return tuple([out.reshape(self.grid_shape) for out in outputs.T])

        raise RuntimeError(f"Cannot process function output of type {type(outputs)}")

    def fix_dim_preset(
        self,
        results: BOResults,
        preset: Literal["min", "mid", "next_acq"] = "min",
        itr: int = -1,
    ) -> None:
        """
        Fixes non-active dimensions using a preset strategy.

        Parameters
        ----------
        results : BOResults
            The Bayesian optimization results to extract fixed values from.
        preset : {"min", "mid", "next_acq"}
            The strategy for choosing fixed dimension values:
            - "min": Use the global minimum location
            - "mid": Use the midpoint of the bounds
            - "next_acq": Use the next acquisition point
        itr : int
            The iteration index to extract data from. Defaults to -1 (last iteration).

        Raises
        ------
        RuntimeError
            If preset is "min" but minimum data is not available for the given iteration.
        ValueError
            If preset is not one of the valid options.
        """
        if preset.lower() == "next_acq":
            x_fix = results.get_next_acq(itr)[0]
        elif preset.lower() == "min":
            x_fix = results.select("x_glmin", itr=itr)
            if x_fix is None:
                raise RuntimeError(f'Minumum data not available for iteration {itr}')
        elif preset.lower() == "mid":
            x_fix = np.mean(self.bounds, axis=0)
        else:
            raise ValueError(f"Unknown fix preset: {preset}")

        new_vals = dict(zip(self.fixed_dims, x_fix[self.fixed_dims]))
        self.fix_dims(new_vals)
