# Outcome Reward Model — 只评估最终结果
class OutcomeRewardModel(nn.Module):
    def __init__(self, base_model_name: str = "mistralai/Mistral-7B-v0.1"):
        super().__init__()
        self.base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
        self.reward_head = nn.Linear(self.base_model.config.hidden_size, 1)
        
    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
        # 取最后一个 token 的 hidden state 作为序列表示
        last_hidden = outputs.hidden_states[-1][:, -1, :]
        reward = self.reward_head(last_hidden)  # shape: (batch, 1)
        return reward.squeeze(-1)

# Process Reward Model — 评估每一步推理过程
class ProcessRewardModel(nn.Module):
    def __init__(
        self, 
        base_model_name: str = "mistralai/Mistral-7B-v0.1",
        step_separator: str = "\n\nStep:"
    ):
        super().__init__()
        self.base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            output_hidden_states=True
        )
        self.step_reward_head = nn.Linear(self.base_model.config.hidden_size, 1)
        self.step_separator = step_separator
        
    def forward(
        self, 
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        step_positions: Optional[List[int]] = None  # 每个步骤结束的位置
    ) -> Dict[str, torch.Tensor]:
        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = outputs.hidden_states[-1]
        
        if step_positions is not None:
            # 提取每个步骤结束位置的 hidden state
            step_rewards = []
            for batch_idx, positions in enumerate(step_positions):
                step_hiddens = hidden_states[batch_idx, positions, :]
                rewards = self.step_reward_head(step_hiddens)
                step_rewards.append(rewards)
            return {"step_rewards": step_rewards}
        else:
            # 回退到 outcome reward
            last_hidden = hidden_states[:, -1, :]
            reward = self.step_reward_head(last_hidden)
            return {"outcome_reward": reward.squeeze(-1)}

@dataclass
class StepLabel:
    """单步标注数据结构"""
    step_index: int
    step_text: str
    is_correct: Optional[bool]  # None = 不确定
    confidence: float  # 0.0 ~ 1.0
    reasoning: str  # 为什么正确/错误

