"""
hai_world.py —— 第 8 讲：一个可编程的"模拟用户 + 智能体"小环境
=====================================================================

【这个文件做什么】
    第 8 讲（Hao Zhu 的客座讲座）讨论：LLM 不只是回答问题，而是和人一起完成任务时，怎样评估它、怎样从交互里学习。
    真人实验又慢又贵，这里造一个"知道真相"的小环境：

        花园小店：店里有 43 种商品（花园世界的名词），每种有 类别 / 颜色 / 大小 / 价位 四个属性。
        用户心里想要某种属性组合（例如"一种红色的、小的、便宜的水果"），一开始只说出一部分（"我想要一种水果"）。
        智能体每一轮可以：
            ask  问一个属性（"你想要什么颜色？"）
            tell 告诉用户它打算怎么做（"我会根据你的回答挑一个"）——不推进任务，只是沟通
            act  推荐一件商品，对话结束
        用户按自己的"性格"回应：会不会答错、能忍受几个问题、每个问题有多烦、在不在乎被告知。

    这个环境里，**任务有没有完成**和**用户满不满意**是两件可以分开测量的事，所以能精确地展示：
    只看任务完成率的评估会漏掉什么、模拟用户和"真实"用户的差别会怎样改变结论、怎样从反馈里学习。

【主要对象】
    CATALOG / prior_over_goals()            商品目录与"用户想要什么"的先验
    UserProfile / sample_users(population)   用户的性格参数；三种人群：理想化模拟、LLM 式模拟、"真实"人群
    Agent 子类                                直接做（AutoAgent）、问到底（AskAllAgent）、
                                              按信息价值提问（VOIAgent，可选"先告知"）、健忘（ForgetfulAgent）、
                                              话痨（ChattyAgent）
    run_episode(agent, user, rng)             跑一段对话，返回 Episode（完整对话记录 + 各项指标）
    initiative_entropy(episodes)              CoGym 的"主动性熵"

【LLM 联系】
    Collaborative Gym（Shao et al. 2025）、SOTOPIA（Zhou et al. 2024）、τ-bench 等真实环境用 LLM 扮演用户和智能体，
    结构和这里一样：环境状态 + 智能体动作 + 用户（真人或模拟）的回应 + 事后的评估。
"""

import math
import os
import random
import sys
from dataclasses import dataclass, field

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from toy_world import NOUN_CAT, NOUN_DIST, NOUNS  # noqa: E402

# =====================================================================
# 1. 商品目录
# =====================================================================
ATTRS = {"color": ["red", "green", "yellow", "white"], "size": ["small", "big"], "price": ["cheap", "dear"]}
QUESTION = {"color": "what color do you want ?", "size": "what size do you want ?",
            "price": "do you want a cheap one or a dear one ?"}


def _make_catalog(seed=3):
    """每个名词随机（但固定）地分到一组属性。【返回】{名词: {"category":…, "color":…, "size":…, "price":…}}。"""
    rng = random.Random(seed)
    cat = {}
    for c, words in NOUNS.items():
        for w in words:
            cat[w] = {"category": c, **{a: rng.choice(vals) for a, vals in ATTRS.items()}}
    return cat


CATALOG = _make_catalog()
CATEGORY_PRIOR = {"fruit": 0.35, "veg": 0.25, "flower": 0.2, "animal": 0.2}


def item_prob(noun):
    """用户想要这件商品的先验概率 = 类别的概率 × 类别内的 Zipf 频率（和花园世界的名词频率一致）。"""
    c = NOUN_CAT[noun]
    return CATEGORY_PRIOR[c] * NOUN_DIST[c][noun]


def goal_of(noun):
    """一件商品的属性组合（用户的"目标"就是一个属性组合：任何属性完全匹配的商品都算对）。"""
    d = CATALOG[noun]
    return (d["category"], d["color"], d["size"], d["price"])


def goal_dict(goal):
    return dict(zip(("category", "color", "size", "price"), goal))


def prior_over_goals():
    """属性组合的先验：把同一组合下所有商品的概率加起来。【返回】{目标: 概率}。"""
    pri = {}
    for n in CATALOG:
        g = goal_of(n)
        pri[g] = pri.get(g, 0.0) + item_prob(n)
    return pri


PRIOR = prior_over_goals()


