第十四章 Generative Agents与仿真¶

"""Generative Agents — 基础数据模型"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timedelta

@dataclass
class MemoryRecord:
    """记忆流中的一条记录"""
    id: int
    timestamp: datetime
    content: str                         # 记忆内容
    memory_type: str = "observation"     # observation / reflection / plan / dialogue
    importance: int = 5                  # 重要性评分 1-10
    embedding: list[float] | None = None # 向量嵌入，使用联合类型语法(Python 3.10+)，等价于Optional[list[float]]

    # 元数据
    location: str = ""
    related_agents: list[str] = field(default_factory=list)
    parent_ids: list[int] = field(default_factory=list)  # 反思的来源记忆

@dataclass
class AgentProfile:
    """Agent的个人档案"""
    name: str
    age: int
    occupation: str
    personality: str          # 性格描述
    background: str           # 背景故事
    relationships: dict[str, str] = field(default_factory=dict)
    # e.g. {"李明": "大学同学，关系亲密", "王芳": "邻居，偶尔打招呼"}

    daily_routine: str = ""   # 典型日常作息

@dataclass
class Location:
    """地图上的位置"""
    name: str
    x: int
    y: int
    area_type: str = "public"   # public / private / work
    description: str = ""
    occupants: list[str] = field(default_factory=list)  # 当前在这里的Agent

@dataclass
class DayPlan:
    """一天的日程"""
    date: str
    schedule: list[tuple[str, str]] = field(default_factory=list)
    # e.g. [("08:00", "在家吃早餐"), ("09:00", "去咖啡店工作"), ...]

    current_action: str = ""
    current_location: str = ""

"""Memory Stream — 记忆流完整实现"""

import math
import numpy as np
from openai import AsyncOpenAI
from datetime import datetime, timedelta

client = AsyncOpenAI()

class MemoryStream:
    """记忆流：存储、评分、检索Agent的所有记忆

    注意：当前实现使用内存列表存储记忆，所有读写为同步操作。
    生产环境应替换为异步数据库（如 asyncpg / motor）以避免在 async 上下文中阻塞事件循环。
    """

    def __init__(
        self,
        agent_name: str,
        alpha: float = 1.0,    # 近因性权重
        beta: float = 1.0,     # 重要性权重
        gamma: float = 1.0,    # 相关性权重
        decay: float = 0.99,   # 时间衰减系数
    ):
        self.agent_name = agent_name
        self.memories: list[MemoryRecord] = []
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.decay = decay
        self._next_id = 1
        # 反思触发：重要性累积阈值
        self._importance_accumulator = 0.0
        self.reflection_threshold = 50.0

    # ────────── 添加记忆 ──────────

    async def add_memory(
        self,
        content: str,
        timestamp: datetime,
        memory_type: str = "observation",
        location: str = "",
        related_agents: list[str] | None = None,
        parent_ids: list[int] | None = None,
    ) -> MemoryRecord:
        """添加一条新记忆到记忆流"""
        # 1. 评估重要性
        importance = await self._rate_importance(content)

        # 2. 生成嵌入向量
        embedding = await self._get_embedding(content)

        # 3. 创建记忆记录
        record = MemoryRecord(
            id=self._next_id,
            timestamp=timestamp,
            content=content,
            memory_type=memory_type,
            importance=importance,
            embedding=embedding,
            location=location,
            related_agents=related_agents or [],
            parent_ids=parent_ids or [],
        )
        self.memories.append(record)
        self._next_id += 1

        # 4. 累积重要性（用于触发反思）
        self._importance_accumulator += importance

        return record

    # ────────── 检索记忆 ──────────

    async def retrieve(
        self,
        query: str,
        current_time: datetime,
        top_k: int = 10,
    ) -> list[MemoryRecord]:
        """三维加权检索最相关的记忆"""
        if not self.memories:
            return []

        query_embedding = await self._get_embedding(query)

        scored: list[tuple[float, MemoryRecord]] = []  # 存储(综合得分, 记忆记录)的元组列表
        for mem in self.memories:
            # 近因性 (指数衰减)
            hours_ago = (current_time - mem.timestamp).total_seconds() / 3600
            recency = self.decay ** hours_ago

            # 重要性 (归一化到0-1)
            importance = mem.importance / 10.0

            # 相关性 (余弦相似度)
            relevance = self._cosine_similarity(query_embedding, mem.embedding)

            # 加权总分
            score = (
                self.alpha * recency
                + self.beta * importance
                + self.gamma * relevance
            )
            scored.append((score, mem))

        # 按总分排序，返回Top-K
        scored.sort(key=lambda x: x[0], reverse=True)  # 按三维加权总分降序排列
        return [mem for _, mem in scored[:top_k]]

    # ────────── 重要性评分 ──────────

    async def _rate_importance(self, content: str) -> int:
        """让LLM对记忆内容的重要性打分（1-10）"""
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    "你是一个记忆重要性评估器。给定一段记忆内容，"
                    "请评估其对于一个人的日常生活的重要性。\n"
                    "评分标准：\n"
                    "1-2: 完全日常（如「刷牙」「走路去厨房」）\n"
                    "3-4: 一般日常（如「和同事聊了天气」）\n"
                    "5-6: 有一定意义（如「学到了新技能」「认识了新朋友」）\n"
                    "7-8: 重要（如「被上司表扬」「发现重要信息」）\n"
                    "9-10: 非常重要（如「求婚成功」「亲人去世」）\n"
                    "只输出一个整数，不要任何其他文字。"
                )},
                {"role": "user", "content": content},
            ],
            temperature=0.1,
            max_tokens=5,
        )
        try:
            score = int(response.choices[0].message.content.strip())
            return max(1, min(10, score))
        except ValueError:
            return 5  # 默认中等重要性

    # ────────── 嵌入向量 ──────────

    async def _get_embedding(self, text: str) -> list[float]:
        """获取文本的嵌入向量"""
        response = await client.embeddings.create(
            model="text-embedding-3-small",
            input=text,
        )
        return response.data[0].embedding

    @staticmethod
    def _cosine_similarity(a: list[float] | None, b: list[float] | None) -> float:
        """计算两个向量的余弦相似度"""
        if a is None or b is None:
            return 0.0
        # np.array将list转为numpy数组；np.linalg.norm计算L2范数(向量长度)；np.dot计算点积
        a_arr = np.array(a)
        b_arr = np.array(b)
        denom = np.linalg.norm(a_arr) * np.linalg.norm(b_arr)
        if denom == 0:
            return 0.0
        return float(np.dot(a_arr, b_arr) / denom)

    # ────────── 辅助方法 ──────────

    def should_reflect(self) -> bool:
        """判断是否应该触发反思"""
        return self._importance_accumulator >= self.reflection_threshold

    def reset_importance_accumulator(self):
        self._importance_accumulator = 0.0

    def get_recent(self, n: int = 20) -> list[MemoryRecord]:
        """获取最近N条记忆"""
        return sorted(self.memories, key=lambda m: m.timestamp, reverse=True)[:n]  # sorted不修改原列表，[:n]取前n条

async def memory_demo():
    """记忆流使用示例"""
    ms = MemoryStream(agent_name="陈思远")

    base_time = datetime(2025, 6, 15, 8, 0)

    # 添加一天中的记忆
    await ms.add_memory(
        "起床，感觉今天精神不错",
        base_time, memory_type="observation", location="家",
    )
    await ms.add_memory(
        "在咖啡店遇到了老朋友林悦，她告诉我下周六有社区聚会",
        base_time + timedelta(hours=2), memory_type="dialogue",
        location="星巴克", related_agents=["林悦"],
    )
    await ms.add_memory(
        "上午完成了项目报告的初稿",
        base_time + timedelta(hours=3), memory_type="observation",
        location="办公室",
    )
    await ms.add_memory(
        "午饭时听说公司可能要裁员",
        base_time + timedelta(hours=5), memory_type="observation",
        location="食堂",
    )

    # 检索与"社区活动"相关的记忆
    current_time = base_time + timedelta(hours=6)
    results = await ms.retrieve("社区活动和聚会", current_time, top_k=3)

    print("检索结果：")
    for mem in results:
        print(f"  [{mem.importance}/10] {mem.content}")

"""Reflection — 反思机制实现"""

class ReflectionEngine:
    """反思引擎：从记忆中提取高阶认知"""

    def __init__(self, memory_stream: MemoryStream):
        self.memory = memory_stream
        self.llm = AsyncOpenAI()

    async def reflect(self, current_time: datetime) -> list[MemoryRecord]:
        """执行一次反思过程，返回生成的反思记忆"""
        # Step 1: 获取最近的重要记忆
        recent = self.memory.get_recent(n=30)
        if len(recent) < 5:
            return []  # 记忆不足，无法反思

        recent_text = "\n".join(
            f"[{m.timestamp.strftime('%H:%M')}] {m.content}" for m in recent[:20]
        )

        # Step 2: 生成反思问题
        questions = await self._generate_reflection_questions(recent_text)

        # Step 3: 对每个问题检索相关记忆并生成反思
        new_reflections: list[MemoryRecord] = []
        for question in questions[:3]:     # 最多3个反思
            # 检索与该问题相关的记忆
            relevant = await self.memory.retrieve(question, current_time, top_k=8)
            evidence_text = "\n".join(f"- {m.content}" for m in relevant)
            parent_ids = [m.id for m in relevant]

            # 生成反思
            reflection_content = await self._generate_reflection(
                question, evidence_text,
            )

            # 将反思作为新记忆存入记忆流
            record = await self.memory.add_memory(
                content=reflection_content,
                timestamp=current_time,
                memory_type="reflection",
                parent_ids=parent_ids,
            )
            new_reflections.append(record)

        # 重置重要性累积器
        self.memory.reset_importance_accumulator()
        return new_reflections

    async def _generate_reflection_questions(self, memories_text: str) -> list[str]:
        """基于最近记忆生成反思问题"""
        response = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"基于{self.memory.agent_name}最近的经历，"
                    "提出3个值得深入思考的高层问题。"
                    "这些问题应该关注模式、关系和自我认知。"
                    "每行一个问题，不要编号。"
                )},
                {"role": "user", "content": f"最近的经历:\n{memories_text}"},
            ],
            temperature=0.7,
            max_tokens=200,
        )
        questions = [
            q.strip() for q in response.choices[0].message.content.strip().split("\n")
            if q.strip()
        ]
        return questions

    async def _generate_reflection(
        self,
        question: str,
        evidence: str,
    ) -> str:
        """基于证据生成一条反思"""
        response = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"你是{self.memory.agent_name}的内心思考。"
                    "基于过去的经历，对以下问题进行深入反思。"
                    "输出一句简洁的高层认知（不超过50字）。"
                )},
                {"role": "user", "content": (
                    f"问题: {question}\n\n"
                    f"相关经历:\n{evidence}"
                )},
            ],
            temperature=0.5,
            max_tokens=100,
        )
        return response.choices[0].message.content.strip()

"""Planning System — 规划系统实现"""

import json
from openai import AsyncOpenAI

client = AsyncOpenAI()

class PlanningSystem:
    """Agent规划系统：生成和调整日程"""

    def __init__(self, profile: AgentProfile, memory: MemoryStream):
        self.profile = profile
        self.memory = memory

    async def create_daily_plan(
        self,
        date: str,
        current_time: datetime,
    ) -> DayPlan:
        """创建一天的日程计划"""
        # 检索与日程规划相关的记忆（习惯、承诺等）
        relevant_memories = await self.memory.retrieve(
            f"{self.profile.name}的日常作息和今天的安排",
            current_time,
            top_k=10,
        )
        memory_context = "\n".join(f"- {m.content}" for m in relevant_memories)

        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"你是{self.profile.name}，{self.profile.age}岁，"
                    f"职业是{self.profile.occupation}。\n"
                    f"性格: {self.profile.personality}\n"
                    f"背景: {self.profile.background}\n"
                    f"日常作息: {self.profile.daily_routine}\n\n"
                    "请为自己制定今天的日程。输出JSON格式：\n"
                    '{"schedule": [["08:00", "活动描述"], ["09:00", "..."], ...]}'
                )},
                {"role": "user", "content": (
                    f"日期: {date}\n"
                    f"相关记忆:\n{memory_context}\n\n"
                    "请制定今天的日程（从起床到睡觉）。"
                )},
            ],
            response_format={"type": "json_object"},
            temperature=0.8,
        )

        data = json.loads(response.choices[0].message.content)
        return DayPlan(
            date=date,
            schedule=[tuple(item) for item in data["schedule"]],
        )

    async def react_to_event(
        self,
        event: str,
        current_plan: DayPlan,
        current_time: datetime,
    ) -> tuple[str, DayPlan]:
        """对突发事件做出反应，可能调整计划

        Returns:
            (reaction, updated_plan): Agent的反应和调整后的计划
        """
        # 检索相关记忆
        relevant = await self.memory.retrieve(event, current_time, top_k=5)
        memory_context = "\n".join(f"- {m.content}" for m in relevant)

        remaining_schedule = [
            item for item in current_plan.schedule
            if item[0] >= current_time.strftime("%H:%M")
        ]
        schedule_text = "\n".join(f"{t}: {a}" for t, a in remaining_schedule)

        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"你是{self.profile.name}。性格: {self.profile.personality}\n"
                    "一个突发事件发生了。请决定你的反应和是否调整计划。\n"
                    "输出JSON: {\"reaction\": \"你的反应描述\", "
                    "\"adjust_plan\": true/false, "
                    "\"new_schedule\": [[\"时间\", \"活动\"], ...] 如果adjust_plan为true}"
                )},
                {"role": "user", "content": (
                    f"突发事件: {event}\n\n"
                    f"相关记忆:\n{memory_context}\n\n"
                    f"当前剩余计划:\n{schedule_text}"
                )},
            ],
            response_format={"type": "json_object"},
            temperature=0.7,
        )

        data = json.loads(response.choices[0].message.content)
        reaction = data["reaction"]

        if data.get("adjust_plan"):
            # 保留已执行的计划项，替换剩余部分
            executed = [
                item for item in current_plan.schedule
                if item[0] < current_time.strftime("%H:%M")
            ]
            new_items = [tuple(item) for item in data.get("new_schedule", [])]
            current_plan.schedule = executed + new_items

        return reaction, current_plan

"""Agent Interaction — Agent间对话与交互"""

class DialogueSystem:
    """Agent间对话系统"""

    def __init__(self):
        self.llm = AsyncOpenAI()

    async def should_initiate_conversation(
        self,
        agent: AgentProfile,
        other: AgentProfile,
        agent_memory: MemoryStream,
        current_time: datetime,
        context: str = "",
    ) -> bool:
        """判断Agent是否应该主动与对方对话"""
        # 检索与对方相关的记忆
        relevant = await agent_memory.retrieve(
            f"与{other.name}的互动", current_time, top_k=5,
        )
        memory_text = "\n".join(f"- {m.content}" for m in relevant)

        relationship = agent.relationships.get(other.name, "不太熟悉")

        response = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"你是{agent.name}。你现在遇到了{other.name}。\n"
                    f"你与{other.name}的关系: {relationship}\n"
                    f"你的性格: {agent.personality}\n"
                    f"场景: {context}\n"
                    "你会主动打招呼吗？只回答 yes 或 no。"
                )},
                {"role": "user", "content": f"相关记忆:\n{memory_text}"},
            ],
            temperature=0.5,
            max_tokens=5,
        )
        return "yes" in response.choices[0].message.content.lower()

    async def generate_dialogue(
        self,
        agent_a: AgentProfile,
        agent_b: AgentProfile,
        memory_a: MemoryStream,
        memory_b: MemoryStream,
        current_time: datetime,
        location: str,
        max_turns: int = 6,
    ) -> list[tuple[str, str]]:
        """生成两个Agent之间的对话

        Returns:
            [(speaker_name, utterance), ...]
        """
        # 为双方检索相关记忆
        context_a = await memory_a.retrieve(
            f"与{agent_b.name}聊天", current_time, top_k=5,
        )
        context_b = await memory_b.retrieve(
            f"与{agent_a.name}聊天", current_time, top_k=5,
        )

        rel_a = agent_a.relationships.get(agent_b.name, "不太熟悉")
        rel_b = agent_b.relationships.get(agent_a.name, "不太熟悉")

        dialogue: list[tuple[str, str]] = []
        messages = [
            {"role": "system", "content": (
                f"模拟{agent_a.name}和{agent_b.name}在{location}的对话。\n\n"
                f"{agent_a.name}: {agent_a.personality}，与{agent_b.name}的关系: {rel_a}\n"
                f"  相关记忆: {'; '.join(m.content for m in context_a[:3])}\n\n"
                f"{agent_b.name}: {agent_b.personality}，与{agent_a.name}的关系: {rel_b}\n"
                f"  相关记忆: {'; '.join(m.content for m in context_b[:3])}\n\n"
                "对话要自然、符合各自性格。每轮输出格式:\n"
                f"{agent_a.name}: [说话内容]\n{agent_b.name}: [说话内容]\n"
                "当对话自然结束时，输出 [END]"
            )},
        ]

        for turn in range(max_turns // 2):
            messages.append({"role": "user", "content": "请继续对话。"})
            response = await self.llm.chat.completions.create(
                model="gpt-4o-mini",
                messages=messages,
                temperature=0.8,
                max_tokens=200,
            )

            text = response.choices[0].message.content
            messages.append({"role": "assistant", "content": text})

            if "[END]" in text:
                text = text.replace("[END]", "").strip()

            # 解析发言
            # 注意：基于换行符和前缀的解析较脆弱，LLM输出格式可能不一致。
            # 生产环境建议要求LLM以JSON数组格式输出对话，再用json.loads解析。
            for line in text.strip().split("\n"):
                line = line.strip()
                if not line:
                    continue
                if line.startswith(f"{agent_a.name}:"):
                    utterance = line[len(agent_a.name) + 1:].strip()
                    dialogue.append((agent_a.name, utterance))
                elif line.startswith(f"{agent_b.name}:"):
                    utterance = line[len(agent_b.name) + 1:].strip()
                    dialogue.append((agent_b.name, utterance))

            if "[END]" in response.choices[0].message.content:
                break

        return dialogue

async def propagate_information(
    dialogue: list[tuple[str, str]],
    memory_a: MemoryStream,
    memory_b: MemoryStream,
    agent_a_name: str,
    agent_b_name: str,
    current_time: datetime,
    location: str,
):
    """将对话记录存入双方的记忆流——信息通过对话在Agent间传播"""
    # 将完整对话摘要存入记忆
    dialogue_text = "\n".join(f"{name}: {text}" for name, text in dialogue)

    summary_a = f"在{location}和{agent_b_name}聊天。" + (
        dialogue[-1][1] if dialogue else ""  # [-1]取最后一条，[1]取元组第二个元素(对话内容)
    )
    summary_b = f"在{location}和{agent_a_name}聊天。" + (
        dialogue[0][1] if dialogue else ""
    )

    await memory_a.add_memory(
        content=summary_a,
        timestamp=current_time,
        memory_type="dialogue",
        location=location,
        related_agents=[agent_b_name],
    )

    await memory_b.add_memory(
        content=summary_b,
        timestamp=current_time,
        memory_type="dialogue",
        location=location,
        related_agents=[agent_a_name],
    )

"""Environment — 2D网格世界"""

from __future__ import annotations
from dataclasses import dataclass, field
from datetime import datetime, timedelta

@dataclass
class WorldLocation:
    """世界中的一个位置/区域"""
    name: str
    x: int
    y: int
    area_type: str              # home / work / shop / park / public
    description: str = ""
    occupants: set[str] = field(default_factory=set)

class GridWorld:
    """简化的2D网格世界"""

    def __init__(self, width: int = 20, height: int = 20):
        self.width = width
        self.height = height
        self.locations: dict[str, WorldLocation] = {}
        self.agent_positions: dict[str, str] = {}  # agent_name -> location_name
        self.time: datetime = datetime(2025, 6, 15, 6, 0)  # 模拟开始时间

    def add_location(self, loc: WorldLocation):
        self.locations[loc.name] = loc

    def move_agent(self, agent_name: str, location_name: str):
        """移动Agent到指定位置"""
        # 注意：若 location_name 不在 self.locations 中，移动会被静默忽略。
        # 生产环境建议添加边界校验，抛出异常或返回 False 以便调用方感知失败。
        # 离开旧位置
        old_loc = self.agent_positions.get(agent_name)
        if old_loc and old_loc in self.locations:
            self.locations[old_loc].occupants.discard(agent_name)  # discard不会在元素不存在时抛异常，比remove更安全
        # 进入新位置
        if location_name in self.locations:
            self.locations[location_name].occupants.add(agent_name)
            self.agent_positions[agent_name] = location_name

    def get_agents_at(self, location_name: str) -> set[str]:
        """获取某位置的所有Agent"""
        loc = self.locations.get(location_name)
        return loc.occupants if loc else set()

    def get_nearby_agents(self, agent_name: str) -> set[str]:
        """获取与某Agent在同一位置的其他Agent"""
        loc = self.agent_positions.get(agent_name)
        if not loc:
            return set()
        return self.get_agents_at(loc) - {agent_name}

    def advance_time(self, minutes: int = 30):
        """推进模拟时间"""
        self.time += timedelta(minutes=minutes)

    def get_time_str(self) -> str:
        return self.time.strftime("%H:%M")

    def describe_location(self, location_name: str) -> str:
        """描述一个位置的当前状态"""
        loc = self.locations.get(location_name)
        if not loc:
            return "未知位置"
        people = "、".join(loc.occupants) if loc.occupants else "无人"  # 用中文顿号连接当前位置上的所有Agent名称
        return f"{loc.name}（{loc.description}）— 当前: {people}"

    def get_world_state(self) -> str:
        """获取世界当前状态的文字描述"""
        lines = [f"⏰ 时间: {self.get_time_str()}\n"]
        for name, loc in self.locations.items():
            if loc.occupants:
                people = ", ".join(loc.occupants)
                lines.append(f"📍 {name}: {people}")
        return "\n".join(lines)

def create_smallville() -> GridWorld:
    """创建一个迷你小镇地图"""
    world = GridWorld(width=10, height=10)

    locations = [
        WorldLocation("陈思远的家", 1, 1, "home", "一间温馨的公寓"),
        WorldLocation("林悦的家", 3, 1, "home", "装饰精美的小屋"),
        WorldLocation("张伟的家", 5, 1, "home", "宽敞的两室公寓"),
        WorldLocation("王芳的家", 7, 1, "home", "带花园的房子"),
        WorldLocation("星光咖啡店", 4, 4, "shop", "社区里最受欢迎的咖啡店"),
        WorldLocation("社区公园", 5, 6, "park", "有长椅和喷泉的小公园"),
        WorldLocation("创新工作室", 2, 5, "work", "开放式联合办公空间"),
        WorldLocation("社区图书馆", 7, 5, "public", "安静的阅读空间"),
        WorldLocation("小镇广场", 4, 8, "public", "居民聚会的中心广场"),
    ]
    for loc in locations:
        world.add_location(loc)

    return world

class TimeManager:
    """时间管理器：控制模拟的时间推进"""

    def __init__(self, start_time: datetime):
        self.current_time = start_time
        self.time_step = timedelta(minutes=30)  # 每次推进30分钟

    def advance(self) -> datetime:
        self.current_time += self.time_step
        return self.current_time

    @property
    def hour(self) -> int:
        return self.current_time.hour

    @property
    def is_daytime(self) -> bool:
        return 7 <= self.hour <= 22

    @property
    def period(self) -> str:
        h = self.hour
        if h < 7:
            return "凌晨"
        elif h < 9:
            return "早晨"
        elif h < 12:
            return "上午"
        elif h < 14:
            return "中午"
        elif h < 17:
            return "下午"
        elif h < 19:
            return "傍晚"
        elif h < 22:
            return "晚上"
        else:
            return "深夜"

"""
Mini赛博小镇 — Generative Agents 完整实现
==========================================
创建4个具有不同性格的Agent，模拟一天的生活。