class AutoStepLabeler:
    """
    自动化过程监督标签生成器
    使用多种策略自动标注每条推理的每一步
    """
    
    def __init__(
        self,
        judge_model: Any,  # LLM judge
        consistency_samples: int = 5,  # 一致性采样数量
        mc_dropout_samples: int = 10,  # MC Dropout 不确定性估计
    ):
        self.judge = judge_model
        self.consistency_samples = consistency_samples
        self.mc_dropout_samples = mc_dropout_samples
    
    def label_by_rollout_consistency(self, steps: List[str], ground_truth: str) -> List[StepLabel]:
        """
        策略1: Rollout 一致性法
        
        对于推理链中的第 k 步，从此步骤开始生成多个可能的后续推理，
        统计哪些后续路径能到达正确答案。如果大多数路径都正确，
        则第 k 步是正确的；否则是错误的。
        """
        labels = []
        for i, step in enumerate(steps):
            # 从第 i 步之后开始 rollout
            rollout_results = []
            for _ in range(self.consistency_samples):
                continuation = self._generate_continuation(steps[:i+1])
                is_correct = self._check_answer(continuation, ground_truth)
                rollout_results.append(is_correct)
            
            # 统计一致性
            correct_ratio = sum(rollout_results) / len(rollout_results)
            confidence = abs(correct_ratio - 0.5) * 2  # 0.5 → 0, 1.0 → 1
            
            is_correct = correct_ratio > 0.5
            labels.append(StepLabel(
                step_index=i,
                step_text=step,
                is_correct=is_correct if confidence > 0.2 else None,
                confidence=confidence,
                reasoning=f"Rollout consistency: {correct_ratio:.0%} paths reach correct answer"
            ))
        
        return labels
    
    def label_by_mc_dropout_uncertainty(self, steps: List[str], ground_truth: str) -> List[StepLabel]:
        """
        策略2: MC Dropout 不确定性法
        
        在推理的每一步启用 dropout，多次前向传播观察隐藏状态的变化。
        高不确定性 = 推理可能存在错误。
        """
        self.judge.train()  # 启用 dropout
        labels = []
        
        for i, step in enumerate(steps):
            step_representations = []
            
            for _ in range(self.mc_dropout_samples):
                # 获取第 i 步的 hidden state
                rep = self._get_step_representation(steps[:i+1], i)
                step_representations.append(rep)
            
            # 计算 representation 的方差
            reps_tensor = torch.stack(step_representations)
            variance = reps_tensor.var(dim=0).mean().item()
            
            # 高方差 = 模型对此步不确定 = 可能错误
            confidence = min(variance * 10, 1.0)
            is_correct = variance < 0.05  # 阈值可调
            
            labels.append(StepLabel(
                step_index=i,
                step_text=step,
                is_correct=is_correct if confidence > 0.3 else None,
                confidence=confidence,
                reasoning=f"MC Dropout variance: {variance:.4f}"
            ))
        
        self.judge.eval()  # 恢复评估模式
        return labels
    
    def label_by_mcts(self, steps: List[str], ground_truth: str, num_simulations: int = 50) -> List[StepLabel]:
        """
        策略3: MCTS 蒙特卡洛树搜索
        
        使用 MCTS 探索推理树，每个节点（步骤）的访问频率和胜率
        提供了高质量的过程监督信号。
        """
        class MCTSNode:
            def __init__(self, step_sequence: List[str], parent=None):
                self.step_sequence = step_sequence
                self.parent = parent
                self.visits = 0
                self.wins = 0
                self.children = []
            
            @property
            def win_rate(self) -> float:
                return self.wins / max(self.visits, 1)
            
            @property
            def ucb_score(self, exploration: float = 1.4) -> float:
                if self.visits == 0:
                    return float("inf")
                exploitation = self.win_rate
                exploration_term = exploration * math.sqrt(
                    math.log(self.parent.visits) / self.visits
                )
                return exploitation + exploration_term
        
        # 构建推理树
        root = MCTSNode(steps)
        
        for _ in range(num_simulations):
            # Selection
            node = root
            while node.children:
                node = max(node.children, key=lambda n: n.ucb_score)
            
            # Expansion
            next_steps = self._generate_possible_next_steps(node.step_sequence)
            for next_step in next_steps[:3]:  # 限制分支
                child = MCTSNode(node.step_sequence + [next_step], parent=node)
                node.children.append(child)
            
            # Simulation
            if node.children:
                node = random.choice(node.children)
            
            # Rollout
            result = self._rollout(node.step_sequence, ground_truth)
            
            # Backpropagation
            while node:
                node.visits += 1
                if result:
                    node.wins += 1
                node = node.parent
        
        # 根据节点访问频率给每一步打分
        labels = []
        for i, step in enumerate(steps):
            node = self._find_node(root, steps[:i+1])
            if node and node.visits > 0:
                labels.append(StepLabel(
                    step_index=i,
                    step_text=step,
                    is_correct=node.win_rate > 0.5,
                    confidence=abs(node.win_rate - 0.5) * 2,
                    reasoning=f"MCTS win rate: {node.win_rate:.0%} ({node.visits} visits)"
                ))
        
        return labels

