from scipy.stats import norm
import numpy as np

def expected_improvement(mu, sigma, best_f, xi=0.0):
    """
    xi 是一个探索系数，xi > 0 时更倾向于探索。
    论文里通常设 xi=0 或 xi=0.01。
    """
    Z = (mu - best_f - xi) / (sigma + 1e-9)
    ei = (mu - best_f - xi) * norm.cdf(Z) + sigma * norm.pdf(Z)
    ei[sigma < 1e-10] = 0.0   # 完全确定的点 EI 为 0
    return ei

import numpy as np
from scipy.linalg import cholesky, solve_triangular
from scipy.optimize import minimize
from scipy.stats import norm


# ---- 核函数 ----

def matern_52(X1, X2, length_scale=1.0, signal_var=1.0):
    """Matern ν=5/2 核"""
    dist = np.sqrt(((X1[:, None] - X2[None, :]) ** 2).sum(-1))
    sqrt5d = np.sqrt(5) * dist / length_scale
    return signal_var * (1 + sqrt5d + sqrt5d**2 / 3) * np.exp(-sqrt5d)


# ---- GP 核心：拟合 + 预测 ----

class GaussianProcess:
    def __init__(self, noise=1e-6):
        self.noise = noise
        self.length_scale = 1.0
        self.signal_var = 1.0
        self.X_train = None
        self.L = None      # Cholesky 分解
        self.alpha = None  # K^{-1} y

    def _build_K(self, X1, X2):
        return matern_52(X1, X2, self.length_scale, self.signal_var)

    def _log_marginal_likelihood(self, params):
        """
        最大化对数边际似然来学习核超参数。
        params = [log(length_scale), log(signal_var)]
        """
        l, s = np.exp(params[0]), np.exp(params[1])
        n = len(self.X_train)
        K = matern_52(self.X_train, self.X_train, l, s) + self.noise * np.eye(n)
        try:
            L = cholesky(K, lower=True)
        except np.linalg.LinAlgError:
            return 1e10
        alpha = solve_triangular(L.T, solve_triangular(L, self.y_train, lower=True))
        # log p(y|X) = -½ y^T K^{-1} y - ½ log|K| - n/2 log(2π)
        lml = (-0.5 * self.y_train @ alpha
               - np.sum(np.log(np.diag(L)))
               - 0.5 * n * np.log(2 * np.pi))
        return -lml   # minimize 负 LML

    def fit(self, X, y, n_restarts=5):
        self.X_train = X
        self.y_train = y
        best_lml, best_params = np.inf, [0.0, 0.0]
        for _ in range(n_restarts):
            p0 = np.random.uniform([-2, -2], [2, 2])
            res = minimize(self._log_marginal_likelihood, p0,
                           method="L-BFGS-B", bounds=[(-5,5),(-5,5)])
            if res.fun < best_lml:
                best_lml, best_params = res.fun, res.x
        self.length_scale, self.signal_var = np.exp(best_params)
        n = len(X)
        K = self._build_K(X, X) + self.noise * np.eye(n)
        self.L = cholesky(K, lower=True)
        self.alpha = solve_triangular(
            self.L.T, solve_triangular(self.L, y, lower=True)
        )

    def predict(self, X_test):
        """返回预测均值 μ 和标准差 σ"""
        K_star = self._build_K(X_test, self.X_train)   # (m, n)
        mu = K_star @ self.alpha                         # μ = K_* K^{-1} y

        # σ² = k(x*,x*) - K_* K^{-1} K_*^T
        v = solve_triangular(self.L, K_star.T, lower=True)  # L^{-1} K_*^T
        k_diag = matern_52(X_test, X_test,
                           self.length_scale, self.signal_var).diagonal()
        var = k_diag - np.sum(v**2, axis=0)
        return mu, np.sqrt(np.clip(var, 0, None))


# ---- Acquisition Function ----

def expected_improvement(mu, sigma, best_f, xi=0.01):
    Z = (mu - best_f - xi) / (sigma + 1e-9)
    ei = (mu - best_f - xi) * norm.cdf(Z) + sigma * norm.pdf(Z)
    ei[sigma < 1e-10] = 0.0
    return ei


# ---- 主循环 ----

def bayesian_optimization(f, bounds, n_init=5, n_iter=20):
    """
    f: 黑盒目标函数（越大越好）
    bounds: [(low1, high1), (low2, high2), ...]
    """
    dim = len(bounds)
    bounds_arr = np.array(bounds)

    # 初始随机探索
    X_obs = np.random.uniform(bounds_arr[:, 0], bounds_arr[:, 1], size=(n_init, dim))
    y_obs = np.array([f(x) for x in X_obs])

    gp = GaussianProcess()

    for i in range(n_iter):
        # 1. 拟合 GP
        gp.fit(X_obs, y_obs)

        # 2. 在候选点上计算 EI
        X_cand = np.random.uniform(
            bounds_arr[:, 0], bounds_arr[:, 1], size=(5000, dim)
        )
        mu, sigma = gp.predict(X_cand)
        ei = expected_improvement(mu, sigma, best_f=y_obs.max())

        # 3. 选 EI 最高的点
        x_next = X_cand[np.argmax(ei)]

        # 4. 评估
        y_next = f(x_next)
        X_obs = np.vstack([X_obs, x_next])
        y_obs = np.append(y_obs, y_next)

        print(f"Iter {i+1:2d}: x={x_next}, f={y_next:.4f}, best={y_obs.max():.4f}")

    best_idx = np.argmax(y_obs)
    return X_obs[best_idx], y_obs[best_idx]


