[PostgreSQL]自适应实验设计：汤普森采样与多臂老虎机

# 创建Python 3.10环境
conda create -n thompson_sampling python=3.10 -y
conda activate thompson_sampling

# 安装核心库
pip install numpy scipy matplotlib pandas seaborn
pip install tqdm  # 用于进度条显示

# 可选：Jupyter支持
pip install jupyterlab

import numpy as np
from scipy.stats import beta
from typing import List, Tuple, Optional
import matplotlib.pyplot as plt
from tqdm import tqdm

class ThompsonSamplingBandit:
    """
    贝努利奖励的汤普森采样多臂老虎机实现
    
    特性：
    - 支持动态添加臂
    - 提供置信区间估计
    - 内置性能监控
    - 可配置的先验参数
    """
    
    def __init__(self, n_arms: int, prior_alpha: float = 1.0, prior_beta: float = 1.0):
        """
        初始化老虎机实例
        
        参数：
        -----------
        n_arms : int
            臂的数量
        prior_alpha : float, default=1.0
            Beta先验的α参数（伪成功次数）
        prior_beta : float, default=1.0
            Beta先验的β参数（伪失败次数）
        """
        self.n_arms = n_arms
        self.prior_alpha = prior_alpha
        self.prior_beta = prior_beta
        
        # 初始化后验参数：每个臂维护一个Beta分布
        # alphas[i] = 成功次数 + prior_alpha
        # betas[i] = 失败次数 + prior_beta
        self.alphas = np.full(n_arms, prior_alpha, dtype=np.float64)
        self.betas = np.full(n_arms, prior_beta, dtype=np.float64)
        
        # 记录历史数据用于分析
        self.total_pulls = 0
        self.arm_pulls = np.zeros(n_arms, dtype=np.int64)
        self.arm_rewards = np.zeros(n_arms, dtype=np.int64)
        self.history = []  # (time, arm, reward)三元组
        
        print(f"Initialized ThompsonSamplingBandit with {n_arms} arms")
        print(f"Prior: Beta({prior_alpha}, {prior_beta}) for all arms\n")
    
    def select_arm(self) -> int:
        """
        选择下一个要拉的臂
        
        实现逻辑：
        1. 从每个臂的当前后验分布中采样一个奖励估计值
        2. 选择采样值最大的臂
        
        返回：
        --------
        int : 选中的臂索引 (0到n_arms-1)
        """
        # 从Beta分布采样：为每个臂生成一个随机奖励估计
        # 使用.rvs()方法从每个臂的当前后验中抽取一个样本
        samples = np.random.beta(self.alphas, self.betas)
        
        # 选择采样值最大的臂
        # np.argmax返回第一个最大值的索引，符合概率匹配原则
        chosen_arm = np.argmax(samples)
        
        return int(chosen_arm)
    
    def reward(self, arm: int, reward: int) -> None:
        """
        更新选择臂的后验分布
        
        参数：
        -----------
        arm : int
            被拉动的臂索引
        reward : int {0, 1}
            观察到的奖励（0=失败，1=成功）
            
        更新规则：
        - 成功（reward=1）：alphas[arm] += 1
        - 失败（reward=0）：betas[arm] += 1
        """
        if arm not in range(self.n_arms):
            raise ValueError(f"Invalid arm index {arm}. Must be in [0, {self.n_arms-1}]")
        if reward not in [0, 1]:
            raise ValueError(f"Reward must be 0 or 1, got {reward}")
        
        # 更新后验参数：共轭先验使得更新极其简单
        if reward == 1:
            self.alphas[arm] += 1.0
            self.arm_rewards[arm] += 1
        else:
            self.betas[arm] += 1.0
        
        # 更新统计量
        self.arm_pulls[arm] += 1
        self.total_pulls += 1
        
        # 记录历史
        self.history.append((self.total_pulls, arm, reward))
    
    def get_arm_stats(self, arm: int) -> dict:
        """
        获取指定臂的详细统计信息
        
        返回：
        --------
        dict 包含：
        - 后验分布参数
        - 估计转化率
        - 置信区间
        - 拉动次数
        """
        if arm not in range(self.n_arms):
            raise ValueError(f"Invalid arm index {arm}")
        
        alpha = self.alphas[arm]
        beta_param = self.betas[arm]
        pulls = self.arm_pulls[arm]
        successes = self.arm_rewards[arm]
        
        # 后验均值作为点估计
        estimated_rate = alpha / (alpha + beta_param)
        
        # 95%可信区间（使用Beta分布的分位数）
        ci_lower = beta.ppf(0.025, alpha, beta_param)
        ci_upper = beta.ppf(0.975, alpha, beta_param)
        
        return {
            "arm": arm,
            "alpha": alpha,
            "beta": beta_param,
            "pulls": pulls,
            "successes": successes,
            "estimated_rate": estimated_rate,
            "ci_95": (ci_lower, ci_upper)
        }
    
    def get_overall_performance(self) -> dict:
        """
        计算整体性能指标
        
        返回：
        --------
        dict 包含：
        - 总拉动次数
        - 总奖励
        - 整体转化率
        - 各臂选择分布
        """
        total_rewards = np.sum(self.arm_rewards)
        overall_rate = total_rewards / self.total_pulls if self.total_pulls > 0 else 0
        
        return {
            "total_pulls": self.total_pulls,
            "total_rewards": total_rewards,
            "overall_conversion_rate": overall_rate,
            "arm_selection_distribution": self.arm_pulls / self.total_pulls if self.total_pulls > 0 else np.zeros(self.n_arms)
        }