class StepLabelQualityEnsurer:
    """
    过程监督标签质量保证
    
    多策略交叉验证 + 置信度校准 + 人工抽检接口
    """
    
    def __init__(self):
        self.strategies = {
            "rollout": AutoStepLabeler().label_by_rollout_consistency,
            "mc_dropout": AutoStepLabeler().label_by_mc_dropout_uncertainty,
            "mcts": AutoStepLabeler().label_by_mcts,
        }
    
    def cross_validate_labels(self, steps: List[str], ground_truth: str) -> List[StepLabel]:
        """
        多策略交叉验证：只有多个策略一致同意的标签才被采用
        """
        all_labels = {}
        
        for strategy_name, label_fn in self.strategies.items():
            try:
                labels = label_fn(steps, ground_truth)
                for label in labels:
                    if label.step_index not in all_labels:
                        all_labels[label.step_index] = []
                    all_labels[label.step_index].append(label)
            except Exception as e:
                logger.warning(f"Strategy {strategy_name} failed: {e}")
        
        # 融合多策略结果
        merged_labels = []
        for step_idx, label_list in all_labels.items():
            valid_labels = [l for l in label_list if l.is_correct is not None]
            
            if len(valid_labels) >= 2:
                # 多数投票
                correct_votes = sum(1 for l in valid_labels if l.is_correct)
                majority = correct_votes > len(valid_labels) / 2
                avg_confidence = sum(l.confidence for l in valid_labels) / len(valid_labels)
                
                merged_labels.append(StepLabel(
                    step_index=step_idx,
                    step_text=label_list[0].step_text,
                    is_correct=majority,
                    confidence=avg_confidence,
                    reasoning=f"Cross-validated by {len(valid_labels)} strategies"
                ))
            elif len(valid_labels) == 1:
                # 单一策略，但置信度降低
                label = valid_labels[0]
                label.confidence *= 0.6
                merged_labels.append(label)
        
        return merged_labels
    
    def calibrate_with_human_feedback(
        self, 
        auto_labels: List[StepLabel], 
        human_labels: List[StepLabel]
    ) -> CalibrationResult:
        """
        使用人工标注校准自动标签的置信度
        """
        from sklearn.calibration import calibration_curve
        
        auto_confidences = [l.confidence for l in auto_labels if l.is_correct is not None]
        human_agreements = []
        
        for auto_label in auto_labels:
            matching_human = next(
                (hl for hl in human_labels if hl.step_index == auto_label.step_index),
                None
            )
            if matching_human:
                agreement = auto_label.is_correct == matching_human.is_correct
                human_agreements.append(1.0 if agreement else 0.0)
        
        if len(human_agreements) < 10:
            return CalibrationResult(
                calibration_score=0.0,
                samples_count=len(human_agreements),
                note="Insufficient human labels for calibration"
            )
        
        # 计算 ECE (Expected Calibration Error)
        prob_true, prob_pred = calibration_curve(
            human_agreements, auto_confidences[:len(human_agreements)], n_bins=5
        )
        ece = np.mean(np.abs(prob_true - prob_pred))
        
        return CalibrationResult(
            calibration_score=1.0 - ece,
            samples_count=len(human_agreements),
            calibration_factor=1.0 / (1.0 - ece + 1e-6) if ece > 0.1 else 1.0
        )

