import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import Matern, WhiteKernel, ConstantKernel
from scipy.stats import norm
import random
import warnings
from typing import List, Callable, Optional, Dict

def non_dominated_indices(acq_matrix: np.ndarray) -> List[int]:
    """
    Given an array of shape (N, M) with M acquisitions per candidate (for M objectives),
    return the list of row indices that are non-dominated (Pareto set).
    We assume we want to *maximize* each dimension of acq_matrix.

    A point j dominates point i if:
        acq_matrix[j, :] >= acq_matrix[i, :]
    in all dimensions, 
    and acq_matrix[j, :] > acq_matrix[i, :]
    in at least one dimension.
    """
    n_points = acq_matrix.shape[0]
    dominated = np.zeros(n_points, dtype=bool)

    for i in range(n_points):
        if dominated[i]:
            continue
        for j in range(n_points):
            if i == j or dominated[i]:
                continue
            # Check if candidate j dominates candidate i
            # For maximization: j >= i in all dims, and j > i in at least one
            better_equal = np.all(acq_matrix[j] >= acq_matrix[i])
            strictly_better = np.any(acq_matrix[j] > acq_matrix[i])
            if better_equal and strictly_better:
                dominated[i] = True
                break

    return [idx for idx in range(n_points) if not dominated[idx]]