def sample_goal(rng):
    """按先验抽一个目标。"""
    goals = list(PRIOR)
    return rng.choices(goals, weights=[PRIOR[g] for g in goals])[0]


def item_for(goal):
    """属性组合对应的最常见商品（推荐时说出它的名字）。"""
    cands = [n for n in CATALOG if goal_of(n) == goal]
    return max(cands, key=item_prob)


# =====================================================================
# 2. 用户
# =====================================================================

@dataclass
class UserProfile:
    """
    一个用户的"性格"。

    【字段】
        noise：回答属性问题时答错的概率（记错、没想清楚、随口说）；
        patience：最多能忍受几个问题，超过就离开（任务失败）；None 表示无限耐心；
        q_cost：每个问题带来的不满（满意度扣分）；
        wants_update：在不在乎智能体"先告诉我你打算怎么做"（在乎的人，没被告知会扣分）；
        reveal：开口时除了类别，再主动多说一个属性的概率；
        volunteer：被问到一个属性时，顺便多说一个没被问到的属性的概率（"主动"的用户）。
    """
    noise: float = 0.0
    patience: int = None
    q_cost: float = 0.0
    wants_update: bool = False
    reveal: float = 0.3
    volunteer: float = 0.0


def sample_users(population, n, seed):
    """
    三种人群：
        "ideal"：理想化的模拟用户——从不答错、无限耐心、不嫌问题烦、不在乎被告知；
        "llm"  ：LLM 扮演的用户——几乎不答错、非常配合（常常主动多说）、耐心很好、不太在乎被告知；
        "real" ：一个更像真人的人群——会答错、耐心有限且因人而异、嫌问题烦的程度因人而异、
                 一半的人在乎被告知、很少主动多说。
    【返回】list[UserProfile]。
    【边界】这些参数是编的，不是测出来的；它们的作用是让"模拟用户和真实用户不一样"这件事可以被量化地讨论。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        if population == "ideal":
            out.append(UserProfile())
        elif population == "llm":
            out.append(UserProfile(noise=0.02, patience=None, q_cost=0.02, wants_update=rng.random() < 0.2,
                                   reveal=0.5, volunteer=0.5))
        elif population == "real":
            out.append(UserProfile(noise=0.1, patience=rng.choice([1, 2, 2, 3]),
                                   q_cost=rng.choice([0.05, 0.1, 0.2]), wants_update=rng.random() < 0.5,
                                   reveal=0.3, volunteer=0.1))
        else:
            raise ValueError(population)
    return out


# =====================================================================
# 3. 对话
# =====================================================================

@dataclass
class Episode:
    """
    一段对话的完整记录与结果。

    【字段】
        turns：[(说话人, 动作, 文字, 是否主动发起)]，说话人是 "user" 或 "agent"；
        goal：用户心里的目标；recommended：推荐的商品（没推荐就是 None）；
        success：推荐的商品属性是否完全匹配目标；quit：用户是否因为问题太多而离开；
        n_questions：智能体问了几个问题；informed：智能体推荐前有没有"告知"；
        satisfaction：用户的满意度（见 run_episode）；asked_known：问了用户已经说过的属性的次数；
        confidence：推荐时智能体自认为的成功概率（它的后验）。
    """
    turns: list = field(default_factory=list)
    goal: tuple = None
    recommended: str = None
    success: bool = False
    quit: bool = False
    n_questions: int = 0
    informed: bool = False
    satisfaction: float = 0.0
    asked_known: int = 0
    user: UserProfile = None
    confidence: float = None


def run_episode(agent, user, rng, goal=None, max_turns=8):
    """
    让 agent 和 user 对话，直到推荐了商品、用户离开或者轮数用完。

    【参数】agent：Agent 对象；user：UserProfile；rng：random.Random；goal：可指定目标（默认按先验抽）。
    【返回】Episode。
    【满意度】 satisfaction = 1[成功] − q_cost × 问题数 − 0.3 × 1[在乎被告知但没被告知]
                             − 0.5 × 1[因为问题太多而离开]
            ——"任务完成"只是满意度的一部分。
    """
    goal = goal or sample_goal(rng)
    gd = goal_dict(goal)
    ep = Episode(goal=goal, user=user)
    said = {"category": gd["category"]}
    opening = f"i want a {gd['category']}"
    if rng.random() < user.reveal:
        a = rng.choice(list(ATTRS))
        said[a] = gd[a]
        opening = f"i want a {gd[a]} {gd['category']}" if a != "price" else f"i want a {gd['category']} , a {gd[a]} one"
    ep.turns.append(("user", "request", opening, True))
    agent.reset(dict(said))
    for _ in range(max_turns):
        act, arg = agent.step()
        if act == "ask":
            ep.n_questions += 1
            if arg in said:
                ep.asked_known += 1
            ep.turns.append(("agent", "ask", QUESTION[arg], True))
            if user.patience is not None and ep.n_questions > user.patience:
                ep.turns.append(("user", "quit", "this takes too long , bye", True))
                ep.quit = True
                break
            ans = gd[arg]
            if rng.random() < user.noise:
                ans = rng.choice([v for v in ATTRS[arg] if v != gd[arg]])
            reply = {arg: ans}
            if rng.random() < user.volunteer:
                rest = [a for a in ATTRS if a not in said and a != arg]
                if rest:
                    b = rng.choice(rest)
                    reply[b] = gd[b]
            said.update(reply)
            text = " and ".join(f"{v}" for v in reply.values())
            ep.turns.append(("user", "answer", text, len(reply) > 1))
            agent.observe(reply)
        elif act == "tell":
            ep.informed = True
            ep.turns.append(("agent", "tell", arg, True))
            ep.turns.append(("user", "ack", "ok", False))
        elif act == "act":
            ep.recommended = arg
            ep.success = goal_of(arg) == goal
            ep.confidence = getattr(agent, "conf", None)
            ep.turns.append(("agent", "act", f"here is the {arg} , {' '.join(CATALOG[arg][a] for a in ATTRS)}", True))
            break
    s = float(ep.success) - user.q_cost * ep.n_questions
    if user.wants_update and not ep.informed:
        s -= 0.3
    if ep.quit:
        s -= 0.5
    ep.satisfaction = s
    return ep


def initiative_entropy(episodes):
    """
    主动性熵（Collaborative Gym 的 Initiative Entropy，两方时归一化到 [0, 1]）：
        p_i = 一方"主动发起"的话语数 / 双方主动发起的话语总数，H = −Σ p_i log p_i / log 2。
    【返回】所有对话的平均值。某一方完全没有主动发起时定义为 0（论文里的约定）。
    【直觉】1 = 双方主动性均衡；0 = 一方包办。
    """
    hs = []
    for ep in episodes:
        cnt = {"user": 0, "agent": 0}
        for spk, _, _, init in ep.turns:
            if init:
                cnt[spk] += 1
        tot = sum(cnt.values())
        if min(cnt.values()) == 0:
            hs.append(0.0)
            continue
        hs.append(-sum((c / tot) * math.log(c / tot) for c in cnt.values()) / math.log(2))
    return sum(hs) / len(hs)


def summarize(episodes):
    """一组对话的平均指标。【返回】dict。"""
    n = len(episodes)
    return {"success": sum(e.success for e in episodes) / n,
            "questions": sum(e.n_questions for e in episodes) / n,
            "quit": sum(e.quit for e in episodes) / n,
            "informed": sum(e.informed for e in episodes) / n,
            "satisfaction": sum(e.satisfaction for e in episodes) / n,
            "initiative": initiative_entropy(episodes)}


# =====================================================================
# 4. 智能体
# =====================================================================

class Agent:
    """
    智能体的接口：reset(用户开口说了哪些属性) → 反复 step() 得到 (动作, 参数)，ask 之后收到 observe(回答)。

    【信念】所有智能体都维护"用户已经说过的属性"self.known，并据此算出各个目标的后验：
            与已知属性矛盾的目标概率为 0，其余按先验归一化（假设用户回答没错——这本身就是一个可能出错的假设）。
    """
    name = "agent"

    def reset(self, said):
        self.known = dict(said)
        self.told = False

    def observe(self, reply):
        self.known.update(reply)

    def posterior(self):
        post = {g: p for g, p in PRIOR.items()
                if all(goal_dict(g)[a] == v for a, v in self.known.items())}
        z = sum(post.values())
        if z == 0:                                  # 用户的回答互相矛盾（答错了）：退回只按类别
            post = {g: p for g, p in PRIOR.items() if g[0] == self.known["category"]}
            z = sum(post.values())
        return {g: p / z for g, p in post.items()}

    def best(self):
        """后验概率最大的目标，以及推荐它时的成功概率。"""
        post = self.posterior()
        g = max(post, key=post.get)
        return g, post[g]

    def recommend(self):
        """推荐后验最大的目标对应的商品；把"自认为的成功概率"记在 self.conf 里（第 02 号脚本看它准不准）。"""
        g, self.conf = self.best()
        return "act", item_for(g)


class AutoAgent(Agent):
    """直接做：不问、不告知，按已知信息推荐最可能的商品（"完全自主"的智能体）。"""
    name = "直接做"

    def step(self):
        return self.recommend()


class AskAllAgent(Agent):
    """问到底：把颜色、大小、价位中没听过的全部问一遍，再推荐。"""
    name = "问到底"

    def step(self):
        for a in ATTRS:
            if a not in self.known:
                return "ask", a
        return self.recommend()


class VOIAgent(Agent):
    """
    按信息价值提问：只在"问了之后成功率预期提高得够多"时才问。

    【思路】当前推荐的成功率 = 后验最大的那个目标的概率 p*。问属性 a 之后，对每种可能的回答 v，
            后验会更新，推荐的成功率变成新的 p*(v)；期望值 E_v[p*(v)] 减去当前的 p*，就是"问 a 的信息价值"。
            信息价值最大的那个问题，如果价值 > threshold（智能体心里"一个问题值多少"），就问；否则直接推荐。
    【参数】threshold：提问的门槛；tell：推荐之前要不要先告知用户（CoGym 发现主动沟通对协作很重要）。
    """

    def __init__(self, threshold=0.1, tell=False):
        self.threshold, self.tell = threshold, tell
        self.name = f"按信息价值问（门槛 {threshold}）" + ("+告知" if tell else "")

    def value_of(self, a, post):
        cur = max(post.values())
        by_v = {}
        for g, p in post.items():
            by_v.setdefault(goal_dict(g)[a], {})[g] = p
        exp = sum(max(sub.values()) for sub in by_v.values())    # Σ_v P(v) · max_g P(g|v) = Σ_v max_g P(g, v)
        return exp - cur

    def step(self):
        post = self.posterior()
        gains = {a: self.value_of(a, post) for a in ATTRS if a not in self.known}
        if gains:
            a = max(gains, key=gains.get)
            if gains[a] > self.threshold:
                return "ask", a
        if self.tell and not self.told:
            self.told = True
            return "tell", "i will pick the best match for what you told me"
        return self.recommend()


class ForgetfulAgent(AskAllAgent):
    """
    健忘：和"问到底"一样，但不记得用户开口时说过什么（只记住类别）——用户说了"红色的水果"，它还会问"什么颜色？"。
    对应 CoGym 总结的失败类型"缺乏情境感知：把每个请求当成孤立的任务，不顾对话的上下文"。
    """
    name = "健忘（问到底，但忘了开场白）"

    def reset(self, said):
        super().reset({"category": said["category"]})


class ChattyAgent(AskAllAgent):
    """
    话痨：和"问到底"一样问遍所有属性，还会先告知，每句话都加很多客套话（文字更长、更"礼貌"）。
    用来检验"按文字表面打分的评审"会不会被它骗。
    """
    name = "话痨（问到底 + 客套）"

    def step(self):
        act, arg = super().step()
        if act == "act" and not self.told:
            self.told = True
            return "tell", "thank you so much for your patience ! i really appreciate it , i will now carefully pick the perfect one for you"
        return act, arg


def default_agents():
    """第 8 讲反复比较的四个智能体。"""
    return [AutoAgent(), AskAllAgent(), VOIAgent(0.1), VOIAgent(0.1, tell=True)]


def evaluate(agent, users, seed, n_rep=1):
    """让一个智能体和一批用户各对话 n_rep 次。【返回】Episode 列表。"""
    rng = random.Random(seed)
    return [run_episode(agent, u, rng) for _ in range(n_rep) for u in users]


__all__ = ["ATTRS", "CATALOG", "PRIOR", "UserProfile", "sample_users", "Episode", "run_episode",
           "initiative_entropy", "summarize", "Agent", "AutoAgent", "AskAllAgent", "VOIAgent", "ChattyAgent", "ForgetfulAgent",
           "default_agents", "evaluate", "goal_of", "goal_dict", "item_for", "sample_goal", "item_prob"]