class PRMEnhancedPPOTrainer:
    """
    集成 Process Reward Model 的 PPO 训练器
    
    将过程奖励（每一步）与最终结果奖励（outcome）加权融合，
    引导策略更关注正确推理过程而非仅仅最终答案。
    """
    
    def __init__(
        self,
        policy_model: nn.Module,
        process_reward_model: ProcessRewardModel,
        value_model: nn.Module,
        outcome_weight: float = 0.3,     # 最终结果奖励权重
        process_weight: float = 0.7,     # 过程奖励权重
        kl_coef: float = 0.1,
        ppo_epochs: int = 4,
        batch_size: int = 64,
        learning_rate: float = 5e-6,
    ):
        self.policy = policy_model
        self.prm = process_reward_model
        self.value = value_model
        self.outcome_weight = outcome_weight
        self.process_weight = process_weight
        self.kl_coef = kl_coef
        self.ppo_epochs = ppo_epochs
        self.batch_size = batch_size
        self.optimizer = torch.optim.AdamW(
            list(self.policy.parameters()),
            lr=learning_rate
        )
    
    def compute_hybrid_reward(
        self,
        sequences: torch.Tensor,         # (batch, seq_len)
        attention_mask: torch.Tensor,    # (batch, seq_len)
        step_positions: List[List[int]], # 每一步在序列中的位置
        ground_truth_correct: torch.Tensor  # (batch,) 最终答案是否正确
    ) -> Tuple[torch.Tensor, Dict]:
        """
        计算混合奖励 = outcome_weight × outcome_reward + process_weight × process_reward
        
        同时返回各分量用于分析和日志
        """
        # 1. 计算 outcome reward（最终结果奖励）
        outcome_rewards = torch.where(
            ground_truth_correct > 0,
            torch.tensor(1.0, device=sequences.device),
            torch.tensor(-1.0, device=sequences.device)
        )
        
        # 2. 计算 process reward（过程奖励）
        with torch.no_grad():
            prm_output = self.prm(sequences, attention_mask, step_positions)
        
        # 对每一步的 process reward 做均值
        process_rewards = []
        for batch_idx, positions in enumerate(step_positions):
            step_rewards = prm_output["step_rewards"][batch_idx]
            # 使用 soft 聚合：越靠后的步骤权重越高
            weights = torch.linspace(0.5, 1.5, len(step_rewards), device=step_rewards.device)
            weights = weights / weights.sum()
            weighted_reward = (step_rewards * weights).sum()
            process_rewards.append(weighted_reward)
        
        process_rewards = torch.stack(process_rewards)
        
        # 3. 混合奖励
        hybrid_rewards = (
            self.outcome_weight * outcome_rewards +
            self.process_weight * process_rewards
        )
        
        reward_components = {
            "outcome_reward": outcome_rewards.mean().item(),
            "process_reward": process_rewards.mean().item(),
            "hybrid_reward": hybrid_rewards.mean().item(),
        }
        
        return hybrid_rewards, reward_components
    
    def train_step(
        self,
        query: torch.Tensor,           # prompt
        response: torch.Tensor,        # model response (推理链)
        response_mask: torch.Tensor,   # response attention mask
        step_positions: List[List[int]],
        ground_truth: torch.Tensor,    # 最终答案正确标记
        old_logprobs: torch.Tensor,    # 旧策略的 log probabilities
    ) -> Dict[str, float]:
        
        # 计算混合奖励
        rewards, reward_info = self.compute_hybrid_reward(
            response, response_mask, step_positions, ground_truth
        )
        
        # Advantage 估计
        with torch.no_grad():
            values = self.value(response, response_mask).squeeze(-1)
            advantages = rewards - values
        
        # PPO 更新
        for _ in range(self.ppo_epochs):
            new_logits = self.policy(query, response, response_mask).logits
            new_logprobs = self._get_logprobs(new_logits, response, response_mask)
            
            ratio = torch.exp(new_logprobs - old_logprobs)
            clipped_ratio = torch.clamp(ratio, 0.8, 1.2)
            
            policy_loss = -torch.min(
                ratio * advantages,
                clipped_ratio * advantages
            ).mean()
            
            kl_loss = (old_logprobs - new_logprobs).mean()
            total_loss = policy_loss + self.kl_coef * kl_loss
            
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 1.0)
            self.optimizer.step()
            self.optimizer.zero_grad()
        
        return {
            "loss": total_loss.item(),
            "policy_loss": policy_loss.item(),
            "kl": kl_loss.item(),
            **reward_info,

        }
    
    def _get_logprobs(
        self, logits: torch.Tensor, tokens: torch.Tensor, mask: torch.Tensor
    ) -> torch.Tensor:
        logprobs = F.log_softmax(logits, dim=-1)
        token_logprobs = torch.gather(
            logprobs[:, :-1, :], 2, tokens[:, 1:].unsqueeze(-1)
        ).squeeze(-1)
        return (token_logprobs * mask[:, 1:]).sum(dim=-1) / mask[:, 1:].sum(dim=-1)

class PRMEnhancedGRPOTrainer:
    """
    GRPO + PRM 训练器
    
    GRPO 在每组采样内计算相对优势，PRM 提供每步的过程奖励。
    组合方案比 PPO 更稳定（不需要 value network），且采样效率更高。
    """
    
    def __init__(
        self,
        policy_model: nn.Module,
        process_reward_model: ProcessRewardModel,
        group_size: int = 8,           # 每组采样数
        top_k_rejected: int = 2,       # 每组排除的样本数
        kl_coef: float = 0.04,
        learning_rate: float = 3e-6,
        beta: float = 0.1,             # GRPO clipping parameter
    ):
        self.policy = policy_model
        self.prm = process_reward_model
        self.group_size = group_size
        self.top_k_rejected = top_k_rejected
        self.kl_coef = kl_coef
        self.beta = beta
        self.optimizer = torch.optim.AdamW(
            self.policy.parameters(), lr=learning_rate
        )
    
    def compute_group_relative_process_rewards(
        self,
        responses: torch.Tensor,        # (group_size, seq_len)
        masks: torch.Tensor,            # (group_size, seq_len)
        step_positions_list: List[List[List[int]]],
        ground_truths: torch.Tensor,    # (group_size,)
    ) -> torch.Tensor:
        """
        计算组内相对过程奖励
        
        核心思路：
        1. 计算每条推理链的 PRM 分数
        2. 在组内归一化（减去组均值，除以组标准差）
        3. 得到相对优势信号
        """
        group_scores = []
        
        for i in range(len(responses)):
            with torch.no_grad():
                prm_out = self.prm(
                    responses[i].unsqueeze(0),
                    masks[i].unsqueeze(0),
                    step_positions_list[i]
                )
            
            step_rewards = prm_out["step_rewards"][0]
            
            # 如果最终答案正确，boost 所有步骤奖励
            if ground_truths[i] > 0:
                step_rewards = step_rewards + 0.3
            
            # 聚合步骤奖励到序列级别
            seq_score = step_rewards.mean().item()
            group_scores.append(seq_score)
        
        group_scores = torch.tensor(group_scores)
        
        # 组内归一化
        mean_score = group_scores.mean()
        std_score = group_scores.std() + 1e-8
        normalized_advantages = (group_scores - mean_score) / std_score
        
        # 可选：剔除 top-k 最低分样本
        _, rejected_indices = torch.topk(
            normalized_advantages, 
            self.top_k_rejected, 
            largest=False
        )
        normalized_advantages[rejected_indices] = 0.0  # 忽略错误样本
        
        return normalized_advantages

