# CMA-ES 和遺傳演算法（GA）、粒子群（PSO）的比較：10 維的三個測試函數，各演算法都只能評估 10,000 次，每個情況跑 5 個亂數種子
import math

import numpy as np

class CMAES:
    """(μ/μ_w, λ)-CMA-ES，參數用 Hansen（2016）教學文的預設值"""
    def __init__(self, x0, sigma0, rng):
        n = self.n = len(x0)
        self.lam = 4 + int(3 * math.log(n)); self.mu = self.lam // 2            # 每代樣本數 λ、選前 μ 個
        w = np.array([math.log(self.mu + 0.5) - math.log(i + 1) for i in range(self.mu)])
        self.w = w / w.sum(); me = self.mueff = 1 / np.sum(self.w ** 2)
        self.cs = (me + 2) / (n + me + 5)                                     # 步長的學習率與阻尼
        self.ds = 1 + 2 * max(0.0, math.sqrt((me - 1) / (n + 1)) - 1) + self.cs
        self.cc = (4 + me / n) / (n + 4 + 2 * me / n)                         # 共變異矩陣的學習率
        self.c1 = 2 / ((n + 1.3) ** 2 + me)
        self.cmu = min(1 - self.c1, 2 * (me - 2 + 1 / me) / ((n + 2) ** 2 + me))
        self.chin = math.sqrt(n) * (1 - 1 / (4 * n) + 1 / (21 * n * n))      # E‖N(0, I)‖
        self.m, self.sigma, self.rng = np.array(x0, float), sigma0, rng
        self.C, self.ps, self.pc, self.g = np.eye(n), np.zeros(n), np.zeros(n), 0
        self.D, self.B = np.ones(n), np.eye(n)

    def ask(self):
        Z = np.array([[float(self.rng.normal()) for _ in range(self.n)] for _ in range(self.lam)])
        self.Y = Z @ np.diag(self.D) @ self.B.T                               # y ~ N(0, C)
        return self.m + self.sigma * self.Y

    def tell(self, fits):
        n, ys = self.n, self.Y[np.argsort(fits, kind="stable")[:self.mu]]
        yw = self.w @ ys
        self.m = self.m + self.sigma * yw                                     # 1. 平均移到好樣本的加權平均
        self.ps = (1 - self.cs) * self.ps + math.sqrt(self.cs * (2 - self.cs) * self.mueff) * (self.B @ ((self.B.T @ yw) / self.D))
        self.g += 1
        hs = float(np.linalg.norm(self.ps) / math.sqrt(1 - (1 - self.cs) ** (2 * self.g)) < (1.4 + 2 / (n + 1)) * self.chin)
        self.pc = (1 - self.cc) * self.pc + hs * math.sqrt(self.cc * (2 - self.cc) * self.mueff) * yw
        self.C = ((1 - self.c1 - self.cmu) * self.C                          # 2. 共變異矩陣：rank-1（演化路徑）＋ rank-μ
                  + self.c1 * (np.outer(self.pc, self.pc) + (1 - hs) * self.cc * (2 - self.cc) * self.C)
                  + self.cmu * sum(wi * np.outer(y, y) for wi, y in zip(self.w, ys)))
        self.C = (self.C + self.C.T) / 2
        self.sigma *= math.exp((self.cs / self.ds) * (np.linalg.norm(self.ps) / self.chin - 1))  # 3. 步長：路徑太長就放大
        d2, self.B = np.linalg.eigh(self.C); self.D = np.sqrt(np.maximum(d2, 1e-300))

def cma(f, n, budget, seed):
    rng = np.random.default_rng(seed)
    es = CMAES(rng.uniform(-5, 5, n), 2.0, rng)
    best, ev = math.inf, 0
    while ev + es.lam <= budget:
        X = es.ask(); fits = np.array([f(x) for x in X]); ev += es.lam
        best = min(best, fits.min()); es.tell(fits)
    return best

def ga(f, n, budget, seed, pop=50):
    """實數編碼 GA：錦標賽選擇、BLX-0.5 交配（機率 0.9）、高斯突變（每個基因 1/n）、保留最好的 1 個"""
    rng = np.random.default_rng(100 + seed); lo, hi = -5.0, 5.0
    P = rng.uniform(lo, hi, (pop, n)); F = np.array([f(x) for x in P]); ev = pop
    while ev + pop - 1 <= budget:
        def pick():
            i, j = rng.integers(pop, size=2)
            return P[i] if F[i] < F[j] else P[j]
        kids = [P[int(np.argmin(F))].copy()]
        while len(kids) < pop:
            a, b = pick(), pick()
            if rng.random() < 0.9:
                l, h = np.minimum(a, b), np.maximum(a, b); d = h - l
                c = rng.uniform(l - 0.5 * d, h + 0.5 * d)
            else:
                c = a.copy()
            c = np.clip(c + (rng.random(n) < 1 / n) * rng.normal(0, 0.1 * (hi - lo), n), lo, hi)
            kids.append(c)
        P = np.array(kids); F = np.r_[F.min(), [f(x) for x in P[1:]]]; ev += pop - 1
    return float(F.min())

def pso(f, n, budget, seed, swarm=30, w=0.7298, c=1.49618):
    """標準 PSO：慣性 0.7298、c₁ = c₂ = 1.49618，速度限制在 0.2 倍範圍"""
    rng = np.random.default_rng(200 + seed); lo, hi, vm = -5.0, 5.0, 0.2 * 10
    X = rng.uniform(lo, hi, (swarm, n)); V = rng.uniform(-0.2, 0.2, (swarm, n)) * (hi - lo)
    F = np.array([f(x) for x in X]); ev = swarm; Pb, Fb = X.copy(), F.copy()
    while ev + swarm <= budget:
        g = Pb[int(np.argmin(Fb))]
        V = np.clip(w * V + c * rng.random((swarm, n)) * (Pb - X) + c * rng.random((swarm, n)) * (g - X), -vm, vm)
        X = np.clip(X + V, lo, hi); F = np.array([f(x) for x in X]); ev += swarm
        Pb[F < Fb], Fb[F < Fb] = X[F < Fb], F[F < Fb]
    return float(Fb.min())

n = 10
Q, _ = np.linalg.qr(np.random.default_rng(0).normal(size=(n, n)))      # 隨機旋轉
wts = 1e6 ** (np.arange(n) / (n - 1))
funcs = {
    "Rosenbrock（彎曲的狹長山谷）": lambda x: float(np.sum(100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2)),
    "旋轉橢球（條件數 10⁶）": lambda x: float(np.sum(wts * (Q @ x) ** 2)),
    "Rastrigin（很多區域最小值）": lambda x: float(10 * n + np.sum(x ** 2 - 10 * np.cos(2 * np.pi * x))),
}
print("10 維、每次最多評估 10,000 次、5 個種子；表中是找到的最小值（真正的最小值都是 0）")
for name, f in funcs.items():
    print(name)
    for alg, run in (("CMA-ES", cma), ("GA", ga), ("PSO", pso)):
        r = np.array([run(f, n, 10000, s) for s in range(5)])
        print(f"  {alg:6s} 中位數 {np.median(r):10.3g}　最好 {r.min():10.3g}　最差 {r.max():10.3g}")