class BanditVisualizer:
    """
    老虎机算法可视化工具
    
    提供三种视图：
    1. 后验分布演化：展示信念更新过程
    2. 累积遗憾曲线：评估算法性能
    3. 臂选择动态：观察探索-利用权衡
    """
    
    def __init__(self, bandit: ThompsonSamplingBandit, true_rates: List[float]):
        """
        参数：
        -----------
        bandit : ThompsonSamplingBandit
            已训练的老虎机实例
        true_rates : List[float]
            各臂的真实转化率（用于计算遗憾）
        """
        self.bandit = bandit
        self.true_rates = np.array(true_rates)
        self.optimal_rate = np.max(true_rates)
        self.optimal_arm = np.argmax(true_rates)
    
    def plot_posterior_distributions(self, figsize=(15, 5), arms=None):
        """
        绘制后验分布密度图
        
        技术细节：
        - 使用SciPy的beta.pdf计算密度
        - 填充95%高密度区域（HDI）
        - 标注真实转化率垂直线
        """
        if arms is None:
            arms = range(self.bandit.n_arms)
        
        n_plots = len(arms)
        fig, axes = plt.subplots(1, n_plots, figsize=figsize, squeeze=False)
        axes = axes.flatten()
        
        x = np.linspace(0, 1, 1000)
        
        for idx, arm in enumerate(arms):
            stats = self.bandit.get_arm_stats(arm)
            alpha, beta_param = stats["alpha"], stats["beta"]
            
            # 计算Beta分布密度
            y = beta.pdf(x, alpha, beta_param)
            
            # 绘制分布曲线
            axes[idx].plot(x, y, 'b-', linewidth=2, 
                          label=f'Posterior: Beta({alpha:.1f}, {beta_param:.1f})')
            
            # 填充95%置信区域
            ci_lower, ci_upper = stats["ci_95"]
            mask = (x >= ci_lower) & (x <= ci_upper)
            axes[idx].fill_between(x[mask], y[mask], alpha=0.3, color='blue')
            
            # 标注真实转化率
            axes[idx].axvline(self.true_rates[arm], color='red', linestyle='--', 
                             linewidth=2, label=f'True Rate: {self.true_rates[arm]:.3f}')
            
            # 装饰
            axes[idx].set_title(f'Arm {arm} (Pulls: {stats["pulls"]})', fontsize=12, fontweight='bold')
            axes[idx].set_xlabel('Conversion Rate', fontsize=10)
            axes[idx].set_ylabel('Density', fontsize=10)
            axes[idx].legend(fontsize=8)
            axes[idx].grid(True, alpha=0.3)
        
        plt.suptitle('Posterior Distributions of Arm Conversion Rates', 
                    fontsize=14, fontweight='bold')
        plt.tight_layout()
        plt.show()
    
    def plot_cumulative_regret(self):
        """
        绘制累积遗憾曲线
        
        计算逻辑：
        1. 计算每个时间步的瞬时遗憾：optimal_rate - chosen_arm_rate
        2. 累积求和得到总遗憾
        3. 与理论下界对比
        """
        history = np.array(self.bandit.history)
        if len(history) == 0:
            print("No history to plot")
            return
        
        times = history[:, 0].astype(int)
        chosen_arms = history[:, 1].astype(int)
        
        # 计算每个时间步的遗憾
        instant_regret = self.optimal_rate - self.true_rates[chosen_arms]
        cumulative_regret = np.cumsum(instant_regret)
        
        fig, ax = plt.subplots(figsize=(12, 6))
        
        # 绘制累积遗憾曲线
        ax.plot(times, cumulative_regret, 'b-', linewidth=2, 
                label=f'Observed Regret (Final: {cumulative_regret[-1]:.2f})')
        
        # 绘制理论最优臂的虚拟线（零遗憾）
        ax.axhline(y=0, color='green', linestyle='--', alpha=0.5, label='Optimal (Zero Regret)')
        
        # 添加对数参考线
        log_ref = 5 * np.log(times + 1)
        ax.plot(times, log_ref, 'r--', alpha=0.5, label='O(log T) Reference')
        
        ax.set_xlabel('Time Steps', fontsize=12)
        ax.set_ylabel('Cumulative Regret', fontsize=12)
        ax.set_title('Cumulative Regret Over Time', fontsize=14, fontweight='bold')
        ax.legend(fontsize=10)
        ax.grid(True, alpha=0.3)
        
        # 标注最终遗憾
        final_regret = cumulative_regret[-1]
        ax.annotate(f'Final Regret: {final_regret:.2f}',
                   xy=(times[-1], final_regret),
                   xytext=(times[-1]*0.7, final_regret*1.2),
                   arrowprops=dict(arrowstyle='->', color='red'),
                   fontsize=11, color='red')
        
        plt.tight_layout()
        plt.show()
    
    def plot_arm_selection_evolution(self, window_size=100):
        """
        绘制臂选择比例随时间演化
        
        技术实现：
        - 使用滑动窗口计算选择频率
        - 堆叠面积图展示探索→利用的过渡
        """
        history = np.array(self.bandit.history)
        if len(history) == 0:
            print("No history to plot")
            return
        
        times = history[:, 0].astype(int)
        chosen_arms = history[:, 1].astype(int)
        
        # 计算滑动窗口内的选择比例
        n_steps = len(times)
        selection_matrix = np.zeros((n_steps, self.bandit.n_arms))
        
        for i in range(n_steps):
            start = max(0, i - window_size)
            window_arms = chosen_arms[start:i+1]
            for arm in range(self.bandit.n_arms):
                selection_matrix[i, arm] = np.mean(window_arms == arm)
        
        fig, ax = plt.subplots(figsize=(12, 6))
        
        # 绘制堆叠面积图
        colors = plt.cm.Set3(np.linspace(0, 1, self.bandit.n_arms))
        ax.stackplot(times, selection_matrix.T, labels=[f'Arm {i}' for i in range(self.bandit.n_arms)],
                     colors=colors, alpha=0.7)
        
        # 标注最优臂
        ax.axhline(y=0.95, color='red', linestyle='--', alpha=0.5, 
                   label=f'Target: 95% on Arm {self.optimal_arm}')
        
        ax.set_xlabel(f'Time Steps (Window Size: {window_size})', fontsize=12)
        ax.set_ylabel('Selection Proportion', fontsize=12)
        ax.set_title('Evolution of Arm Selection Proportions', fontsize=14, fontweight='bold')
        ax.legend(loc='upper right', fontsize=9)
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()

