# Stack 1: Naive Retry — 问题比你想的更严重
class NaiveRetryAgent:
    def __init__(self, llm, max_retries=3):
        self.llm = llm
        self.max_retries = max_retries

    def run(self, task: str) -> str:
        for attempt in range(self.max_retries):
            response = self.llm.generate(task)
            # "自我审查"——问题所在：模型没有额外信息来判断对错
            review = self.llm.generate(
                f"Review this answer for correctness:\n"
                f"Task: {task}\nAnswer: {response}\n"
                f"Is it correct? If not, provide the correct answer."
            )
            if "correct" in review.lower() and "not" not in review.lower():
                return response
            # 用审查结果替代原始答案——但审查本身可能更错
            response = review
        return response

# Stack 2: Self-Refine — 结构化三角色循环
class SelfRefineAgent:
    def __init__(self, llm, max_rounds=3):
        self.llm = llm
        self.max_rounds = max_rounds

    def generate(self, task: str) -> str:
        output = self.llm.generate(f"Generate: {task}")
        for _ in range(self.max_rounds):
            # Critic角色：给出具体的、结构化的反馈
            critique = self.llm.generate(
                f"Act as a critic. Evaluate this output:\n"
                f"Task: {task}\nOutput: {output}\n\n"
                f"Provide specific feedback on:\n"
                f"1. Accuracy issues\n2. Completeness gaps\n"
                f"3. Style/quality improvements\n"
                f"Be concrete — cite specific parts that need change."
            )
            # Refiner角色：只改Critic指出的部分
            output = self.llm.generate(
                f"Based on this feedback, refine the output:\n"
                f"Original: {output}\nFeedback: {critique}\n"
                f"Only change what the feedback specifically addresses."
            )
        return output

# Stack 3: Reflexion — 语言强化学习的完整实现
import json
from dataclasses import dataclass, field

@dataclass
class ReflexionMemory:
    entries: list = field(default_factory=list)

    def add(self, task: str, trajectory: str,
            reward: float, reflection: str):
        self.entries.append({
            "task_summary": task[:100],
            "reward": reward,
            "reflection": reflection,
            "trajectory_snippet": trajectory[-500:]
        })

    def retrieve(self, current_task: str, top_k: int = 3) -> str:
        """基于任务相似度检索相关反思"""
        # 生产环境应使用向量检索
        relevant = sorted(
            self.entries,
            key=lambda x: len(
                set(current_task.split()) &
                set(x["task_summary"].split())
            ), reverse=True
        )[:top_k]
        return "\n".join(
            f"Past experience (reward={e['reward']}):\n"
            f"{e['reflection']}\n"
            for e in relevant
        )

class ReflexionAgent:
    def __init__(self, llm, evaluator, max_episodes=5):
        self.llm = llm
        self.evaluator = evaluator
        self.memory = ReflexionMemory()
        self.max_episodes = max_episodes

    def act(self, task: str) -> str:
        past_reflections = self.memory.retrieve(task)
        context = f"Task: {task}\n"
        if past_reflections:
            context += (
                f"\n--- Past Reflections ---\n"
                f"{past_reflections}\n"
                f"--- End Reflections ---\n\n"
                f"Learn from these past experiences.\n"
            )
        trajectory = self._run_episode(context)
        reward = self.evaluator.score(task, trajectory)

        if reward < 1.0:
            reflection = self.llm.generate(
                f"You attempted this task and scored {reward}.\n"
                f"Trajectory:\n{trajectory}\n\n"
                f"Analyze your failures:\n"
                f"1. What specific mistakes did you make?\n"
                f"2. Why did each mistake occur?\n"
                f"3. What would you do differently?"
            )
            self.memory.add(task, trajectory, reward, reflection)
        return trajectory

    def _run_episode(self, context: str) -> str:
        return self.llm.generate(context)

# Stack 4: LATS 核心循环（简化版）
import math
from typing import Optional

class MCTSNode:
    def __init__(self, state: str, parent=None):
        self.state = state
        self.parent = parent
        self.children = []
        self.visits = 0
        self.value = 0.0
        self.reflection: Optional[str] = None

    def ucb(self, exploration=1.414) -> float:
        if self.visits == 0:
            return float('inf')
        exploit = self.value / self.visits
        explore = exploration * math.sqrt(
            math.log(self.parent.visits) / self.visits
        )
        return exploit + explore

    def best_child(self) -> 'MCTSNode':
        return max(self.children, key=lambda c: c.ucb())

class LATSAgent:
    def __init__(self, llm, evaluator, n_candidates=3,
                 max_depth=10, n_simulations=20):
        self.llm = llm
        self.evaluator = evaluator
        self.n_candidates = n_candidates
        self.max_depth = max_depth
        self.n_simulations = n_simulations

    def run(self, task: str) -> str:
        root = MCTSNode(state=task)
        for _ in range(self.n_simulations):
            node = self._select(root)
            if self._is_terminal(node):
                reward = self.evaluator.score(node.state)
            else:
                children = self._expand(node)
                reward = self._simulate(
                    children[0]  # 简化：取第一个子节点模拟
                )
            self._backpropagate(node, reward)
        return self._best_path(root)

    def _expand(self, node: MCTSNode) -> list:
        candidates = self.llm.generate(
            f"Generate {self.n_candidates} different next steps:\n"
            f"Current state: {node.state}\n"
            f"Provide each as a separate numbered option.",
            n=self.n_candidates
        )
        for c in candidates:
            child = MCTSNode(state=c, parent=node)
            # LATS特色：用LLM自我评估节点价值
            child.reflection = self.llm.generate(
                f"Evaluate this intermediate state:\n{c}\n"
                f"Rate progress toward solving the task (1-10)."
            )
            node.children.append(child)
        return node.children

    def _simulate(self, node: MCTSNode) -> float:
        """快速rollout：用LLM完成剩余步骤并评估"""
        rollout = self.llm.generate(
            f"Quickly complete from this state:\n{node.state}"
        )
        return self.evaluator.score(rollout)

    def _backpropagate(self, node: MCTSNode, reward: float):
        while node:
            node.visits += 1
            node.value += reward
            node = node.parent

    def _best_path(self, root: MCTSNode) -> str:
        node = root
        while node.children:
            node = max(node.children, key=lambda c: c.visits)
        return node.state