class StepLevelBestOfN:
    """
    步骤级 Best-of-N 采样
    
    在推理的每一步，使用 PRM 从多个候选步骤中选择最佳步骤，
    而不是在整个推理结束时做 BoN。显著提升采样效率。
    """
    
    def __init__(
        self,
        policy: nn.Module,
        prm: ProcessRewardModel,
        candidates_per_step: int = 5,  # 每步候选数
        max_steps: int = 20,
        threshold: float = 0.3,        # 接受阈值
    ):
        self.policy = policy
        self.prm = prm
        self.candidates_per_step = candidates_per_step
        self.max_steps = max_steps
        self.threshold = threshold
    
    def generate(self, prompt: str) -> Tuple[List[str], float]:
        """步骤级 BoN 生成"""
        steps = []
        current_text = prompt
        
        for step_idx in range(self.max_steps):
            candidates = []
            
            for _ in range(self.candidates_per_step):
                # 采样一个候选步骤
                candidate = self._sample_next_step(current_text)
                candidates.append(candidate)
            
            # 使用 PRM 评估每个候选
            best_candidate = None
            best_score = -float("inf")
            
            for candidate in candidates:
                candidate_steps = steps + [candidate]
                score = self._evaluate_steps(candidate_steps)
                
                if score > best_score:
                    best_score = score
                    best_candidate = candidate
            
            if best_candidate is None or best_score < self.threshold:
                break
            
            steps.append(best_candidate)
            current_text = current_text + "\n" + best_candidate
            
            # 如果生成结束，退出
            if self._is_final_step(best_candidate):
                break
        
        return steps, best_score
    
    def _evaluate_steps(self, steps: List[str]) -> float:
        """使用 PRM 评估步骤序列的健康度"""
        # Tokenize
        text = "\n".join(steps)
        tokens = self.prm.tokenizer(text, return_tensors="pt")
        
        with torch.no_grad():
            output = self.prm(tokens.input_ids, tokens.attention_mask, None)
        
        # 如果没有指定步骤位置，用自动检测
        step_positions = self._detect_step_positions(tokens.input_ids[0], steps)
        
        with torch.no_grad():
            detailed_output = self.prm(
                tokens.input_ids, tokens.attention_mask, [step_positions]
            )
        
        step_rewards = detailed_output["step_rewards"][0]
        
        # 最新的步骤权重最高
        weights = torch.linspace(0.5, 1.0, len(step_rewards))
        weights = weights / weights.sum()
        
        return (step_rewards * weights).sum().item()

@dataclass
class ProcessSupervisionPipelineConfig:
    """过程监督管线配置"""
    # 数据生成
    auto_label_strategies: List[str] = field(
        default_factory=lambda: ["rollout", "mc_dropout", "mcts"]
    )
    min_cross_validation_count: int = 2
    human_verification_ratio: float = 0.05
    
    # 模型
    prm_base_model: str = "mistralai/Mistral-7B-v0.1"
    policy_model: str = "Qwen/Qwen2.5-7B-Instruct"
    
    # 训练
    training_method: str = "grpo"  # "ppo" or "grpo"
    group_size: int = 8
    process_weight: float = 0.7
    outcome_weight: float = 0.3
    
    # 评估
    eval_datasets: List[str] = field(
        default_factory=lambda: ["math500", "gsm8k", "agentbench"]
    )
    eval_metrics: List[str] = field(
        default_factory=lambda: [
            "step_accuracy", "outcome_accuracy", "process_consistency"
        ]
    )
    
    # 部署
    prm_cache_size: int = 10000
    prm_cache_ttl: int = 3600
    
    # 监控
    log_frequency: int = 10
    checkpoint_frequency: int = 100
    wandb_project: str = "process-supervision"