def run_simulation(n_arms=5, n_steps=5000, true_rates=None, prior_alpha=1.0, prior_beta=1.0, seed=42):
    """
    执行完整的老虎机模拟实验
    
    参数：
    -----------
    n_arms : int
        臂的数量
    n_steps : int
        实验总步数
    true_rates : List[float], optional
        各臂的真实转化率。若为None，则随机生成
    prior_alpha, prior_beta : float
        Beta先验参数
    
    返回：
    --------
    tuple : (bandit, visualizer, metrics)
    """
    np.random.seed(seed)
    
    # 生成真实转化率（确保有明显优劣差异）
    if true_rates is None:
        true_rates = np.random.beta(2, 5, n_arms)
        true_rates[0] = 0.5  # 确保有一个较好的臂
    
    print("="*60)
    print("SIMULATION CONFIGURATION")
    print("="*60)
    print(f"Number of arms: {n_arms}")
    print(f"Time steps: {n_steps}")
    print(f"True conversion rates: {[f'{r:.4f}' for r in true_rates]}")
    print(f"Optimal arm: {np.argmax(true_rates)} (rate: {np.max(true_rates):.4f})")
    print(f"Prior: Beta({prior_alpha}, {prior_beta})")
    print("="*60 + "\n")
    
    # 初始化老虎机
    bandit = ThompsonSamplingBandit(n_arms, prior_alpha, prior_beta)
    
    # 模拟实验过程
    print("Running simulation...")
    for step in tqdm(range(n_steps), desc="Experiment Progress"):
        # 选择臂
        arm = bandit.select_arm()
        
        # 模拟真实环境反馈（按真实转化率生成奖励）
        reward = np.random.binomial(1, true_rates[arm])
        
        # 更新后验
        bandit.reward(arm, reward)
    
    print("\nSimulation completed!")
    
    # 创建可视化器
    visualizer = BanditVisualizer(bandit, true_rates)
    
    # 计算性能指标
    metrics = bandit.get_overall_performance()
    final_regret = calculate_regret(bandit, true_rates)
    
    print("\n" + "="*60)
    print("FINAL PERFORMANCE METRICS")
    print("="*60)
    print(f"Total pulls: {metrics['total_pulls']}")
    print(f"Total rewards: {metrics['total_rewards']}")
    print(f"Overall conversion rate: {metrics['overall_conversion_rate']:.4f}")
    print(f"Final cumulative regret: {final_regret:.2f}")
    print(f"Arm selection distribution: {[f'{p:.2%}' for p in metrics['arm_selection_distribution']]}")
    print("="*60)
    
    return bandit, visualizer, metrics