# Stack 5: AgentPRM — 过程奖励驱动的自纠
class AgentPRMSelfCorrection:
    def __init__(self, actor_llm, prm_model,
                 n_samples=5, threshold=0.7):
        self.actor = actor_llm
        self.prm = prm_model  # 过程奖励模型
        self.n_samples = n_samples
        self.threshold = threshold

    def run(self, task: str) -> str:
        # 阶段1：生成N条候选轨迹
        trajectories = []
        for i in range(self.n_samples):
            traj = self._generate_trajectory(task)
            trajectories.append(traj)

        # 阶段2：用PRM对每条轨迹逐步打分
        scored = []
        for traj in trajectories:
            step_scores = []
            for step_idx, step in enumerate(traj.steps):
                # Promise：这一步声称做什么
                promise = self.actor.generate(
                    f"What does this step promise to accomplish?\n"
                    f"Step: {step.content}\n"
                    f"State: {step.context}"
                )
                # Progress：实际推进了多少
                progress_score = self.prm.score(
                    task=task,
                    step=step.content,
                    promise=promise,
                    history=traj.steps[:step_idx]
                )
                step_scores.append(progress_score)

            # 轨迹总分 = min(逐步分数) — 木桶原理
            traj_score = min(step_scores) if step_scores else 0
            scored.append((traj, traj_score, step_scores))

        # 阶段3：选择最佳轨迹
        scored.sort(key=lambda x: x[1], reverse=True)
        best_traj, best_score, best_step_scores = scored[0]

        # 阶段4：对低于阈值的步骤进行定向修正
        corrected_steps = list(best_traj.steps)
        for i, score in enumerate(best_step_scores):
            if score < self.threshold:
                corrected = self.actor.generate(
                    f"This step was rated {score:.2f}/1.0.\n"
                    f"Original step: {corrected_steps[i].content}\n"
                    f"Specific issue: low progress toward goal.\n"
                    f"Rewrite this step to be more effective."
                )
                corrected_steps[i] = corrected

        return corrected_steps

    def _generate_trajectory(self, task: str):
        return self.actor.generate(task)

graph TD
    subgraph "Stack 1: Naive Retry"
        A1[执行任务] --> A2{自我审查}
        A2 -->|通过| A3[返回结果]
        A2 -->|失败| A1
    end

    subgraph "Stack 2: Self-Refine"
        B1[Generator<br>生成初始输出] --> B2[Critic<br>结构化评价]
        B2 --> B3[Refiner<br>定向修改]
        B3 --> B2
    end

    subgraph "Stack 3: Reflexion"
        C1[Actor执行] --> C2[Evaluator<br>外部评分]
        C2 -->|失败| C3[Self-Reflection<br>语言化失败分析]
        C3 --> C4[Memory Buffer<br>存储反思]
        C4 --> C1
    end

    subgraph "Stack 4: LATS/MCTS"
        D1[Selection<br>UCB选节点] --> D2[Expansion<br>LLM生成候选]
        D2 --> D3[Simulation<br>快速rollout]
        D3 --> D4[Backpropagation<br>回传更新]
        D4 --> D1
    end

    subgraph "Stack 5: PRM"
        E1[并行采样N条轨迹] --> E2[PRM逐步打分]
        E2 --> E3[选最优轨迹]
        E3 --> E4[低分步骤定向修正]
        E4 --> E5[最终输出]
    end

    A3 --> B1
    B3 --> C1
    C4 --> D1
    D4 --> E1

    style A1 fill:#ffcccc
    style B1 fill:#ffe0b2
    style C1 fill:#fff9c4
    style D1 fill:#c8e6c9
    style E1 fill:#bbdefb

graph LR
    START[任务类型?] --> Q1{需要精确<br>推理?}
    Q1 -->|否| Q2{有外部<br>验证器?}
    Q1 -->|是| Q3{预算允许<br>PRM训练?}

    Q2 -->|是| R1[Stack 1<br>简单重试+外部信号]
    Q2 -->|否| R2[Stack 2<br>Self-Refine]

    Q3 -->|是| Q4{延迟<br>敏感?}
    Q3 -->|否| Q5{有记忆<br>需求?}

    Q4 -->|是| R5[Stack 5<br>PRM并行采样]
    Q4 -->|否| R4[Stack 4<br>LATS深度搜索]

    Q5 -->|是| R3[Stack 3<br>Reflexion]
    Q5 -->|否| R2

    style R1 fill:#ffcccc
    style R2 fill:#ffe0b2
    style R3 fill:#fff9c4
    style R4 fill:#c8e6c9
    style R5 fill:#bbdefb