import torch
import torch.nn.functional as F
from typing import List, Optional

class EarlyStopReasoner:
    def __init__(
        self,
        model,
        tokenizer,
        entropy_threshold: float = 0.3,
        min_steps: int = 10,
        max_steps: int = 1000,
        window_size: int = 5,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.entropy_threshold = entropy_threshold
        self.min_steps = min_steps
        self.max_steps = max_steps
        self.window_size = window_size
        
    def _estimate_answer_entropy(
        self, 
        logits: torch.Tensor, 
        generated_ids: List[int]
    ) -> float:
        """估计当前推理状态下最终答案的条件熵"""
        
        # 获取候选答案 token 的概率分布
        probs = F.softmax(logits[0, -1, :], dim=-1)
        
        # 只考虑词汇表顶部的 K 个候选
        top_k = 50
        top_probs, top_indices = torch.topk(probs, top_k)
        
        # 计算部分熵 (partial entropy)
        entropy = -torch.sum(top_probs * torch.log(top_probs + 1e-10))
        
        # 归一化到 [0, 1]
        max_entropy = torch.log(torch.tensor(top_k, dtype=torch.float))
        normalized_entropy = entropy / max_entropy
        
        return normalized_entropy.item()
    
    @torch.no_grad()
    def generate(self, prompt: str) -> tuple[str, int, float]:
        inputs = self.tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"]
        attention_mask = inputs.get("attention_mask")
        
        # 将输入移到模型所在设备
        device = next(self.model.parameters()).device
        input_ids = input_ids.to(device)
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)
        
        past_key_values = None
        generated = input_ids[0].tolist()
        entropy_history = []
        early_stopped = False
        
        for step in range(self.max_steps):
            with torch.no_grad():
                outputs = self.model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    past_key_values=past_key_values,
                    use_cache=True,
                )
                
                logits = outputs.logits
                past_key_values = outputs.past_key_values
                
                # 计算不确定性
                if step >= self.min_steps:
                    entropy = self._estimate_answer_entropy(logits, generated)
                    entropy_history.append(entropy)
                    
                    # 使用滑动窗口平均平滑不确定性估计
                    if len(entropy_history) >= self.window_size:
                        avg_entropy = sum(
                            entropy_history[-self.window_size:]
                        ) / self.window_size
                        
                        if avg_entropy < self.entropy_threshold:
                            early_stopped = True
                            break
                
                # 贪婪解码（简化，实际可以使用采样）
                next_token = logits[0, -1].argmax(dim=-1).unsqueeze(0)
                input_ids = next_token.unsqueeze(0)
                generated.append(next_token.item())
                
                if next_token.item() == self.tokenizer.eos_token_id:
                    break
        
        decoded = self.tokenizer.decode(generated)
        total_steps = len(entropy_history) if early_stopped else step + 1
        
        return decoded, total_steps, entropy_history[-1] if entropy_history else 1.0

import torch
import torch.nn as nn

class RewardPredictor(nn.Module):
    """轻量级奖励预测器，基于当前的 hidden state 预测最终答案正确概率"""
    
    def __init__(self, hidden_size: int = 4096):
        super().__init__()
        self.predictor = nn.Sequential(
            nn.Linear(hidden_size, 512),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(512, 128),
            nn.GELU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )
        
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # hidden_states: [batch, seq_len, hidden_size]
        # 使用最后一个位置的 hidden state
        last_hidden = hidden_states[:, -1, :]  # [batch, hidden_size]
        reward = self.predictor(last_hidden)   # [batch, 1]
        return reward.squeeze(-1)              # [batch]

class AdaptiveReasoningController:
    """
    多级自适应推理控制器
    策略 1: 简单问题直接输出（< 200 tokens）
    策略 2: 中等复杂度问题使用 Early Stopping
    策略 3: 复杂问题使用完整推理 + Re-FORC
    """
    
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
        self.es_reasoner = EarlyStopReasoner(model, tokenizer)
        
    def estimate_complexity(self, prompt: str) -> str:
        """快速估算问题复杂度"""
        # 使用一个小型分类器或启发式规则
        prompt_len = len(prompt.split())
        
        if prompt_len < 20 and "?" in prompt:
            return "simple"
        elif "prove" in prompt.lower() or "explain" in prompt.lower():
            return "complex"
        else:
            return "medium"
    
    def generate(self, prompt: str) -> tuple[str, int]:
        complexity = self.estimate_complexity(prompt)
        
        if complexity == "simple":
            # 直接贪婪解码，限制最大 token 数
            output = self._direct_generate(prompt, max_tokens=256)
            return output, output["usage"]["completion_tokens"]
        
        elif complexity == "medium":
            # Early Stopping
            text, steps, final_entropy = self.es_reasoner.generate(prompt)
            return text, steps
        
        else:
            # 完整推理
            return self._full_reasoning(prompt)

# vLLM 配置示例：启用 prefix caching + 动态批处理
from vllm import LLM, SamplingParams

llm = LLM(
    model="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    tensor_parallel_size=2,
    max_model_len=32768,
    enable_prefix_caching=True,    # 缓存公共前缀
    max_num_batched_tokens=65536,   # 支持长序列批处理
    gpu_memory_utilization=0.90,   # 留出 KV cache 空间
)

sampling_params = SamplingParams(
    temperature=0.0,
    max_tokens=16384,
    early_stopping=True,           # 使用 vLLM 的 early stopping
    stop_token_ids=[151643],       # 自定义停止 token
)

def prepare_reasoning_training_data(trajectories):
    """
    准备训练数据：在推理链中插入 <STOP_REASONING> token
    当模型已经得出正确答案时，插入停止标记
    """
    processed = []
    for prompt, reasoning_chain, answer, is_correct in trajectories:
        if not is_correct:
            continue  # 只使用正确轨迹
        
        # 找到第一次出现正确答案的位置
        first_correct_pos = find_first_correct_position(
            reasoning_chain, answer
        )
        
        # 在该位置后插入 <STOP_REASONING> 和最终答案
        augmented_chain = (
            reasoning_chain[:first_correct_pos]
            + " <STOP_REASONING> "
            + answer
        )
        
        processed.append({
            "prompt": prompt,
            "completion": augmented_chain,
        })
    
    return processed

while True:
    thought = llm.generate(f"当前状态：{state}，下一步做什么？")
    action = parse_action(thought)       # 执行工具调用
    result = execute_tool(action)
    state += f"\n执行结果：{result}"