def calculate_regret(bandit, true_rates):
    """计算累积遗憾"""
    history = np.array(bandit.history)
    if len(history) == 0:
        return 0.0
    
    chosen_arms = history[:, 1].astype(int)
    optimal_rate = np.max(true_rates)
    instant_regret = optimal_rate - true_rates[chosen_arms]
    return np.sum(instant_regret)

# 执行实验
bandit, visualizer, metrics = run_simulation(
    n_arms=5,
    n_steps=5000,
    true_rates=[0.12, 0.08, 0.15, 0.10, 0.25],  # Arm 4最优
    prior_alpha=1.0,
    prior_beta=1.0,
    seed=42
)

# 可视化分析
visualizer.plot_posterior_distributions()
visualizer.plot_cumulative_regret()
visualizer.plot_arm_selection_evolution()

class ECommerceRecommendationBandit:
    """
    电商推荐专用老虎机
    
    扩展功能：
    1. 分层抽样（新用户 vs 老用户）
    2. 位置偏置矫正
    3. 非平稳性检测与响应
    4. 最小曝光量保证（每臂≥5万）
    """
    
    def __init__(self, n_arms, min_pulls_per_arm=50000, 
                 prior_alpha=1.0, prior_beta=19.0):  # 先验均值5%
        self.n_arms = n_arms
        self.min_pulls_per_arm = min_pulls_per_arm
        
        # 分层后验：新用户和老用户分开建模
        self.new_user_bandit = ThompsonSamplingBandit(n_arms, prior_alpha, prior_beta)
        self.old_user_bandit = ThompsonSamplingBandit(n_arms, prior_alpha, prior_beta)
        
        # 位置偏置参数：Slot 0-5的点击率倍数
        self.position_bias = np.array([1.4, 1.3, 1.2, 1.0, 0.9, 0.8])
        
        # 非平稳性检测：滑动窗口统计
        self.rolling_window = []
        self.window_size = 1000
        
        print(f"Initialized E-commerce Bandit with {n_arms} arms")
        print(f"Min pulls per arm: {min_pulls_per_arm}")
        print(f"Position bias model: {self.position_bias}\n")
    
    def select_arm(self, user_segment='old', position=0, current_time=None):
        """
        带业务约束的臂选择
        
        参数：
        -----------
        user_segment : str, 'new' or 'old'
            用户分群标签
        position : int, 0-5
            推荐位位置
        current_time : datetime, optional
            用于非平稳性检测
        """
        # 强制探索阶段：确保每个臂获得最小曝光
        for arm in range(self.n_arms):
            total_pulls = (self.new_user_bandit.arm_pulls[arm] + 
                          self.old_user_bandit.arm_pulls[arm])
            if total_pulls < self.min_pulls_per_arm:
                return arm
        
        # 选择对应用户分群的老虎机
        bandit = self.new_user_bandit if user_segment == 'new' else self.old_user_bandit
        
        # 基础选择逻辑
        arm = bandit.select_arm()
        
        # 位置偏置矫正：后验分布平移
        # 实际实现中，我们会将观察到的奖励除以位置偏置因子
        # 这里简化处理：仅在选择时考虑
        position_factor = self.position_bias[position]
        
        return arm
    
    def reward(self, arm, reward, user_segment='old', position=0):
        """
        带偏置矫正的奖励更新
        
        技术细节：
        - 将观察到的CTR除以位置偏置因子
        - 确保奖励仍在[0,1]区间内
        """
        # 位置偏置矫正
        position_factor = self.position_bias[position]
        adjusted_reward = reward / position_factor
        
        # 限制在合理范围（防止异常值）
        adjusted_reward = np.clip(adjusted_reward, 0, 1)
        
        # 更新对应分群的后验
        if user_segment == 'new':
            self.new_user_bandit.reward(arm, adjusted_reward)
        else:
            self.old_user_bandit.reward(arm, adjusted_reward)
        
        # 非平稳性检测：记录时间序列数据
        self.rolling_window.append({
            'time': len(self.rolling_window),
            'arm': arm,
            'reward': adjusted_reward,
            'segment': user_segment
        })
        
        # 保持窗口大小
        if len(self.rolling_window) > self.window_size:
            self.rolling_window.pop(0)
    
    def detect_non_stationarity(self, threshold=0.05):
        """
        检测环境是否发生非平稳变化
        
        实现：CUSUM（累积和）检测算法
        返回True时，应考虑重启后验或增加探索率
        """
        if len(self.rolling_window) < self.window_size:
            return False
        
        # 简化的均值偏移检测
        window_rewards = [r['reward'] for r in self.rolling_window]
        mid_point = len(window_rewards) // 2
        
        mean_first_half = np.mean(window_rewards[:mid_point])
        mean_second_half = np.mean(window_rewards[mid_point:])
        
        # 相对变化超过阈值
        relative_change = abs(mean_second_half - mean_first_half) / (mean_first_half + 1e-6)
        
        if relative_change > threshold:
            print(f"ALERT: Non-stationarity detected! Change: {relative_change:.2%}")
            return True
        
        return False