依赖：pip install openai numpy
环境变量：OPENAI_API_KEY
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from datetime import datetime, timedelta

import numpy as np
from openai import AsyncOpenAI

# ─────────────── 数据模型 ───────────────

@dataclass
class Memory:
    id: int
    timestamp: datetime
    content: str
    mem_type: str = "observation"   # observation / reflection / dialogue / plan
    importance: int = 5
    embedding: list[float] | None = None
    related_agents: list[str] = field(default_factory=list)

@dataclass
class Agent:
    name: str
    age: int
    occupation: str
    personality: str
    background: str
    relationships: dict[str, str] = field(default_factory=dict)
    daily_routine: str = ""
    schedule: list[tuple[str, str]] = field(default_factory=list)
    current_action: str = "睡觉"
    current_location: str = ""

@dataclass
class Place:
    name: str
    x: int
    y: int
    description: str
    occupants: set[str] = field(default_factory=set)

# ─────────────── 记忆流 ───────────────

class SimpleMemoryStream:
    """简化版记忆流"""

    def __init__(self, agent_name: str, llm: AsyncOpenAI):
        self.agent_name = agent_name
        self.llm = llm
        self.memories: list[Memory] = []
        self._next_id = 1
        self._imp_accum = 0.0

    async def add(
        self, content: str, timestamp: datetime,
        mem_type: str = "observation",
        related_agents: list[str] | None = None,
    ) -> Memory:
        importance = await self._rate(content)
        embedding = await self._embed(content)
        mem = Memory(
            id=self._next_id, timestamp=timestamp,
            content=content, mem_type=mem_type,
            importance=importance, embedding=embedding,
            related_agents=related_agents or [],
        )
        self.memories.append(mem)
        self._next_id += 1
        self._imp_accum += importance
        return mem

    async def retrieve(
        self, query: str, now: datetime, top_k: int = 8,
    ) -> list[Memory]:
        if not self.memories:
            return []
        q_emb = await self._embed(query)
        scored = []
        for m in self.memories:
            hours = max((now - m.timestamp).total_seconds() / 3600, 0.01)
            recency = 0.99 ** hours
            importance = m.importance / 10.0
            relevance = self._cos_sim(q_emb, m.embedding)
            score = recency + importance + relevance
            scored.append((score, m))
        scored.sort(key=lambda x: x[0], reverse=True)  # 按近因性+重要性+相关性综合得分降序
        return [m for _, m in scored[:top_k]]

    def should_reflect(self) -> bool:
        return self._imp_accum >= 40

    def reset_accum(self):
        self._imp_accum = 0.0

    async def _rate(self, content: str) -> int:
        resp = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    "对以下事件的重要性评分(1-10)。1=极日常,10=改变人生。只输出数字。"
                )},
                {"role": "user", "content": content},
            ],
            temperature=0.1, max_tokens=5,
        )
        try:
            return max(1, min(10, int(resp.choices[0].message.content.strip())))
        except ValueError:
            return 5

    async def _embed(self, text: str) -> list[float]:
        resp = await self.llm.embeddings.create(
            model="text-embedding-3-small", input=text,
        )
        return resp.data[0].embedding

    @staticmethod
    def _cos_sim(a: list[float] | None, b: list[float] | None) -> float:
        if a is None or b is None:
            return 0.0
        a_, b_ = np.array(a), np.array(b)
        d = np.linalg.norm(a_) * np.linalg.norm(b_)
        return float(np.dot(a_, b_) / d) if d > 0 else 0.0  # 余弦相似度，分母为零时返回0避免除零错误