class ProcessSupervisionPipeline:
    """
    生产级 Process Supervision 训练与部署管线
    整合数据生成、模型训练、评估、部署全流程
    """
    
    def __init__(self, config: ProcessSupervisionPipelineConfig):
        self.config = config
        self.data_generator = ProcessSupervisionDataGenerator(config)
        self.trainer = self._init_trainer()
        self.evaluator = ProcessSupervisionEvaluator(config)
        self.prm_server = None
    
    def _init_trainer(self):
        prm = ProcessRewardModel(self.config.prm_base_model)
        
        if self.config.training_method == "ppo":
            value_model = AutoModelForCausalLM.from_pretrained(
                self.config.policy_model
            )
            return PRMEnhancedPPOTrainer(
                policy_model=None,  # set during training
                process_reward_model=prm,
                value_model=value_model,
                outcome_weight=self.config.outcome_weight,
                process_weight=self.config.process_weight,
            )
        else:
            return PRMEnhancedGRPOTrainer(
                policy_model=None,
                process_reward_model=prm,
                group_size=self.config.group_size,
            )
    
    def generate_training_data(
        self, 
        source_questions: List[str], 
        ground_truth_answers: List[str]
    ) -> ProcessSupervisionDataset:
        """自动生成过程监督训练数据"""
        labeler = AutoStepLabeler()
        quality_ensurer = StepLabelQualityEnsurer()
        dataset = ProcessSupervisionDataset()
        
        for question, answer in tqdm(
            zip(source_questions, ground_truth_answers),
            desc="Generating process supervision data"
        ):
            # 1. 生成推理链
            reasoning_steps = self._generate_reasoning(question)
            
            # 2. 自动标注每一步
            cross_validated_labels = quality_ensurer.cross_validate_labels(
                reasoning_steps, answer
            )
            
            # 3. 只保留高置信度标注
            high_conf_labels = [
                l for l in cross_validated_labels 
                if l.confidence > 0.6 and l.is_correct is not None
            ]
            
            if len(high_conf_labels) >= len(reasoning_steps) * 0.5:
                dataset.add(question, reasoning_steps, high_conf_labels, answer)
        
        return dataset
    
    def train(self, dataset: ProcessSupervisionDataset):
        """训练 PRM"""
        logger.info(f"Starting PRM training with {len(dataset)} examples")
        
        for epoch in range(3):
            for batch in dataset.batches(self.config.group_size):
                metrics = self.trainer.train_step(**batch)
                
                # 日志
                if self.trainer.step % self.config.log_frequency == 0:
                    logger.info(f"Step {self.trainer.step}: {metrics}")
                    if hasattr(self.trainer, 'wandb'):
                        self.trainer.wandb.log(metrics)
                
                # Checkpoint
                if self.trainer.step % self.config.checkpoint_frequency == 0:
                    self._save_checkpoint(epoch, self.trainer.step)
        
        logger.info("PRM training complete")
    
    def evaluate(self, prm: ProcessRewardModel) -> EvaluationReport:
        """多维度评估"""
        report = EvaluationReport()
        
        for dataset_name in self.config.eval_datasets:
            dataset = self._load_eval_dataset(dataset_name)
            
            for question, reasoning, ground_truth in dataset:
                # PRM 评估
                prm_scores = self._prm_evaluate(prm, question, reasoning)
                
                # 步骤准确率
                report.add_metric(
                    dataset_name, "step_accuracy",
                    self._compute_step_accuracy(prm_scores, ground_truth)
                )
                
                # 结果准确率
                report.add_metric(
                    dataset_name, "outcome_accuracy",
                    self._compute_outcome_accuracy(prm_scores, ground_truth)
                )
                
                # 过程一致性
                report.add_metric(
                    dataset_name, "process_consistency",
                    self._compute_process_consistency(prm_scores)
                )
        
        return report
    
    def deploy(self, prm_checkpoint_path: str):
        """部署 PRM 为推理服务"""
        self.prm_server = ProcessRewardModelServer(
            checkpoint_path=prm_checkpoint_path,
            cache_size=self.config.prm_cache_size,
            cache_ttl=self.config.prm_cache_ttl,
        )
        return self.prm_server