def simulate_ecommerce_scenario():
    """
    模拟电商推荐场景
    
    场景特征：
    - 时间跨度：72小时（4320个5分钟时段）
    - 流量分布：早高峰(9-11点)、午高峰(14-16点)、晚高峰(20-22点)
    - 用户分群：新用户占比15%，CTR普遍高30%
    - 周末效应：周六日CTR整体+15%
    """
    
    # 真实CTR配置（来自历史数据）
    true_rates = {
        'new': [0.098, 0.082, 0.127, 0.093, 0.115, 0.073],  # 新用户各算法CTR
        'old': [0.075, 0.062, 0.091, 0.067, 0.083, 0.053]   # 老用户各算法CTR
    }
    optimal_rates = {'new': 0.127, 'old': 0.091}  # 各分群最优CTR
    
    # 初始化系统
    bandit = ECommerceRecommendationBandit(n_arms=6, min_pulls_per_arm=50000)
    
    # 时间配置
    total_hours = 72
    time_slots_per_hour = 12  # 每5分钟一个时段
    total_slots = total_hours * time_slots_per_hour
    
    # 流量分布：模拟真实业务波动
    def get_hourly_traffic(hour_of_day, day_of_week):
        """计算每小时流量分布"""
        base_traffic = 1000  # 基础曝光量
        
        # 时段效应
        if 9 <= hour_of_day <= 11:
            multiplier = 2.5
        elif 14 <= hour_of_day <= 16:
            multiplier = 2.0
        elif 20 <= hour_of_day <= 22:
            multiplier = 2.8
        elif 0 <= hour_of_day <= 6:
            multiplier = 0.3
        else:
            multiplier = 1.0
        
        # 周末效应
        if day_of_week >= 5:  # 周六日
            multiplier *= 1.15
        
        return int(base_traffic * multiplier)
    
    # 记录性能指标
    metrics_log = {
        'hour': [],
        'total_rewards': [],
        'cumulative_regret': [],
        'exploration_ratio': [],
        'non_stationarity_alert': []
    }
    
    print("Starting E-commerce Simulation...")
    print(f"Total simulation time: {total_hours} hours ({total_slots} slots)")
    
    current_regret = 0
    exploration_count = 0
    
    for slot in tqdm(range(total_slots), desc="Simulating Hours"):
        hour = slot // time_slots_per_hour
        day = hour // 24
        hour_of_day = hour % 24
        
        # 获取当前时段流量
        traffic = get_hourly_traffic(hour_of_day, day)
        
        # 生成用户请求
        for i in range(traffic):
            # 用户分群
            is_new_user = np.random.random() < 0.15
            segment = 'new' if is_new_user else 'old'
            
            # 推荐位位置（有偏分布）
            position = np.random.choice(6, p=[0.25, 0.20, 0.15, 0.15, 0.15, 0.10])
            
            # 选择算法
            arm = bandit.select_arm(user_segment=segment, position=position)
            
            # 模拟真实奖励
            true_rate = true_rates[segment][arm]
            
            # 添加噪声模拟现实波动
            noise = np.random.normal(0, 0.01)
            effective_rate = np.clip(true_rate + noise, 0, 1)
            
            reward = np.random.binomial(1, effective_rate)
            
            # 更新后验
            bandit.reward(arm, reward, segment, position)
            
            # 计算瞬时遗憾
            optimal = optimal_rates[segment]
            current_regret += (optimal - effective_rate)
            
            # 统计探索行为（强制探索阶段）
            total_pulls_arm = (bandit.new_user_bandit.arm_pulls[arm] + 
                              bandit.old_user_bandit.arm_pulls[arm])
            if total_pulls_arm < 50000:
                exploration_count += 1
        
        # 每小时记录指标
        if slot % time_slots_per_hour == 0:
            # 非平稳性检测
            alert = bandit.detect_non_stationarity()
            
            metrics_log['hour'].append(hour)
            metrics_log['total_rewards'].append(np.sum(bandit.new_user_bandit.arm_rewards + 
                                                     bandit.old_user_bandit.arm_rewards))
            metrics_log['cumulative_regret'].append(current_regret)
            metrics_log['exploration_ratio'].append(exploration_count / ((slot+1) * 1000))
            metrics_log['non_stationarity_alert'].append(alert)
    
    print("\nSimulation completed!")
    return bandit, metrics_log, true_rates