# ---- 测试：用 BO 找 Branin 函数的最优点 ----
if __name__ == "__main__":
    def branin(x):
        # 标准 benchmark，有 3 个全局最优点，最优值约 0.397（取负号变最大化）
        a, b, c = 1, 5.1/(4*np.pi**2), 5/np.pi
        r, s, t = 6, 10, 1/(8*np.pi)
        val = a*(x[1] - b*x[0]**2 + c*x[0] - r)**2 + s*(1-t)*np.cos(x[0]) + s
        return -val   # 转成最大化

    bounds = [(-5, 10), (0, 15)]
    best_x, best_y = bayesian_optimization(branin, bounds, n_init=5, n_iter=30)
    print(f"\n最优点: x={best_x}, f(x)={-best_y:.4f}（Branin 最优约 0.397）")

import numpy as np
from scipy.linalg import cholesky, solve_triangular
from scipy.stats import norm

NUM_LAYERS = 40
MID_START, MID_END = 8, 32   # 先验：只在中间层搜
MAX_SKIP = 20                  # 最多跳 50%

def mock_eval(z, seed=None):
    """
    模拟评估函数，替代"真实跑 self-speculative decoding 测速"。
    真实最优：跳第 12~22 层的奇数层（中间层冗余度最高）。
    """
    rng = np.random.RandomState(seed)
    optimal = set(range(12, 23, 2))
    n_skip = z.sum()
    if n_skip == 0 or n_skip > MAX_SKIP:
        return -1.0
    overlap = len(set(np.where(z == 1)[0]) & optimal)
    return float(np.clip(overlap / len(optimal) + rng.normal(0, 0.05), 0, 1))

def matern_52_binary(Z1, Z2, length_scale=3.0):
    dist = np.sqrt(((Z1[:, None].astype(float) - Z2[None].astype(float))**2).sum(-1))
    r = np.sqrt(5) * dist / length_scale
    return (1 + r + r**2 / 3) * np.exp(-r)

def sample_configs(n, rng):
    configs = []
    for _ in range(n):
        n_skip = rng.randint(1, MAX_SKIP + 1)
        mid = list(range(MID_START, MID_END))
        idx = rng.choice(mid, size=min(n_skip, len(mid)), replace=False)
        z = np.zeros(NUM_LAYERS, dtype=int); z[idx] = 1
        configs.append(z)
    return np.array(configs)

def bo_skip_layers(n_init=8, n_iter=20):
    rng = np.random.RandomState(42)

    # ---- 第一步：随机初始化，收集最初的几个观测 ----
    Z_obs = sample_configs(n_init, rng)
    y_obs = np.array([mock_eval(z, seed=i) for i, z in enumerate(Z_obs)])
    print(f"初始 {n_init} 次随机探索，最优 score={y_obs.max():.3f}，"
          f"方案={list(np.where(Z_obs[y_obs.argmax()]==1)[0])}\n")

    noise = 0.01
    for it in range(n_iter):

        # ---- 第二步：用所有观测点拟合 GP ----
        # GP 记住了每一个 (z, score) 对，核矩阵 K 编码了方案间的"相似程度"
        n = len(y_obs)
        K = matern_52_binary(Z_obs, Z_obs) + noise * np.eye(n)
        L = cholesky(K, lower=True)
        alpha = solve_triangular(L.T, solve_triangular(L, y_obs, lower=True))

        # ---- 第三步：对大量候选配置，用 GP 预测 μ 和 σ ----
        Z_cand = sample_configs(3000, rng)
        Ks = matern_52_binary(Z_cand, Z_obs)
        mu = Ks @ alpha                                          # μ = K_* K^{-1} y
        v = solve_triangular(L, Ks.T, lower=True)
        sigma = np.sqrt(np.clip(1.0 - (v**2).sum(0), 0, None)) # σ² = 1 - K_* K^{-1} K_*^T

        # ---- 第四步：计算 EI，选下一个要评估的点 ----
        best_f = y_obs.max()
        Z_ei = (mu - best_f) / (sigma + 1e-9)
        ei = (mu - best_f) * norm.cdf(Z_ei) + sigma * norm.pdf(Z_ei)

        best_idx = np.argmax(ei)
        z_next = Z_cand[best_idx]

        # ---- 打印 GP 此刻在"想什么" ----
        if it < 5 or it % 5 == 4:
            top3 = np.argsort(ei)[-3:][::-1]
            print(f"Iter {it+1:2d} | 当前最优 score={best_f:.3f}")
            print(f"  GP 最看好的前3个候选：")
            for rank, i in enumerate(top3):
                layers = list(np.where(Z_cand[i]==1)[0])
                print(f"    #{rank+1} skip={layers}")
                print(f"       μ={mu[i]:.3f}  σ={sigma[i]:.3f}  EI={ei[i]:.4f}")

        # ---- 第五步：实际评估，把新观测加进去 ----
        y_next = mock_eval(z_next, seed=n_init + it)
        Z_obs = np.vstack([Z_obs, z_next])
        y_obs = np.append(y_obs, y_next)

        if it < 5 or it % 5 == 4:
            print(f"  → 评估结果: score={y_next:.3f}，"
                  f"新最优={y_obs.max():.3f}\n")

    best_z = Z_obs[y_obs.argmax()]
    return list(np.where(best_z==1)[0]), y_obs.max()

best_layers, best_score = bo_skip_layers()
print(f"\n最终最优跳层方案: {best_layers}")
print(f"最优 score: {best_score:.4f}")