
import numpy as np
import matplotlib.pyplot as plt
import pickle
from pathlib import Path
from typing import Union, Optional, Tuple

def plot_bo_model(
    model_source: Union[str, Path, object],
    bounds: Optional[Union[Tuple[float, float, float, float], np.ndarray]] = None,
    grid_size: int = 100,
    output_path: Optional[Union[str, Path]] = None,
    show: bool = True,
    dims: Tuple[int, int] = (0, 1),
    fixed_values: Optional[np.ndarray] = None
):
    """
    Plots the Gaussian Process mean and standard deviation for a Bayesian Optimization model.
    For N > 2, plots a 2D slice defined by `dims`, fixing other dimensions.

    Args:
        model_source: Path to a pickled BO model or the BO object itself.
        bounds: Plotting bounds. 
               - If 4-tuple: (x_min, x_max, y_min, y_max) for the 2 plotted dims.
               - If None: inferred from model bounds or data.
        grid_size: Resolution of the grid.
        output_path: If provided, saves the plot to this path.
        show: Whether to display the plot.
        dims: Tuple of (x_dim_index, y_dim_index) to plot.
        fixed_values: Values for the FIXED dimensions. 
                      If None, defaults to the coordinates of the BEST observed point.
                      Must be of shape (n_features,). Values at `dims` indices are ignored.
    """
    
    # 1. Load Model
    if isinstance(model_source, (str, Path)):
        import pickle
        with open(model_source, "rb") as f:
            bo = pickle.load(f)
    else:
        bo = model_source
        
    if not hasattr(bo, 'surrogate_models') or not bo.surrogate_models:
        print(f"Error: Object {type(bo)} does not appear to have 'surrogate_models' or list is empty.")
        return

    # 2. Determine Dimensionality & Bounds
    # Try to get training data shape to determine N features
    if hasattr(bo, 'X_train') and bo.X_train is not None:
         n_features = bo.X_train.shape[1]
    elif hasattr(bo, 'bounds') and bo.bounds is not None:
         n_features = len(bo.bounds)
    else:
         print("Error: Could not determine n_features from X_train or bounds.")
         return

    # Validate dims
    d1, d2 = dims
    if d1 < 0 or d1 >= n_features or d2 < 0 or d2 >= n_features:
        print(f"Error: dims {dims} out of range for n_features={n_features}")
        return
        
    # Get Plotting Bounds for the Slice
    if bounds is None:
        # Infer from BO bounds
        if hasattr(bo, 'bounds') and bo.bounds is not None:
             b = bo.bounds
             # (x_min, x_max, y_min, y_max)
             bounds = (b[d1, 0], b[d1, 1], b[d2, 0], b[d2, 1])
        elif hasattr(bo, 'X_train') and bo.X_train is not None:
             # Infer from data min/max with some padding
             min_d1, max_d1 = bo.X_train[:, d1].min(), bo.X_train[:, d1].max()
             min_d2, max_d2 = bo.X_train[:, d2].min(), bo.X_train[:, d2].max()
             pad1 = (max_d1 - min_d1) * 0.1 if max_d1 != min_d1 else 1.0
             pad2 = (max_d2 - min_d2) * 0.1 if max_d2 != min_d2 else 1.0
             bounds = (min_d1 - pad1, max_d1 + pad1, min_d2 - pad2, max_d2 + pad2)
        else:
             bounds = (-5, 5, -5, 5) # Default fallback

    # 3. Handle Fixed Values for N > 2
    # If not provided, find the best observed point
    current_fixed = np.zeros(n_features)
    
    if fixed_values is not None:
        if len(fixed_values) != n_features:
             print(f"Warning: fixed_values shape {fixed_values.shape} != n_features {n_features}. Padding/truncating.")
             current_fixed[:min(len(fixed_values), n_features)] = fixed_values[:min(len(fixed_values), n_features)]
        else:
             current_fixed = np.array(fixed_values)
    else:
        # Default: Best Point
        if hasattr(bo, 'X_train') and hasattr(bo, 'Y_train') and bo.Y_train is not None:
             # Assuming single obj or using first obj for "best"
             # BO minimizes usually? Or maximizes acquisition. 
             # EZGA minimizes F. So lowest Y is best.
             y_vals = bo.Y_train
             if y_vals.ndim > 1:
                 # Scalarize if multi-obj
                 scores = bo._aggregate_objectives(y_vals)
                 best_idx = np.argmin(scores)
             else:
                 best_idx = np.argmin(y_vals)
             
             current_fixed = bo.X_train[best_idx].copy()
             print(f"Slice fixed at Best Observed Point (idx={best_idx}): F={y_vals[best_idx]}")
    
    # 4. Create Grid
    x_grid = np.linspace(bounds[0], bounds[1], grid_size)
    y_grid = np.linspace(bounds[2], bounds[3], grid_size)
    X_mesh, Y_mesh = np.meshgrid(x_grid, y_grid)
    
    # Flatten mesh
    mesh_flat = np.vstack([X_mesh.ravel(), Y_mesh.ravel()]).T # (N_grid, 2)
    
    # Construct full query points
    # Start with all fixed
    query_points = np.tile(current_fixed, (mesh_flat.shape[0], 1))
    
    # Overwrite the active slice dimensions
    query_points[:, d1] = mesh_flat[:, 0]
    query_points[:, d2] = mesh_flat[:, 1]
    
    # 5. Predict
    gp = bo.surrogate_models[0]
    means, stds = gp.predict(query_points, return_std=True)
    
    # Handle outputs
    if means.ndim > 1 and means.shape[1] > 1:
        mean = means[:, 0]
        std = stds if stds.ndim == 1 else stds[:, 0]
    else:
        mean = means
        std = stds

    # Reshape back to grid
    Z_mean = mean.reshape(grid_size, grid_size)
    Z_std = std.reshape(grid_size, grid_size)
    
    # 6. Plot
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # Extent for imshow
    extent = bounds
    
    # A. Mean
    im1 = axes[0].imshow(Z_mean, extent=extent, origin='lower', cmap='viridis', aspect='auto')
    axes[0].set_title(f"GP Mean (Slice dims {d1}, {d2})")
    axes[0].set_xlabel(f"x[{d1}]")
    axes[0].set_ylabel(f"x[{d2}]")
    plt.colorbar(im1, ax=axes[0])
    
    # B. Uncertainty
    im2 = axes[1].imshow(Z_std, extent=extent, origin='lower', cmap='plasma', aspect='auto')
    axes[1].set_title(f"GP Uncertainty (Slice dims {d1}, {d2})")
    axes[1].set_xlabel(f"x[{d1}]")
    axes[1].set_ylabel(f"x[{d2}]")
    plt.colorbar(im2, ax=axes[1])
    
    # C. Overlay Data Points (Projected)
    # Only plot points that are somewhat "close" in the other dimensions? 
    # Or just plot local optimization path?
    # For now, simply project ALL points.
    if hasattr(bo, 'X_train') and bo.X_train is not None:
         x_pts = bo.X_train[:, d1]
         y_pts = bo.X_train[:, d2]
         
         # Optional: Highlight the "Best" point used for fixing
         axes[0].scatter(x_pts, y_pts, c='k', marker='.', alpha=0.3, label='Observed (Projected)')
         
         # Highlight the "Slice Center" (Best Point)
         axes[0].scatter([current_fixed[d1]], [current_fixed[d2]], c='red', marker='*', s=100, label='Slice Center (Best)')
         axes[0].legend()
         
         axes[1].scatter(x_pts, y_pts, c='k', marker='.', alpha=0.3)
         axes[1].scatter([current_fixed[d1]], [current_fixed[d2]], c='red', marker='*', s=100)
    
    # Kernel Info
    try:
        kernel = gp.kernel_
        if hasattr(kernel, "k1") and hasattr(kernel.k1, "k2"):
             # Matern assumption
             ls = kernel.k1.k2.length_scale
        else:
             ls = "?"
        
        info_str = f"Kernel: {kernel}\nLength Scales: {ls}"
        fig.suptitle(info_str, fontsize=9)
    except:
        pass

    plt.tight_layout()
    
    if output_path:
        plt.savefig(output_path)
        print(f"Plot saved to {output_path}")
        
    if show:
        plt.show()
    else:
        plt.close()