# 运行完整仿真
bandit, metrics_log, true_rates = simulate_ecommerce_scenario()

def adaptive_forgetting_update(self, arm, reward, user_segment='old', decay_factor=0.999):
    """
    带指数衰减的后验更新
    
    对历史数据施加衰减因子，使算法适应非平稳环境
    θ_t ~ Beta(α_t * decay, β_t * decay)
    """
    bandit = self.new_user_bandit if user_segment == 'new' else self.old_user_bandit
    
    # 先衰减现有后验
    bandit.alphas[arm] *= decay_factor
    bandit.betas[arm] *= decay_factor
    
    # 再更新新观测
    bandit.reward(arm, reward)

class NonStationaryThompsonSampling(ThompsonSamplingBandit):
    def __init__(self, n_arms, gamma=0.99, *args, **kwargs):
        super().__init__(n_arms, *args, **kwargs)
        self.gamma = gamma  # 衰减因子
    
    def reward(self, arm, reward):
        # 先衰减历史后验
        self.alphas[arm] *= self.gamma
        self.betas[arm] *= self.gamma
        
        # 再添加新观测
        super().reward(arm, reward)

class LinearThompsonSampling:
    def __init__(self, n_arms, n_features, regularization=1.0):
        self.n_arms = n_arms
        self.n_features = n_features
        
        # 每臂维护一个贝叶斯线性回归模型
        self.posterior_means = [np.zeros(n_features) for _ in range(n_arms)]
        self.posterior_covs = [np.eye(n_features) * regularization for _ in range(n_arms)]
        self.arm_pulls = np.zeros(n_arms)
    
    def select_arm(self, context_vector):
        """
        context_vector: shape (n_features,)
        """
        max_sample = -float('inf')
        chosen_arm = 0
        
        for arm in range(self.n_arms):
            # 从后验采样权重
            mean = self.posterior_means[arm]
            cov = self.posterior_covs[arm]
            
            # 多元正态采样
            beta_sample = np.random.multivariate_normal(mean, cov)
            
            # 预测奖励
            predicted_reward = context_vector.dot(beta_sample)
            
            if predicted_reward > max_sample:
                max_sample = predicted_reward
                chosen_arm = arm
        
        return chosen_arm
    
    def reward(self, arm, context_vector, reward):
        # 更新后验：等价于贝叶斯线性回归
        # 使用Sherman-Morrison公式高效更新逆矩阵
        cov = self.posterior_covs[arm]
        mean = self.posterior_means[arm]
        
        # 更新协方差：(X^T X + λI)^{-1}
        x = context_vector.reshape(-1, 1)
        denom = 1 + x.T.dot(cov).dot(x)
        cov_update = cov - cov.dot(x).dot(x.T).dot(cov) / denom
        
        # 更新均值：(X^T X + λI)^{-1} X^T y
        mean_update = mean + cov_update.dot(x).dot(reward - x.T.dot(mean))
        
        self.posterior_covs[arm] = cov_update
        self.posterior_means[arm] = mean_update.flatten()
        self.arm_pulls[arm] += 1

class SafeThompsonSampling(ThompsonSamplingBandit):
    def __init__(self, *args, safety_threshold=0.03, **kwargs):
        super().__init__(*args, **kwargs)
        self.safety_threshold = safety_threshold
    
    def select_arm(self):
        # 标准TS采样
        samples = np.random.beta(self.alphas, self.betas)
        
        # 应用安全过滤
        posterior_means = self.alphas / (self.alphas + self.betas)
        safe_arms = np.where(posterior_means >= self.safety_threshold)[0]
        
        if len(safe_arms) == 0:
            # 全部不满足，选择后验均值最高的（保守策略）
            return np.argmax(posterior_means)
        else:
            # 在安全臂中选择采样值最高的
            safe_samples = samples[safe_arms]
            return safe_arms[np.argmax(safe_samples)]

# 分布式更新伪代码
def distributed_update(arm, reward_batch):
    """
    reward_batch: List[int]，批量奖励
    """
    # 从Redis读取当前后验
    alpha, beta = redis_client.hmget(f"arm:{arm}", "alpha", "beta")
    
    # 计算批量更新
    successes = sum(reward_batch)
    failures = len(reward_batch) - successes
    
    # 原子递增（避免竞态条件）
    redis_client.hincrbyfloat(f"arm:{arm}", "alpha", successes)
    redis_client.hincrbyfloat(f"arm:{arm}", "beta", failures)
    
    # 异步持久化到Cassandra（带时间戳）
    cassandra_client.execute(
        "INSERT INTO bandit_updates (arm, timestamp, alpha_delta, beta_delta) "
        "VALUES (%s, %s, %s, %s)",
        (arm, time.time(), successes, failures)
    )

package main

import (
    "net/http"
    "github.com/go-redis/redis/v8"
    "context"
    "encoding/json"
    "log"
    "math/rand"
)