class ProcessRewardModelServer:
    """
    PRM 推理服务
    
    为生产环境的 Agent 提供实时的步骤级奖励评估
    """
    
    def __init__(
        self,
        checkpoint_path: str,
        cache_size: int = 10000,
        cache_ttl: int = 3600,
        device: str = "cuda",
    ):
        self.model = ProcessRewardModel.from_pretrained(checkpoint_path)
        self.model.to(device)
        self.model.eval()
        self.device = device
        self.cache = TTLCache(maxsize=cache_size, ttl=cache_ttl)
        self.request_counter = Counter()
    
    @torch.no_grad()
    def score_step(
        self,
        conversation: List[Dict],
        step_text: str,
        step_index: int,
        use_cache: bool = True,
    ) -> Dict:
        """
        对单步推理进行评分
        
        Args:
            conversation: 对话/推理历史
            step_text: 当前步骤文本
            step_index: 步骤序号
            
        Returns:
            {"step_score": float, "confidence": float, "analysis": str}
        """
        cache_key = hash((json.dumps(conversation[-3:]), step_text))
        
        if use_cache and cache_key in self.cache:
            return self.cache[cache_key]
        
        # 构建上下文
        context = self._format_conversation(conversation)
        full_text = context + "\n" + step_text
        
        inputs = self.model.tokenizer(full_text, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        outputs = self.model(**inputs)
        
        # 提取步骤位置（倒数第二段）
        step_start = len(context.split()) + 1
        step_end = inputs["input_ids"].shape[1] - 1
        positions = [[step_end]]
        
        detailed = self.model(inputs["input_ids"], inputs["attention_mask"], positions)
        step_score = detailed["step_rewards"][0][0].item()
        
        # 置信度估计（基于表示层的方差）
        step_score = np.clip(step_score, 0.0, 1.0)
        confidence = self._estimate_confidence(inputs["input_ids"], positions)
        
        result = {
            "step_score": step_score,
            "confidence": confidence,
            "analysis": self._generate_analysis(step_text, step_score, confidence),
            "is_correct": step_score > 0.5,
        }
        
        self.cache[cache_key] = result
        self.request_counter["total_evaluations"] += 1
        
        return result
    
    def analyze_trajectory(
        self,
        conversation: List[Dict],
        steps: List[str],
    ) -> TrajectoryAnalysis:
        """
        对完整推理轨迹进行分析
        返回每步评分 + 错误定位
        """
        step_scores = []
        step_confidences = []
        
        for i, step in enumerate(steps):
            result = self.score_step(conversation[:-(len(steps)-i)], step, i)
            step_scores.append(result["step_score"])
            step_confidences.append(result["confidence"])
        
        # 检测错误步骤
        error_steps = self._detect_error_steps(step_scores, step_confidences)
        
        # 生成改进建议
        improvement_suggestions = self._generate_suggestions(
            steps, step_scores, error_steps
        )
        
        return TrajectoryAnalysis(
            step_scores=step_scores,
            step_confidences=step_confidences,
            overall_score=np.mean(step_scores),
            error_steps=error_steps,
            improvements=improvement_suggestions,
        )
    
    def _detect_error_steps(
        self, scores: List[float], confidences: List[float]
    ) -> List[ErrorStep]:
        """检测推理链中的错误步骤"""
        errors = []
        threshold = 0.4
        
        for i, (score, conf) in enumerate(zip(scores, confidences)):
            if score < threshold and conf > 0.6:
                errors.append(ErrorStep(
                    step_index=i,
                    severity=1.0 - score,
                    confidence=conf,
                    likely_cause=self._classify_error(scores, i),
                ))
        
        return errors

    
    def _classify_error(self, scores: List[float], error_idx: int) -> str:
        """分类错误类型"""
        if error_idx > 0:
            prev_score = scores[error_idx - 1]
            # 如果前一步分数正常，突然变低 = 逻辑跳转错误
            if prev_score > 0.6:
                return "logic_jump_error"
            # 如果前一步也低 = 累积错误
            else:
                return "accumulated_error"
        else:
            return "initial_step_error"