# ------------------------------------------------------------------
# Modified BayesianOptimization Class with Constraints
# ------------------------------------------------------------------
class BayesianOptimization:
    def __init__(self,
                 surrogate_models=None,
                 active_model_key='gp',
                 gp_alpha=1e-1,
                 weights=None,
                 bounds=None,
                 n_candidates:int=1,
                 n_objectives: int=1,
                 constraints: Optional[List] = None,
                 discrete_levels: Optional[List[np.ndarray]] = None, 
                 logic: str = "all",
                 temperature:float = None,
                 random_state: int = 42):
        """
        A Bayesian Optimization framework that supports:
        
        - Multiple surrogate models (default is a Gaussian Process Regressor).
        - Various acquisition functions (EI, UCB, exploration, exploitation, and a new Boltzmann-based approach).
        - Constraints on the feature space.
        - Discrete or continuous sampling of candidate points.
        
        Parameters
        ----------
        surrogate_models : dict, optional
            Dictionary of surrogate models, keyed by string. 
            Defaults to a simple GP if None is provided.
        active_model_key : str, optional
            Key to indicate which surrogate model to use for predictions/acquisitions.
            Default is 'gp'.
        acquisition_func : callable, optional
            A custom acquisition function. If None, defaults to `self.expected_improvement`.
        gp_alpha : float, optional
            Regularization alpha for the default Gaussian Process. Default is 1e-6.
        weights : np.ndarray, optional
            If doing multi-objective optimization, specify the weights for each objective
            to reduce it to a single objective.
        bounds : np.ndarray, optional
            An array of shape (n_features, 2) specifying the search space bounds. 
        strategy : str, optional
            One of "ei", "ucb", "exploration", "exploitation", or the new "boltzmann". 
            Default is "ei".
        n_candidates : int, optional
            Number of candidates to recommend on each call. Default is 1.
        constraints : List[FeatureConstraint], optional
            List of constraint objects. Each must have `is_valid(x)` returning bool.
        discrete_levels : List[np.ndarray], optional
            If not None, each entry in the list is a 1D array with possible discrete 
            values for the corresponding dimension.
        logic : str, optional
            Either "all" or "any", controlling whether we require all constraints 
            to be satisfied or at least one. Default is "all".
        temperature : float optional
        """
        np.random.seed(random_state)
        random.seed(random_state)

        valid_modes = ["all", "any"]
        if logic not in valid_modes:
            raise ValueError(
                f"Invalid logic='{logic}'. Must be one of {valid_modes}."
            )
        self.logic = logic
        self.n_objectives = n_objectives
        self.bounds = bounds
        self.discrete_levels = discrete_levels
        self.constraints = constraints if constraints else []

        self.active_model_key = active_model_key
        self.weights = weights
        self._trained_objective_values = None
        self.n_candidates = n_candidates
        self.temperature = temperature

        # 1. If no surrogate models are provided, create a default Gaussian Process.
        # Build a dictionary of GPs, one per objective: "obj0", "obj1", ...
        self.surrogate_models: list[str, GaussianProcessRegressor] = []
        for i in range(n_objectives):
            kernel = (ConstantKernel(1.0, (1e-3, 1e3)) *
                      Matern(nu=2.5, length_scale=10.0, length_scale_bounds=(1e-2, 1e3)) +
                      WhiteKernel(noise_level=1e-6, noise_level_bounds=(1e-8, 1e-1)))
            gp = GaussianProcessRegressor(kernel=kernel,
                                          alpha=gp_alpha,
                                          normalize_y=True,
                                          n_restarts_optimizer=5,
                                          random_state=random_state)
            self.surrogate_models.append(gp)



    def _aggregate_objectives(self, objectives: np.ndarray) -> np.ndarray:
        # same as your original code
        if objectives.ndim == 1:
            return objectives
        n_obj = objectives.shape[1]
        if self.weights is None:
            return np.mean(objectives, axis=1)
        else:
            if len(self.weights) != n_obj:
                raise ValueError(f"Mismatch: 'weights' length {len(self.weights)} != # of objectives {n_obj}.")
            return objectives @ self.weights

    def fit(self, X: np.ndarray, Y: np.ndarray):
        """
        Fit each GP to its respective objective column.
        
        Parameters
        ----------
        X : np.ndarray, shape (n_samples, n_features)
        Y : np.ndarray, shape (n_samples, n_objectives)

        Raises
        ------
        ValueError
            If features or objectives have incompatible shapes or dimensions.
        """
        if X.ndim!=2:
            raise ValueError("X must be 2D: (n_samples, n_features).")
        if Y.ndim==1:
            # e.g. single objective as (n_samples,)
            Y = Y.reshape(-1,1)
        if Y.shape[1]!= self.n_objectives:
            raise ValueError(f"Y has {Y.shape[1]} objectives, but n_objectives={self.n_objectives}.")

        self.X_train = X
        self.Y_train = Y

        for i, model in enumerate(self.surrogate_models):
            model.fit(X, Y[:, i])

    def _predict(self, X: np.ndarray, objective_idx: int ):
        """
        Predict mean and standard deviation of the objective at the given features.
        
        Parameters
        ----------
        X : np.ndarray
            Shape (n_points, n_features).
        objective_idx

        Returns
        -------
        tuple of np.ndarray
            (mean, std) each of shape (n_points,).
        
        Raises
        ------
        NotImplementedError
            If the active surrogate model does not support uncertainty estimation.
        """
        model = self.surrogate_models[objective_idx]
        try:
            mean, std = model.predict(X, return_std=True)
        except TypeError:
            # e.g. for RandomForestRegressor
            if hasattr(model, "estimators_"):
                predictions = np.array([est.predict(X) for est in model.estimators_])
                mean = np.mean(predictions, axis=0)
                std = np.std(predictions, axis=0)
            else:
                raise NotImplementedError(
                    "The active surrogate model does not support uncertainty estimation."
                )
        return mean, std
    
    def validate(self, features: np.ndarray) -> bool:
        """
        Checks whether the provided feature vector satisfies
        the constraints according to the specified logic.
        
        Returns
        -------
        bool
            True if constraints pass, False otherwise.
        """
        if self.logic == "all":
            return all(constraint(features) for constraint in self.constraints)
        elif self.logic == "any":
            return any(constraint(features) for constraint in self.constraints)
        return False

    # ----------------------------------------------------------------
    # Modified multi_objective_acquisition_matrix
    # ----------------------------------------------------------------
    def multi_objective_acquisition_matrix(self,
                                           X: np.ndarray,
                                           T: float=1.0,
                                           T0: float=0.5,
                                           k: float=10.0) -> np.ndarray:
        """
        Compute an M-dimensional "Boltzmann-like" acquisition for each candidate in X,
        where M = number of objectives. We treat each objective separately,
        then store them in columns.

        For each objective:
          - Predict mean, std
          - min-max normalize them
          - combine (1-mean_norm) + (std_norm) with logistic weighting
          - fill into the corresponding column

        Returns
        -------
        acq_matrix : np.ndarray of shape (X.shape[0], self.n_objectives)
        """
        n_candidates = X.shape[0]
        M = self.n_objectives
        acq_matrix = np.zeros((n_candidates, M), dtype=float)

        for obj_idx in range(M):
            mean, std = self._predict(X, obj_idx)
            min_m, max_m = mean.min(), mean.max()
            min_s, max_s = std.min(), std.max()

            rng_m = max_m - min_m if max_m>min_m else 1e-9
            rng_s = max_s - min_s if max_s>min_s else 1e-9

            mean_norm = (mean - min_m)/rng_m
            std_norm  = (std  - min_s)/rng_s

            # logistic factor
            sig_T = 1.0/(1.0 + np.exp(-k*(T-T0)))
            C_T = 1.0 - sig_T
            B_T = sig_T

            # Minimizing mean => exploitation = (1-mean_norm)
            exploitation = 1.0 - mean_norm
            exploration  = std_norm

            acq = C_T*exploitation + B_T*exploration
            acq_matrix[:, obj_idx] = acq

        return acq_matrix

    # ----------------------------------------------------------------
    # Modified recommend_candidates to handle constraints & discrete
    # ----------------------------------------------------------------
    def recommend_candidates(self,
                             n_candidates: int=1,
                             candidate_multiplier: int=10,
                             T: float=1.0,
                             T0: float=0.5,
                             k: float=10.0,
                             selection_mode: str="pareto",
                             min_distance: float=0.0,
                             discrete_design: bool=True,
                             avoid_repetitions: bool=True) -> np.ndarray:
        """
        Generates candidate_multiplier*n_candidates points, 
        evaluates an M-dim acquisition for each (where M=n_objectives),
        then selects the final n_candidates according to `selection_mode`.

        selection_mode: 
          - "pareto": pick from the Pareto front in M-dim acquisition space,
                      if front is bigger than n_candidates => pick subset by distance;
                      if front is smaller => fill remainder with top dominated by aggregator sum
          - "all-merged": for each candidate, take the max across M acquisitions => single scalar,
                          rank & pick top n with distance-based filtering.

        Returns up to n_candidates points, or fewer if constraints & distance 
        limit the final selection.
        """

        """
        1) Generate candidate_multiplier*n_candidates points
        2) Evaluate multi_objective_acquisition_matrix => shape (N, M)
        3) Extract Pareto front in that M-dim acquisition space
        4) If the front is larger than n_candidates, pick a subset using distance
        5) Return up to n_candidates from that front
        """

        # 1) Candidate pool
        candidate_count = candidate_multiplier*n_candidates
        candidates = []
        used = set()
        max_attempts = candidate_count*100
        attempts=0
        n_features = self.bounds.shape[0]

        while len(candidates)<candidate_count and attempts<max_attempts:
            if discrete_design and self.discrete_levels is not None:
                cand = []
                for dim_idx, (low, high) in enumerate(self.bounds):
                    possible_vals = self.discrete_levels[dim_idx]
                    valid_vals = possible_vals[(possible_vals>=low)&(possible_vals<=high)]
                    if len(valid_vals)==0:
                        break
                    cand.append(np.random.choice(valid_vals))
                if len(cand)<n_features:
                    attempts+=1
                    continue
                cand = np.array(cand, dtype=float)
            elif discrete_design:
                # integer
                cand = np.array([
                    random.randint(int(low), int(high)) for (low, high) in self.bounds
                ], dtype=float)
            else:
                # continuous
                cand = np.array([
                    random.uniform(low, high) for (low, high) in self.bounds
                ], dtype=float)

            if self.validate(cand):
                c_tuple = tuple(cand)
                if avoid_repetitions:
                    if c_tuple not in used:
                        used.add(c_tuple)
                        candidates.append(cand)
                else:
                    candidates.append(cand)
            attempts+=1

        if len(candidates)<n_candidates:
            raise RuntimeError("Not enough valid candidates found.")

        X_cand = np.array(candidates)

        # 2) Evaluate multi-objective acquisitions => shape (N, M)
        acq_matrix = self.multi_objective_acquisition_matrix(X_cand, T, T0, k)

        # 3) Depending on selection_mode, pick final set
        if selection_mode=="pareto":
            # 3a) get Pareto front
            pareto_idx = non_dominated_indices(acq_matrix)
            pareto_points = X_cand[pareto_idx]
            # if the front is exactly or bigger than n_candidates => pick subset
            if len(pareto_points)==0:
                # corner case: no non-dominated => all are dominated
                # fallback => do aggregator
                selection_mode="all-merged"
            elif len(pareto_points)<n_candidates:
                # pick all from front, then fill remainder from dominated
                chosen = list(pareto_points)

                # 3b) aggregator: sum across M dims for each candidate
                sums = np.sum(acq_matrix, axis=1)
                # remove those in front
                front_set = set(pareto_idx)
                dominated_idx = [i for i in range(len(X_cand)) if i not in front_set]
                # sort dominated by sum descending
                dominated_sorted = sorted(dominated_idx, key=lambda i: sums[i], reverse=True)

                def distance(a,b):
                    return np.linalg.norm(a-b)

                # fill the remainder
                for idx in dominated_sorted:
                    if len(chosen)>=n_candidates:
                        break
                    pt = X_cand[idx]
                    # distance check
                    if min_distance>0:
                        too_close = any(distance(pt, c0)<min_distance for c0 in chosen)
                        if too_close:
                            continue
                    chosen.append(pt)

                return np.array(chosen[:n_candidates])
            else:
                # front >= n_candidates => pick subset via distance
                chosen = []
                def distance(a,b):
                    return np.linalg.norm(a-b)
                chosen.append(pareto_points[0])
                for i in range(1, len(pareto_points)):
                    pt = pareto_points[i]
                    if min_distance>0:
                        too_close = any(distance(pt, c0)<min_distance for c0 in chosen)
                        if too_close:
                            continue
                    chosen.append(pt)
                    if len(chosen)>=n_candidates:
                        break
                return np.array(chosen)
        
        # Fallback or alternative approach
        if selection_mode=="all-merged":
            # For each candidate i, define merged_acq_i = max across M dims
            merged_acq = np.max(acq_matrix, axis=1)
            # sort descending
            sorted_idx = np.argsort(-merged_acq)
            chosen = []
            def distance(a,b):
                return np.linalg.norm(a-b)

            for idx in sorted_idx:
                pt = X_cand[idx]
                if len(chosen)==0:
                    chosen.append(pt)
                else:
                    if min_distance>0:
                        too_close = any(distance(pt,c0)<min_distance for c0 in chosen)
                        if too_close:
                            continue
                    chosen.append(pt)
                if len(chosen)>=n_candidates:
                    break
            return np.array(chosen)

        # If we got here with an unexpected selection_mode, fallback
        raise ValueError(f"Unknown selection_mode='{selection_mode}'")

    def evaluate_BO_over_generations(self, temperature: list = None, max_generations: int = 1000) -> None:
        """
        This function simulates the evolution of candidate points over several generations
        using Bayesian Optimization. It plots several figures to help diagnose and understand:
          - The background objective landscape.
          - How candidate points evolve over generations under different temperature settings.
          - The evolution of the predicted mean responses and errors for each objective.
          - The predicted mean surfaces of the final surrogate models.
        
        Parameters
        ----------
        temperature : list
            A list or array of temperature values (one per generation). The temperature influences
            the exploration/exploitation balance in the acquisition function.
        max_generations : int, optional
            Number of generations to simulate. (This function uses the length of `temperature`.)
        
        Returns
        -------
        None
        """
        import matplotlib.pyplot as plt
        import numpy as np

        # Determine the number of generations from the temperature array length.
        max_generations = len(temperature)

        # Define two synthetic objective functions.
        # Y: a sinusoidal function with noise.
        def Y(X1, X2, noise=0.0):
            # Returns a sinusoidal pattern over the domain.
            return np.sin(X1 / 3) + np.sin(X2 / 3) + np.random.rand(X1.shape[0]) * noise

        # Y2: a quadratic function with noise.
        def Y2(X1, X2, noise=0.0):
            return (X1 - 50)**2 / 2500 + (X2 - 50)**2 / 2500 + np.random.rand(X1.shape[0]) * noise

        # Generate synthetic training data.
        X_train = np.random.rand(5, 2) * 100  # 5 samples in a 100x100 domain.
        # Generate training responses for two objectives (each column corresponds to one objective).
        Y_train = np.array([Y(X_train[:, 0], X_train[:, 1], noise=0.5),
                            Y2(X_train[:, 0], X_train[:, 1], noise=1.0)]).T

        # Create a grid over the domain for background visualization.
        x_grid = np.linspace(0, 100, 200)
        y_grid = np.linspace(0, 100, 200)
        XX, YY = np.meshgrid(x_grid, y_grid)
        Z = Y(XX, YY)    # Background objective function for objective 1.
        Z2 = Y2(XX, YY)  # Background objective function for objective 2.

        # Plot the background gradient (for objective 1) as a contour plot.
        plt.figure(figsize=(8, 6))
        bg = plt.contourf(XX, YY, Z, levels=100, cmap='viridis')
        plt.colorbar(bg, label='Objective Value (Background)')
        plt.xlabel('X1')
        plt.ylabel('X2')
        plt.title('Background Objective Landscape (Objective 1)')
        plt.tight_layout()
        plt.show()

        # Initialize arrays to store mean responses and errors over generations.
        n_objectives = 2
        mean_response_array = np.zeros((max_generations, n_objectives))
        error_array = np.zeros((max_generations, n_objectives))

        # Instantiate the Bayesian Optimization model (multi-objective version).
        bo = BayesianOptimization(bounds=np.array([[0, 100], [0, 100]]), n_objectives=n_objectives)
        bo.fit(X_train, Y_train)

        # Loop over generations.
        for i, (gen, T_val) in enumerate(zip(range(1, max_generations + 1), temperature)):
            # Obtain candidate points using a Boltzmann sampling strategy with current temperature T_val.
            new_candidates = bo.recommend_candidates(n_candidates=10, T=T_val)
            # Compute responses for each candidate for both objectives.
            responses = np.array([Y(new_candidates[:, 0], new_candidates[:, 1], noise=0.5),
                                  Y2(new_candidates[:, 0], new_candidates[:, 1], noise=1.0)]).T

            # Scale marker size relative to temperature (for visualization purposes).
            marker_size = 400 * (T_val / max(temperature))

            # Plot candidate points on the background plot.
            plt.figure(figsize=(8, 6))
            plt.contourf(XX, YY, Z, levels=100, cmap='viridis')
            # Note: No legend is added as requested.
            plt.scatter(new_candidates[:, 0],
                        new_candidates[:, 1],
                        c=responses[:, 0],
                        cmap='viridis',
                        s=marker_size,
                        edgecolor='k',
                        alpha=0.8)
            plt.xlabel('X1')
            plt.ylabel('X2')
            plt.title(f'Generation {gen} Candidate Points (T = {T_val:.2f})')
            plt.tight_layout()

            # Optionally annotate each candidate with its generation number.
            # (Comment out if not desired.)
            # for (x, y) in new_candidates:
            #     plt.text(x, y, f'{gen}', fontsize=8, color='white', ha='center', va='center')

            # Augment training data with new candidates and their responses.
            X_train = np.concatenate((X_train, new_candidates), axis=0)
            Y_train = np.concatenate((Y_train, responses), axis=0)
            mean_response_array[i, :] = np.mean(responses, axis=0)
            bo.fit(X_train, Y_train)

            # Compute prediction errors over the grid for both objectives.
            points = np.column_stack((XX.ravel(), YY.ravel()))
            mean_obj1, _ = bo._predict(points, 0)
            mean_grid_obj1 = mean_obj1.reshape(XX.shape)
            error_array[i, 0] = np.sum(np.abs(Z - mean_grid_obj1))

            mean_obj2, _ = bo._predict(points, 1)
            mean_grid_obj2 = mean_obj2.reshape(XX.shape)
            error_array[i, 1] = np.sum(np.abs(Z2 - mean_grid_obj2))

        # Plot evolution of mean response for Objective 1.
        plt.figure(figsize=(8, 6))
        plt.plot(temperature, mean_response_array[:, 0], 'o', markersize=6, color='blue')
        plt.xlabel('Temperature')
        plt.ylabel('Mean Response (Objective 1)')
        plt.title('Evolution of Mean Response for Objective 1')
        plt.tight_layout()
        plt.show()

        # Plot evolution of mean response for Objective 2.
        plt.figure(figsize=(8, 6))
        plt.plot(temperature, mean_response_array[:, 1], 'o', markersize=6, color='red')
        plt.xlabel('Temperature')
        plt.ylabel('Mean Response (Objective 2)')
        plt.title('Evolution of Mean Response for Objective 2')
        plt.tight_layout()
        plt.show()

        # Plot evolution of prediction error for Objective 1.
        plt.figure(figsize=(8, 6))
        plt.plot(error_array[:, 0], 'o', markersize=6, color='blue')
        plt.xlabel('Generation')
        plt.ylabel('Error (Objective 1)')
        plt.title('Prediction Error Evolution for Objective 1')
        plt.tight_layout()
        plt.show()

        # Plot evolution of prediction error for Objective 2.
        plt.figure(figsize=(8, 6))
        plt.plot(error_array[:, 1], 'o', markersize=6, color='red')
        plt.xlabel('Generation')
        plt.ylabel('Error (Objective 2)')
        plt.title('Prediction Error Evolution for Objective 2')
        plt.tight_layout()
        plt.show()

        # Plot final predicted mean surface for Objective 1.
        points = np.column_stack((XX.ravel(), YY.ravel()))
        mean_obj1, _ = bo._predict(points, 0)
        mean_grid_obj1 = mean_obj1.reshape(XX.shape)
        plt.figure(figsize=(8, 6))
        bg = plt.contourf(XX, YY, mean_grid_obj1, levels=100, cmap='viridis')
        plt.colorbar(bg, label='Predicted Mean (Objective 1)')
        plt.xlabel('X1')
        plt.ylabel('X2')
        plt.title('Predicted Mean Surface for Objective 1')
        plt.tight_layout()
        plt.show()

        # Plot final predicted mean surface for Objective 2.
        mean_obj2, _ = bo._predict(points, 1)
        mean_grid_obj2 = mean_obj2.reshape(XX.shape)
        plt.figure(figsize=(8, 6))
        bg = plt.contourf(XX, YY, mean_grid_obj2, levels=100, cmap='viridis')
        plt.colorbar(bg, label='Predicted Mean (Objective 2)')
        plt.xlabel('X1')
        plt.ylabel('X2')
        plt.title('Predicted Mean Surface for Objective 2')
        plt.tight_layout()
        plt.show()

        # This function currently does not return values (generations_array and mutation_rate_array
        # are undefined), so we return None.
        return None