// ArmPosterior holds Beta distribution parameters
type ArmPosterior struct {
    Alpha float64 `json:"alpha"`
    Beta  float64 `json:"beta"`
}

// ThompsonSamplingEngine handles arm selection
type ThompsonSamplingEngine struct {
    redisClient *redis.Client
    ctx         context.Context
}

// NewEngine creates a new TS engine
func NewEngine(redisAddr string) *ThompsonSamplingEngine {
    client := redis.NewClient(&redis.Options{
        Addr:     redisAddr,
        PoolSize: 100, // Connection pool for high concurrency
    })
    return &ThompsonSamplingEngine{
        redisClient: client,
        ctx:         context.Background(),
    }
}

// SelectArm returns the best arm for given context
func (e *ThompsonSamplingEngine) SelectArm(experimentID string, userSegment string) (int, error) {
    // Key pattern: experiment:{id}:segment:{segment}:arm:{arm}
    pattern := "experiment:" + experimentID + ":segment:" + userSegment + ":arm:*"
    
    // Fetch all arm posteriors for this segment
    iter := e.redisClient.Scan(e.ctx, 0, pattern, 0).Iterator()
    
    maxSample := -1.0
    var bestArm int
    var posterior ArmPosterior
    
    for iter.Next(e.ctx) {
        key := iter.Val()
        // Extract arm index from key
        var armIdx int
        _, err := fmt.Sscanf(key, "experiment:%s:segment:%s:arm:%d", &experimentID, &userSegment, &armIdx)
        if err != nil {
            continue
        }
        
        // Get posterior parameters from Redis
        val, err := e.redisClient.HGetAll(e.ctx, key).Result()
        if err != nil {
            continue
        }
        
        // Parse values
        posterior.Alpha, _ = strconv.ParseFloat(val["alpha"], 64)
        posterior.Beta, _ = strconv.ParseFloat(val["beta"], 64)
        
        // Sample from Beta distribution
        // Using BES algorithm (Best-of-class Efficient Sampling)
        sample := e.sampleBeta(posterior.Alpha, posterior.Beta)
        
        if sample > maxSample {
            maxSample = sample
            bestArm = armIdx
        }
    }
    
    if iter.Err() != nil {
        return -1, iter.Err()
    }
    
    return bestArm, nil
}

// sampleBeta implements efficient Beta sampling using the BES algorithm
func (e *ThompsonSamplingEngine) sampleBeta(a, b float64) float64 {
    // Handle edge cases
    if a <= 1.0 && b <= 1.0 {
        return e.sampleBetaSmallParams(a, b)
    }
    
    // Use ratio-of-uniforms method for general case
    return e.sampleBetaRatioOfUniforms(a, b)
}

// sampleBetaSmallParams handles α, β ≤ 1 (common in cold start)
func (e *ThompsonSamplingEngine) sampleBetaSmallParams(a, b float64) float64 {
    // Implementation based on Johnk's method
    u := math.Pow(rand.Float64(), 1.0/a)
    v := math.Pow(rand.Float64(), 1.0/b)
    return u / (u + v)
}

// Health check endpoint
func healthCheckHandler(w http.ResponseWriter, r *http.Request) {
    response := map[string]string{"status": "healthy"}
    json.NewEncoder(w).Encode(response)
}

func main() {
    engine := NewEngine("localhost:6379")
    
    http.HandleFunc("/select", func(w http.ResponseWriter, r *http.Request) {
        experimentID := r.URL.Query().Get("experiment_id")
        userSegment := r.URL.Query().Get("segment")
        
        arm, err := engine.SelectArm(experimentID, userSegment)
        if err != nil {
            http.Error(w, err.Error(), http.StatusInternalServerError)
            return
        }
        
        response := map[string]interface{}{
            "arm": arm,
            "experiment_id": experimentID,
        }
        json.NewEncoder(w).Encode(response)
    })
    
    http.HandleFunc("/health", healthCheckHandler)
    
    log.Fatal(http.ListenAndServe(":8080", nil))
}

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import *

# 定义输入Schema
reward_schema = StructType([
    StructField("experiment_id", StringType()),
    StructField("arm", IntegerType()),
    StructField("reward", DoubleType()),
    StructField("user_segment", StringType()),
    StructField("event_time", TimestampType())
])

def update_posterior(partition_data):
    """
    对每个分区的数据批量更新后验
    使用Redis流水线减少网络开销
    """
    import redis
    r = redis.Redis(host='redis-cluster', port=6379)
    pipe = r.pipeline()
    
    for row in partition_data:
        key = f"experiment:{row.experiment_id}:segment:{row.user_segment}:arm:{row.arm}"
        
        # 原子递增
        if row.reward > 0:
            pipe.hincrbyfloat(key, "alpha", row.reward)
        else:
            pipe.hincrbyfloat(key, "beta", 1.0 - row.reward)
    
    # 执行批量操作
    pipe.execute()