# ─────────────── 世界环境 ───────────────

class World:
    def __init__(self):
        self.places: dict[str, Place] = {}
        self.agent_loc: dict[str, str] = {}
        self.time = datetime(2025, 6, 15, 7, 0)

    def add_place(self, p: Place):
        self.places[p.name] = p

    def move(self, agent: str, place: str):
        old = self.agent_loc.get(agent)
        if old and old in self.places:
            self.places[old].occupants.discard(agent)
        if place in self.places:
            self.places[place].occupants.add(agent)
        self.agent_loc[agent] = place

    def colocated(self, agent: str) -> set[str]:
        loc = self.agent_loc.get(agent, "")
        if loc in self.places:
            return self.places[loc].occupants - {agent}
        return set()

    def advance(self, minutes: int = 60):
        self.time += timedelta(minutes=minutes)

    def state_str(self) -> str:
        lines = [f"⏰ {self.time.strftime('%Y-%m-%d %H:%M')}"]
        for name, p in self.places.items():
            if p.occupants:
                lines.append(f"  📍 {name}: {', '.join(sorted(p.occupants))}")
        return "\n".join(lines)

# ─────────────── 模拟引擎 ───────────────

class Simulation:
    """Mini赛博小镇模拟引擎"""

    def __init__(self):
        self.llm = AsyncOpenAI()
        self.world = World()
        self.agents: dict[str, Agent] = {}
        self.memories: dict[str, SimpleMemoryStream] = {}
        self.log: list[str] = []

    # ── 初始化 ──

    def setup(self):
        """初始化世界和Agent"""
        # 创建地点
        places = [
            Place("陈思远的家", 1, 1, "温馨的公寓"),
            Place("林悦的家", 3, 1, "精致的小屋"),
            Place("张伟的家", 5, 1, "宽敞的公寓"),
            Place("王芳的家", 7, 1, "带花园的房子"),
            Place("星光咖啡店", 4, 4, "社区最受欢迎的咖啡店"),
            Place("社区公园", 5, 6, "有长椅和喷泉的公园"),
            Place("创新工作室", 2, 5, "联合办公空间"),
            Place("社区图书馆", 7, 5, "安静的图书馆"),
        ]
        for p in places:
            self.world.add_place(p)

        # 创建Agent
        agents_data = [
            Agent(
                name="陈思远", age=28, occupation="软件工程师",
                personality="内向但热心，喜欢技术，会主动帮助别人",
                background="从小城市来到大城市工作，热爱编程，最近在学习AI",
                relationships={"林悦": "好朋友，经常一起喝咖啡", "张伟": "同事",
                               "王芳": "邻居，偶尔打招呼"},
                daily_routine="早起，去咖啡店工作，下午去工作室，晚上看书",
            ),
            Agent(
                name="林悦", age=26, occupation="平面设计师",
                personality="外向活泼，善于社交，喜欢组织活动",
                background="艺术学院毕业，做自由职业设计师，是社区活动的积极组织者",
                relationships={"陈思远": "好朋友", "张伟": "认识但不太熟",
                               "王芳": "关系不错，经常聊天"},
                daily_routine="晚起，去咖啡店设计，下午逛公园找灵感，喜欢社交",
            ),
            Agent(
                name="张伟", age=35, occupation="大学教授",
                personality="严谨学术派，但对学生很耐心，偶尔幽默",
                background="物理学博士，在本地大学任教，研究量子计算",
                relationships={"陈思远": "同事关系", "林悦": "不太熟",
                               "王芳": "经常去她的花园买花"},
                daily_routine="早起锻炼，去图书馆备课研究，下午给学生答疑",
            ),
            Agent(
                name="王芳", age=42, occupation="花艺师/社区志愿者",
                personality="温暖亲切，是社区的核心人物，喜欢帮助他人",
                background="经营一个小花店，同时是社区志愿者委员会主席",
                relationships={"陈思远": "邻居，对这个年轻人很照顾",
                               "林悦": "好朋友，经常合作搞社区活动",
                               "张伟": "老客户，卖花给他装饰办公室"},
                daily_routine="早起打理花园，上午在花店，下午做社区志愿工作",
            ),
        ]

        for a in agents_data:
            self.agents[a.name] = a
            self.memories[a.name] = SimpleMemoryStream(a.name, self.llm)
            # 所有Agent从自己家开始
            self.world.move(a.name, f"{a.name}的家")

    # ── 规划 ──

    async def plan_day(self, agent: Agent) -> list[tuple[str, str]]:
        """为Agent生成一天的日程"""
        recent = await self.memories[agent.name].retrieve(
            "今天的安排", self.world.time, top_k=5,
        )
        mem_ctx = "\n".join(f"- {m.content}" for m in recent) or "暂无"

        resp = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"你是{agent.name}，{agent.age}岁{agent.occupation}。\n"
                    f"性格: {agent.personality}\n"
                    f"日常: {agent.daily_routine}\n"
                    "为自己规划从07:00到22:00的日程(每小时一项)。\n"
                    f"可选地点: {', '.join(self.world.places.keys())}\n"
                    "输出JSON: {{\"schedule\": [[\"07:00\", \"活动\", \"地点\"], ...]}}"
                )},
                {"role": "user", "content": f"记忆:\n{mem_ctx}"},
            ],
            response_format={"type": "json_object"},
            temperature=0.8,
        )
        data = json.loads(resp.choices[0].message.content)
        schedule = []
        for item in data.get("schedule", []):
            time_str = item[0] if len(item) > 0 else "08:00"
            action = item[1] if len(item) > 1 else "休息"
            location = item[2] if len(item) > 2 else f"{agent.name}的家"
            schedule.append((time_str, f"{action} @ {location}"))
        agent.schedule = schedule
        return schedule

    # ── 执行行动 ──

    async def execute_action(
        self, agent: Agent, action: str, time_slot: str,
    ):
        """执行一个Agent的行动"""
        # 解析位置
        location = f"{agent.name}的家"
        if " @ " in action:
            parts = action.rsplit(" @ ", 1)
            action_desc = parts[0]
            location = parts[1]
        else:
            action_desc = action

        # 确保位置存在
        if location not in self.world.places:
            # 找最近匹配
            for place_name in self.world.places:
                if location in place_name or place_name in location:
                    location = place_name
                    break
            else:
                location = f"{agent.name}的家"

        # 移动Agent
        self.world.move(agent.name, location)
        agent.current_action = action_desc
        agent.current_location = location

        # 记录观察
        observation = f"{time_slot} - {action_desc}（在{location}）"
        await self.memories[agent.name].add(
            observation, self.world.time, "observation",
        )

        self._print(f"  {agent.name}: {action_desc} 📍{location}")

    # ── 社交互动 ──

    async def check_interactions(self):
        """检查同一位置的Agent是否互动"""
        for name, agent in self.agents.items():
            nearby = self.world.colocated(name)
            for other_name in nearby:
                # 避免重复（只让"名字排序靠前"的一方发起）
                if name > other_name:
                    continue
                other = self.agents[other_name]

                # 决定是否对话
                should_talk = await self._should_talk(agent, other)
                if should_talk:
                    await self._do_conversation(agent, other)

    async def _should_talk(self, a: Agent, b: Agent) -> bool:
        rel = a.relationships.get(b.name, "不太熟")
        resp = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"{a.name}(性格:{a.personality})在{a.current_location}"
                    f"遇到了{b.name}。关系: {rel}。"
                    "会主动聊天吗？只回答yes或no"
                )},
                {"role": "user", "content": "yes or no?"},
            ],
            temperature=0.5, max_tokens=5,
        )
        return "yes" in resp.choices[0].message.content.lower()

    async def _do_conversation(self, a: Agent, b: Agent):
        """执行一次对话"""
        mem_a = await self.memories[a.name].retrieve(
            f"与{b.name}的关系", self.world.time, top_k=3,
        )
        mem_b = await self.memories[b.name].retrieve(
            f"与{a.name}的关系", self.world.time, top_k=3,
        )

        resp = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"模拟{a.name}和{b.name}在{a.current_location}的简短对话(3-4轮)。\n"
                    f"{a.name}: {a.personality}。当前在做: {a.current_action}\n"
                    f"  记忆: {'; '.join(m.content for m in mem_a[:2])}\n"
                    f"{b.name}: {b.personality}。当前在做: {b.current_action}\n"
                    f"  记忆: {'; '.join(m.content for m in mem_b[:2])}\n"
                    "格式: 姓名: 说话内容"
                )},
                {"role": "user", "content": "请生成对话。"},
            ],
            temperature=0.8, max_tokens=300,
        )
        dialogue = resp.choices[0].message.content.strip()

        self._print(f"\n  💬 对话 ({a.current_location}):")
        for line in dialogue.split("\n"):
            if line.strip():
                self._print(f"    {line.strip()}")
        self._print("")

        # 存入双方记忆
        summary_a = f"在{a.current_location}和{b.name}聊了天"
        summary_b = f"在{b.current_location}和{a.name}聊了天"
        await self.memories[a.name].add(
            summary_a, self.world.time, "dialogue", [b.name],
        )
        await self.memories[b.name].add(
            summary_b, self.world.time, "dialogue", [a.name],
        )

    # ── 反思 ──

    async def maybe_reflect(self, agent: Agent):
        """如果累积重要性超过阈值，进行反思"""
        ms = self.memories[agent.name]
        if not ms.should_reflect():
            return

        recent = [m.content for m in ms.memories[-15:]]
        resp = await self.llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": (
                    f"你是{agent.name}。基于最近的经历，写出1-2条深层反思/感悟。"
                    "每条反思独占一行，简洁有力（20字以内）。"
                )},
                {"role": "user", "content": "\n".join(f"- {m}" for m in recent)},
            ],
            temperature=0.7, max_tokens=100,
        )

        reflections = resp.choices[0].message.content.strip().split("\n")
        for ref in reflections:
            ref = ref.strip().lstrip("- ·•")
            if ref:
                await ms.add(ref, self.world.time, "reflection")
                self._print(f"  💭 {agent.name}反思: {ref}")
        ms.reset_accum()

    # ── 主循环 ──

    async def run(self, hours: int = 16):
        """运行模拟"""
        self._print("=" * 60)
        self._print("🏘️ Mini赛博小镇 — 模拟开始")
        self._print("=" * 60)

        # 为所有Agent规划日程
        self._print("\n📋 日程规划:")
        plan_tasks = [self.plan_day(a) for a in self.agents.values()]
        # asyncio.gather并发执行所有协程任务，*解包列表为位置参数，等待全部完成后返回结果列表
        await asyncio.gather(*plan_tasks)
        for a in self.agents.values():
            self._print(f"\n  {a.name}的日程:")
            for t, act in a.schedule[:5]:
                self._print(f"    {t} {act}")
            if len(a.schedule) > 5:
                self._print(f"    ... (共{len(a.schedule)}项)")

        # 按时间步推进模拟
        for step in range(hours):
            current_hour = 7 + step
            if current_hour > 22:
                break
            time_str = f"{current_hour:02d}:00"
            self.world.time = datetime(2025, 6, 15, current_hour, 0)

            self._print(f"\n{'─' * 50}")
            self._print(f"⏰ {time_str}")
            self._print(f"{'─' * 50}")

            # 每个Agent执行当前时间段的行动
            for agent in self.agents.values():
                # 找到对应时间段的行动
                action = "休息"
                for t, act in agent.schedule:
                    if t.startswith(f"{current_hour:02d}"):
                        action = act
                        break
                await self.execute_action(agent, action, time_str)

            # 检查互动
            await self.check_interactions()

            # 检查是否需要反思
            for agent in self.agents.values():
                await self.maybe_reflect(agent)

            # 打印世界状态
            self._print(f"\n{self.world.state_str()}")

        self._print(f"\n{'=' * 60}")
        self._print("🌙 模拟结束 — 一天结束了")
        self._print(f"{'=' * 60}")

        # 打印记忆统计
        self._print("\n📊 记忆统计:")
        for name, ms in self.memories.items():
            types = {}
            for m in ms.memories:
                types[m.mem_type] = types.get(m.mem_type, 0) + 1
            self._print(f"  {name}: {len(ms.memories)}条记忆 {types}")

    def _print(self, msg: str):
        print(msg)
        self.log.append(msg)

    def save_log(self, path: str = "simulation_log.txt"):
        with open(path, "w", encoding="utf-8") as f:
            f.write("\n".join(self.log))
        print(f"\n📄 日志已保存至 {path}")

# ─────────────── 运行入口 ───────────────

async def main():
    sim = Simulation()
    sim.setup()
    await sim.run(hours=16)  # 模拟16小时（07:00 - 22:00）
    sim.save_log()

if __name__ == "__main__":
    asyncio.run(main())