def main():
    spark = SparkSession.builder.appName("ThompsonSamplingUpdater").getOrCreate()
    
    # 从Kafka读取奖励数据
    df = spark.readStream \
        .format("kafka") \
        .option("kafka.bootstrap.servers", "kafka:9092") \
        .option("subscribe", "bandit-rewards") \
        .load() \
        .selectExpr("cast(value as string) as json") \
        .select(from_json(col("json"), reward_schema).alias("data")) \
        .select("data.*")
    
    # 处理延迟到达数据（Watermarking）
    df_with_watermark = df.withWatermark("event_time", "10 minutes")
    
    # 5分钟窗口聚合
    windowed_df = df_with_watermark.groupBy(
        window("event_time", "5 minutes"),
        "experiment_id", "arm", "user_segment"
    ).agg(
        sum("reward").alias("total_reward"),
        count("*").alias("total_events")
    )
    
    # 应用更新函数
    query = windowed_df.writeStream \
        .foreachBatch(lambda df, epoch_id: df.foreachPartition(update_posterior)) \
        .outputMode("update") \
        .start()
    
    query.awaitTermination()

if __name__ == "__main__":
    main()

# Prometheus配置 - 自定义指标导出
scrape_configs:
  - job_name: 'thompson-sampling-decision'
    static_configs:
      - targets: ['decision-service:8080']
    metrics_path: '/metrics'
    scrape_interval: 5s

# 关键PromQL查询
# 1. 各臂选择率
sum(rate(arm_selections_total[5m])) by (arm)

# 2. 累积遗憾增长速率
deriv(cumulative_regret[10m])

# 3. 后验分布熵（探索程度）
- sum(beta_distribution_alpha) by (arm) * log(beta_distribution_alpha) 
- sum(beta_distribution_beta) by (arm) * log(beta_distribution_beta)

def hybrid_allocation(thompson_arm, ab_test_groups, total_traffic, confidence_threshold=0.95):
    """
    混合分配策略
    
    thompson_arm: TS选择的臂
    ab_test_groups: A/B测试分组配置
    total_traffic: 总流量
    
    返回：
    -------
    dict: 各组分配比例
    """
    # 计算TS对最优臂的后验概率
    posterior = bandit.get_arm_stats(thompson_arm)
    alpha, beta = posterior['alpha'], posterior['beta']
    
    # 计算P(θ > θ_all_others)
    # 使用蒙特卡洛采样近似
    samples_thompson = np.random.beta(alpha, beta, 10000)
    samples_others = [np.random.beta(a, b, 10000) for a,b in zip(other_alphas, other_betas)]
    
    prob_best = np.mean(samples_thompson > np.max(samples_others, axis=0))
    
    if prob_best > confidence_threshold:
        # TS高置信，分配大部分流量
        return {
            'thompson_arm': 0.85,
            'ab_control': 0.10,
            'ab_other': 0.05
        }
    else:
        # TS低置信，依赖A/B测试
        return {
            'thompson_arm': 0.50,
            'ab_control': 0.30,
            'ab_other': 0.20
        }

import boto3
import json
from datetime import datetime

def snapshot_posteriors_to_s3(experiment_id):
    """
    将当前后验参数保存至S3
    """
    s3 = boto3.client('s3')
    bucket = 'thompson-sampling-snapshots'
    
    snapshot = {}
    for arm in range(n_arms):
        for segment in ['new', 'old']:
            key = f"experiment:{experiment_id}:segment:{segment}:arm:{arm}"
            posterior = redis_client.hgetall(key)
            snapshot[f"{segment}:{arm}"] = {
                "alpha": float(posterior[b'alpha']),
                "beta": float(posterior[b'beta']),
                "timestamp": datetime.utcnow().isoformat()
            }
    
    # 上传至S3
    s3_key = f"snapshots/{experiment_id}/{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.json"
    s3.put_object(
        Bucket=bucket,
        Key=s3_key,
        Body=json.dumps(snapshot),
        ServerSideEncryption='AES256'
    )
    
    print(f"Snapshot saved: s3://{bucket}/{s3_key}")

# 每小时定时任务
schedule.every().hour.do(snapshot_posteriors_to_s3, experiment_id="exp_123")

# Kubernetes回滚配置
apiVersion: batch/v1
kind: Job
metadata:
  name: bandit-rollback
spec:
  template:
    spec:
      containers:
      - name: rollback
        image: bandit-admin:latest
        command:
        - python
        - rollback.py
        env:
        - name: S3_SNAPSHOT_PATH
          valueFrom:
            configMapKeyRef:
              name: bandit-config
              key: snapshot_path
        - name: REDIS_CLUSTER
          valueFrom:
            secretKeyRef:
              name: redis-secret
              key: cluster_endpoint
      restartPolicy: OnFailure