"""
=============================================================================
第 11 讲 · 工具模块：迷你 DataSTORM —— 规划者 / 执行者、统计摘要、查询一致性、论点、报告；基线；评判
=============================================================================

对应讲义：lectures/lecture11_datastorm.md 第 3~5 节

-----------------------------------------------------------------------------
【DataSTORM 的流程】（论文图 1）
-----------------------------------------------------------------------------
    1. 热启动：先在网上做一轮研究（Co-STORM，第 2 讲），得到一份初步的报告和几条洞见；
    2. 多层探索（最多 m 层）：每一层，
         规划者（Planner）提出若干个自然语言的探索问题 →
         执行者（Executor）把每个问题翻译成 SQL、执行，返回答案 + 每一列的统计摘要（自下而上的归纳）→
         查询一致性模块：对比同一层的 SQL（以及洞见库里已有的），把同一个概念的写法统一，必要时重查 →
         更新"全局洞见库"（容量有限，弱的被强的替换）；
       每 p 层生成（或修正）一次**论点**（thesis），引导之后的探索；
    3. 写报告：大纲 → 分节起草 → 逐句核查引用 → 修改 → 润色。

-----------------------------------------------------------------------------
【这个模块里有什么？】
-----------------------------------------------------------------------------
    Question                探索问题的几种类型（前后对比、月度序列、年度趋势、按省份、季节、行为者）
    Executor                把问题翻译成 SQL：有几处"可以有不同写法"的地方（行为者看哪几列、时间窗口），
                            还会偶尔写错；返回结果、统计摘要
    summarize               每一列的统计摘要（不同值的比例、出现最多的 5 个值、最小 / 最大 / 中位数 / 均值）
    DataSTORM               规划者 + 执行者 + 一致性模块 + 洞见库 + 论点 + 报告；可以关掉任一部件做消融
    web_only / react_agent  两个基线：以网上的研究为主的深度研究智能体；不带论点、统计摘要、一致性的 ReAct 智能体
    judge                   程序化的评判：报告里的洞见对上了几条标准洞见、数据库支撑的比例、查询一致性

-----------------------------------------------------------------------------
【模拟的边界】（非常重要）
-----------------------------------------------------------------------------
没有真实的 LLM。"规划者"按设定的规则从候选问题里挑（有论点时更集中在论点相关的问题上；看到统计摘要里
的异常值会追问）；"执行者"按设定的概率选择不同的写法、写错 SQL；"论点"取洞见库里最强的那条。
SQL 是真实执行的，洞见的数字来自数据库。数字说明的是**机制**，不是真实系统的成绩。
"""

import random
import statistics
from dataclasses import dataclass, field

from conflict_world import COUNTRIES, EVENT_TYPES, POLICY, UNION

PARAMS = {
    "layers": 5, "first_n": 2, "per_layer": 5,   # 论文在 InsightBench 上的设置：m = 5，第一层 2 个问题，之后每层 5 个
    "thesis_every": 2,                           # 每 p 层生成 / 修正一次论点
    "bank_size": 8,                              # 洞见库的容量
    "p_focus_thesis": 0.85,                      # 有论点时，挑和论点相关的问题的概率
    "p_focus_topic": 0.5,                        # 没有论点时，挑和研究主题（国家）相关的问题的概率
    "p_relevant": 0.3,                           # 每次挑问题时，直接挑"研究问题本身指向的分析类型"的概率
    "p_sql_error": 0.06,                         # 执行者把 SQL 写错（例如漏了一个条件）的概率
    "p_factcheck": 0.8,                          # 写报告时，逐句核查能发现"数字和查询结果不符"的概率
}
# 研究主题（仿论文的 ACLED 数据集：每个主题是一个中性的研究问题）。kinds：读了研究问题，自然会先想到的分析类型
TOPICS = [
    {"id": "T1", "country": "Veloria", "text": "Analyze Veloria's total peace policy and its impact on violence.",
     "gold": ["G1", "G2", "G6"], "kinds": ["prepost"]},
    {"id": "T2", "country": "Kestavia", "text": "What drove unrest in Kestavia in 2023?", "gold": ["G3", "G5"],
     "kinds": ["monthly", "actor"]},
    {"id": "T3", "country": "Orundi", "text": "Where is violence concentrated in Orundi?", "gold": ["G4"],
     "kinds": ["province"]},
    {"id": "T4", "country": "Pallasca", "text": "How does the conflict in Pallasca vary over the year?", "gold": ["G7"],
     "kinds": ["season", "monthly"]},
]
VIOLENCE = ["Battles", "Violence against civilians", "Explosions/Remote violence"]


# =============================================================================
# 第 1 部分：探索问题、SQL、统计摘要
# =============================================================================
@dataclass(frozen=True)
class Question:
    """
    一个探索问题（规划者用自然语言提出；这里用结构化的形式表示它的"意思"）。

    【kind】prepost（政策前后对比）/ monthly（月度序列）/ trend（年度趋势）/ province（按省份的死亡分布）/
            season（季节分布）/ actor（某个行为者参与的事件）/ actor_share（某个行为者参与了多大比例的抗议）/
            spike（某个月的细节，追问用）
    """
    kind: str
    country: str
    target: str = None          # 事件类型或行为者
    date: str = None            # 前后对比的分界月份；spike 的月份

    def text(self):
        """问题的自然语言形式。"""
        t = self.target
        return {"prepost": f"How did monthly {t} in {self.country} change before vs after {self.date}?",
                "monthly": f"What are the monthly counts of {t} in {self.country}?",
                "trend": f"How did yearly {t} in {self.country} evolve from 2020 to 2024?",
                "province": f"How are fatalities in {self.country} distributed across provinces?",
                "season": f"What share of {t} in {self.country} happen in December to March?",
                "actor": f"How many events in {self.country} involved {t}?",
                "actor_share": f"What share of Protests in {self.country} involved {t}?",
                "spike": f"What happened with {t} in {self.country} in {self.date}?"}[self.kind]


def summarize(rows, columns):
    """
    每一列的统计摘要（论文第 3.2.1 节："不同值的比例、出现最多的 5 个值；数值列再加最小、最大、中位数、均值"）。

    【为什么】结果表常常太长，只能截断地显示给规划者（前 5 行和后 5 行）。统计摘要能**自下而上**地把截断的
        部分里的异常值暴露出来 —— 例如"最大值远大于均值和中位数"，提示规划者去追问那个月。
    【返回】dict：列名 → 摘要
    """
    out = {}
    for j, col in enumerate(columns):
        vals = [r[j] for r in rows]
        s = {"distinct": len(set(vals)) / max(1, len(vals))}
        counts = {}
        for v in vals:
            counts[v] = counts.get(v, 0) + 1
        s["top"] = sorted(counts.items(), key=lambda kv: -kv[1])[:5]
        nums = [v for v in vals if isinstance(v, (int, float))]
        if nums:
            s.update(min=min(nums), max=max(nums), median=statistics.median(nums), mean=sum(nums) / len(nums))
            s["argmax"] = rows[nums.index(max(nums))][0] if len(columns) > 1 else None
        out[col] = s
    return out


@dataclass
class Result:
    """
    执行者对一个问题的回答。

    【字段】question、sql、rows、columns、stats（统计摘要，可能为 None）、
        choices（这条 SQL 在几个"可以有不同写法"的地方选了什么：{"actor": "a1"/"a1a2"/"all", "window": ...}）、
        error（SQL 是否写错了）
    """
    question: Question
    sql: str
    rows: list
    columns: list
    stats: dict
    choices: dict
    error: bool


class Executor:
    """
    执行者：把一个探索问题翻译成 SQL、执行。

    【论文】执行者是一个 ReAct 式的 Text-to-SQL 智能体（第 8 讲的思路），在一个**独立的上下文**里工作，
        只把自然语言的答案（和 SQL）交回给规划者。
    【模拟】
        · 行为者的问题：看哪几列？只看 actor1 / actor1 和 actor2 / 也看 assoc_actor_1 —— 随机选一种，除非一致性模块
          已经规定了统一的写法（论文的例子：ACLED 的 actor1、assoc_actor_1、actor2、assoc_actor_2 几列，不同的执行者
          处理得不一样，聚合出来的数字完全不同）；
        · 前后对比：前一段从哪个月开始（2020-01 还是政策前两年）—— 同样随机；
        · 以 p_sql_error 的概率写错（漏了国家的条件）。
    """

    def __init__(self, world, rng, p_error):
        self.world, self.rng, self.p_error = world, rng, p_error

    def run(self, q, standard=None, with_stats=True):
        """
        执行一个问题。

        【参数】standard：一致性模块规定的统一写法（{"actor": "all", "window": "2020-01"}），None 表示自由选择
        【返回】Result
        """
        rng = self.rng
        choices = {}
        if q.kind in ("actor", "actor_share"):
            choices["actor"] = (standard or {}).get("actor") or rng.choice(["a1", "a1a2", "all"])
        if q.kind == "prepost":
            choices["window"] = (standard or {}).get("window") or rng.choice(["2020-01", _shift(q.date, -24)])
        error = rng.random() < self.p_error
        where = "1 = 1" if error else f"country = '{q.country}'"           # 写错：漏了国家的条件
        t = q.target
        if q.kind == "prepost":
            sql = (f"SELECT CASE WHEN month < '{q.date}' THEN 'before' ELSE 'after' END AS period, "
                   f"COUNT(*) * 1.0 / COUNT(DISTINCT month) AS per_month FROM events WHERE {where} "
                   f"AND event_type = '{t}' AND month >= '{choices['window']}' GROUP BY period ORDER BY period DESC")
        elif q.kind == "monthly":
            sql = (f"SELECT month, COUNT(*) AS n FROM events WHERE {where} AND event_type = '{t}' "
                   f"GROUP BY month ORDER BY month")
        elif q.kind == "trend":
            sql = (f"SELECT year, COUNT(*) AS n FROM events WHERE {where} AND event_type = '{t}' "
                   f"GROUP BY year ORDER BY year")
        elif q.kind == "province":
            sql = (f"SELECT admin1, SUM(fatalities) AS deaths FROM events WHERE {where} "
                   f"GROUP BY admin1 ORDER BY deaths DESC")
        elif q.kind == "season":
            sql = (f"SELECT SUM(CASE WHEN substr(month, 6, 2) IN ('12', '01', '02', '03') THEN 1 ELSE 0 END) * 1.0 "
                   f"/ COUNT(*) AS dry_share FROM events WHERE {where} AND event_type = '{t}'")
        elif q.kind in ("actor", "actor_share"):
            cols = {"a1": ["actor1"], "a1a2": ["actor1", "actor2"], "all": ["actor1", "assoc_actor_1", "actor2"]}[choices["actor"]]
            cond = " OR ".join(f"{c} = \"{t}\"" for c in cols)
            if q.kind == "actor":
                sql = (f"SELECT event_type, COUNT(*) AS n FROM events WHERE {where} AND ({cond}) "
                       f"GROUP BY event_type ORDER BY n DESC")
            else:
                sql = (f"SELECT SUM(CASE WHEN {cond} THEN 1 ELSE 0 END) * 1.0 / COUNT(*) AS share FROM events "
                       f"WHERE {where} AND event_type = 'Protests'")
        elif q.kind == "spike":
            sql = (f"SELECT admin1, COUNT(*) AS n FROM events WHERE {where} AND event_type = '{t}' "
                   f"AND month = '{q.date}' GROUP BY admin1 ORDER BY n DESC")
        else:
            raise ValueError(q.kind)
        cur = self.world.db.execute(sql)
        self.world.calls += 1
        rows = cur.fetchall()
        columns = [d[0] for d in cur.description]
        stats = summarize(rows, columns) if with_stats else None
        return Result(q, sql, rows, columns, stats, choices, error)


def _join(main, extra):
    """
    把附加条件接到论点的主论断后面："……；but ……"（方向相反：附加的是一个反例或限制）或"……；and ……"（方向相同）。

    【边界】value 不是数字的论断（省份、月份）一律用 and。
    """
    nums = isinstance(main.value, (int, float)) and isinstance(extra.value, (int, float))
    word = "but" if nums and (main.value > 0) != (extra.value > 0) else "and"
    tail = extra.text
    if tail.split()[0] in ("Monthly", "Yearly"):                 # 句首的普通词改成小写；专有名词（国家、组织）不改
        tail = tail[0].lower() + tail[1:]
    return f"{main.text[:-1]}; {word} {tail}"


def _shift(month, k):
    """把 'YYYY-MM' 往前（k < 0）或往后移 k 个月。"""
    y, m = int(month[:4]), int(month[5:])
    total = y * 12 + (m - 1) + k
    return f"{total // 12}-{total % 12 + 1:02d}"


# =============================================================================
# 第 2 部分：从结果里读出洞见
# =============================================================================
@dataclass
class Claim:
    """
    报告里的一条论断。

    【字段】text、key（结构化的描述，评判用）、value（数值）、source（"db" 或 "web"）、strength（重要性）、
        result（来自哪个查询；网上的论断为 None）、wrong（数字是否其实是错的：SQL 写错了，或网上的说法误导）
    """
    text: str
    key: tuple
    value: object
    source: str
    strength: float
    result: Result = None
    wrong: bool = False


def read_insight(res, show_rows=5):
    """
    规划者读一个查询的结果，看能不能得出一条洞见（模拟 LLM 的"读结果"）。

    【思路】按问题的类型，用简单的阈值判断"值得写进报告"的模式：
        · 前后对比：变化超过 20%；              · 年度趋势：增加到 1.8 倍以上或减少一半以上；
        · 按省份：最多的省占一半以上；          · 季节：旱季占一半以上；
        · 行为者：参与的事件数（主要在哪种事件里）；参与的比例（直接报告）；
        · 月度序列：**只有在看得到的时候**才能发现异常的一个月 ——
          要么它在显示的前 5 行、后 5 行里，要么统计摘要的"最大值"远大于均值（≥ 3 倍）。
    【返回】(Claim 或 None, 需要追问的 Question 或 None)
    """
    q, rows = res.question, res.rows
    wrong = res.error
    if not rows:
        return None, None
    if q.kind == "prepost":
        d = {r[0]: r[1] for r in rows}
        if "before" in d and "after" in d and d["before"]:
            change = d["after"] / d["before"] - 1
            if abs(change) >= 0.2:
                return Claim(f"Monthly {q.target} in {q.country} {'fell' if change < 0 else 'rose'} by {abs(change):.0%} "
                             f"after {q.date}.", ("prepost", q.country, q.target), change, "db", abs(change), res, wrong), None
    elif q.kind == "trend":
        first, last = rows[0][1], rows[-1][1]
        if first and (last / first >= 1.8 or last / first <= 0.5):
            return Claim(f"Yearly {q.target} in {q.country} went from {first} ({rows[0][0]}) to {last} ({rows[-1][0]}).",
                         ("trend", q.country, q.target), last / first - 1, "db", abs(last / first - 1) / 2, res, wrong), None
    elif q.kind == "province":
        total = sum(r[1] for r in rows) or 1
        top, share = rows[0][0], rows[0][1] / total
        if share >= 0.5:
            return Claim(f"{share:.0%} of fatalities in {q.country} occurred in {top} province.",
                         ("concentration", q.country, "fatalities"), (top, share), "db", share, res, wrong), None
    elif q.kind == "season":
        share = rows[0][0] or 0
        if share >= 0.5:
            return Claim(f"{share:.0%} of {q.target} in {q.country} happen between December and March.",
                         ("season", q.country, q.target), share, "db", share, res, wrong), None
    elif q.kind == "actor":
        total = sum(r[1] for r in rows)
        main = rows[0][0]
        return Claim(f"{q.target} was involved in {total} events in {q.country}, mostly {main}.",
                     ("actor", q.country, q.target), total, "db", 0.4, res, wrong), None
    elif q.kind == "actor_share":
        share = rows[0][0] or 0
        return Claim(f"{share:.0%} of Protests in {q.country} involved {q.target}.", ("actor_share", q.country, q.target),
                     share, "db", 0.35, res, wrong), None
    elif q.kind == "monthly":
        visible = rows[:show_rows] + rows[-show_rows:]
        st = res.stats
        follow = None
        if st:
            s = st[res.columns[1]]
            if s["max"] >= 3 * s["mean"]:
                follow = Question("spike", q.country, q.target, s["argmax"])     # 统计摘要暴露了异常的一个月
        else:
            mean = sum(r[1] for r in rows) / len(rows)
            peak = max(visible, key=lambda r: r[1])
            if peak[1] >= 3 * mean:
                follow = Question("spike", q.country, q.target, peak[0])         # 碰巧在显示的几行里
        return None, follow
    elif q.kind == "spike":
        total = sum(r[1] for r in rows)
        return Claim(f"{q.target} in {q.country} spiked in {q.date} ({total} events).",
                     ("spike", q.country, q.target), q.date, "db", 0.8, res, wrong), None
    return None, None


# =============================================================================
# 第 3 部分：DataSTORM
# =============================================================================
def candidates(country, hints):
    """
    一个国家可以问的全部探索问题。hints：热启动从网上得到的线索（例如政策开始的月份）。
    """
    out = []
    for et in EVENT_TYPES:
        if hints.get(country, {}).get("date"):
            out.append(Question("prepost", country, et, hints[country]["date"]))
        out += [Question("monthly", country, et), Question("trend", country, et), Question("season", country, et)]
    out.append(Question("province", country))
    if country == "Kestavia":
        out += [Question("actor", country, UNION), Question("actor_share", country, UNION)]
    return out


@dataclass
class Run:
    """一次运行的全部记录：每一层的问题、结果、洞见库、论点的变化、最后的报告。"""
    layers: list = field(default_factory=list)       # [(问题列表, 结果列表)]
    bank: list = field(default_factory=list)
    theses: list = field(default_factory=list)       # [(层, 动作, 论点)]
    thesis_seen: set = field(default_factory=set)    # 生成 / 修正论点时已经考虑过的洞见（判断"有没有新证据"）
    report: list = field(default_factory=list)
    queries: int = 0


class DataSTORM:
    """
    迷你 DataSTORM。

    【用法】run = DataSTORM(world, seed=0).research(topic)；run.report 是报告里的论断（Claim）
    【参数】thesis / inductive / consistency：三个部件，关掉任一个就是论文表 4 的消融
    【规划者的规则】（模拟 LLM 的选择）
        · 第一层：先查"研究问题本身指向的"分析（例如"集中在哪里"→ 按省份的分布），结合热启动的线索
          （例如政策开始的月份）（问题个数 first_n）；之后每次挑问题，也有 p_relevant 的概率直接挑这类分析；
        · 有论点时：以 p_focus_thesis 的概率挑"和论点同一个国家"的问题；其中 60% 的时候挑"同一类分析、换一个
          对象"的问题（检验论点：别的事件类型也是这样吗？有没有反例、有没有附加条件）；
        · 没有论点时：以 p_focus_topic 的概率挑研究主题的国家，否则随便挑一个国家 —— 探索是"散"的；
        · 统计摘要暴露了异常（最大值 ≥ 3 倍均值）时，下一层优先追问那个月。
    """

    def __init__(self, world, seed=0, thesis=True, inductive=True, consistency=True, params=None):
        self.world = world
        self.p = dict(PARAMS)
        if params:
            self.p.update(params)
        self.thesis_on, self.inductive, self.consistency = thesis, inductive, consistency
        self.rng = random.Random(f"datastorm|{seed}")
        self.executor = Executor(world, self.rng, self.p["p_sql_error"])

    # ---- 热启动：网上的研究 ----
    def warm_start(self, topic):
        """
        热启动（论文第 3.1 节，用 Co-STORM）：读和主题相关的网上报道，得到几条线索和几条"来自网上"的论断。

        【返回】(线索 {国家: {"date": ...}}, 网上的论断列表)
        """
        hints, claims = {}, []
        for date, text, meta in self.world.web:
            if meta["country"] != topic["country"]:
                continue
            if meta.get("date"):
                hints.setdefault(meta["country"], {})["date"] = meta["date"]
            claims.append(Claim(text, ("web", meta["country"], date), None, "web", 0.3, None,
                                wrong=bool(meta.get("misleading"))))
        return hints, claims

    # ---- 一层 ----
    def plan(self, topic, hints, layer, thesis, followups, asked):
        """规划者：这一层问哪些问题。"""
        n = self.p["first_n"] if layer == 0 else self.p["per_layer"]
        picks = [q for q in followups if q not in asked][:n]
        pool_topic = [q for q in candidates(topic["country"], hints) if q not in asked]
        relevant = [q for q in pool_topic if q.kind in topic["kinds"]]
        if layer == 0:                                            # 先查研究问题本身指向的、网上提到的事
            first = [q for q in relevant if q.kind != "prepost" or q.target in VIOLENCE[:1]]
            self.rng.shuffle(first)
            picks += first[:n - len(picks)]
        while len(picks) < n:
            fresh = [q for q in relevant if q not in picks]
            if fresh and self.rng.random() < self.p["p_relevant"]:
                pool = fresh                                      # 读了研究问题，直接想到的分析
            elif thesis is not None and self.rng.random() < self.p["p_focus_thesis"]:
                pool = [q for q in candidates(thesis.key[1], hints) if q not in asked and q not in picks]
                same_kind = [q for q in pool if q.kind == thesis.key[0] or (thesis.key[0] == "prepost" and q.kind == "trend")]
                if same_kind and self.rng.random() < 0.6:
                    pool = same_kind                              # 检验论点：同一类分析，换一个对象（有没有反例？）
            elif self.rng.random() < self.p["p_focus_topic"]:
                pool = [q for q in pool_topic if q not in picks]
            else:
                c = self.rng.choice(list(COUNTRIES))
                pool = [q for q in candidates(c, hints) if q not in asked and q not in picks]
            if not pool:
                break
            picks.append(self.rng.choice(pool))
        return picks

    def research(self, topic):
        """完整的一次研究：热启动 → m 层探索（每 p 层修正论点）→ 写报告。【返回】Run"""
        run = Run()
        hints, web_claims = self.warm_start(topic)
        asked, followups, thesis = set(), [], None
        standard = {"actor": "all", "window": "2020-01"} if self.consistency else None
        for layer in range(self.p["layers"]):
            qs = self.plan(topic, hints, layer, thesis, followups, asked)
            followups = []
            results = []
            for q in qs:
                asked.add(q)
                res = self.executor.run(q, standard=standard, with_stats=self.inductive)
                run.queries += 1
                results.append(res)
                claim, follow = read_insight(res)
                if follow:
                    followups.append(follow)
                if claim:
                    self._add(run.bank, claim)
            run.layers.append((qs, results))
            if self.thesis_on and (layer + 1) % self.p["thesis_every"] == 0:
                thesis = self._thesis(run, thesis, layer, topic)
        run.report = self._write(run, web_claims, thesis, topic)
        return run

    def _add(self, bank, claim):
        """更新洞见库：同一个结构化描述只留一条；超过容量时去掉最弱的。"""
        same = [c for c in bank if c.key == claim.key]
        if same:
            if claim.strength > same[0].strength:
                bank.remove(same[0])
            else:
                return
        bank.append(claim)
        bank.sort(key=lambda c: -c.strength)
        del bank[self.p["bank_size"]:]

    def _thesis(self, run, old, layer, topic):
        """
        生成 / 修正论点（论文第 3.2.2 节；附录 E.2：修正时在 sharpen / pivot / confirm 里选一个）。

        【思路】论点的"主论断" = 洞见库里和研究主题（国家）相关的最强的一条。
            · 还没有旧论点 → generate（生成）；
            · 主论断没变：上次以来有新的证据 → sharpen（加上附加条件："……; and/but ……"），没有 → confirm（确认）；
            · 主论断变了：
                - 新的主论断和旧的方向相反（例如旧的是"战斗减少了"，新的是"远程暴力增加了"），
                  或者旧的主论断已经被挤出了洞见库 → pivot（转向：故事的主线换了）；
                - 否则 → sharpen（新的成为主论断，旧的成为附加条件）。
        【边界】洞见库里还没有和研究主题相关的洞见 → 暂时不生成论点（返回 old）
        【返回】新的论点（一条 Claim；run.theses 里记下这一步的动作和论点的全文）
        """
        pool = [c for c in run.bank if c.key[1] == topic["country"]]
        if not pool:
            return old
        best = pool[0]
        new = [c for c in pool if c.key not in run.thesis_seen and c.key != best.key]
        qualifier = None
        if old is None:
            action = "generate"
        elif best.key == old.key:
            action = "sharpen" if new else "confirm"
            qualifier = new[0] if new else None
        else:
            kept = next((c for c in run.bank if c.key == old.key), None)
            opposite = (kept is not None and isinstance(best.value, (int, float)) and isinstance(kept.value, (int, float))
                        and (best.value > 0) != (kept.value > 0))
            action = "pivot" if kept is None or opposite else "sharpen"
            qualifier = kept
        run.thesis_seen.update(c.key for c in pool)
        text = best.text if qualifier is None else _join(best, qualifier)
        run.theses.append((layer + 1, action, text))
        return best

    def _write(self, run, web_claims, thesis, topic):
        """
        写报告（论文图 2：大纲 → 分节起草 → 逐句核查引用 → 修改 → 润色）。

        【思路】报告 = 洞见库里的论断 + 热启动得到的网上论断。
            有论点时，大纲的每一节都要支撑、检验或修正论点，和论点无关的洞见（别的国家）不写进报告；
            没有论点时，洞见库里有什么就写什么 —— 报告是"散"的（论文：没有中心主题的洞见的堆砌）。
            逐句核查：数字和它引用的查询结果不符的论断（SQL 写错了），以 p_factcheck 的概率被发现、删掉；
            和数据库的结果矛盾的网上说法（"暴力全面下降"，而数据显示针对平民的暴力上升）也会被删掉 ——
            前提是报告里有那条数据库的证据。
        """
        claims = list(run.bank)
        if thesis is not None:
            claims = [c for c in claims if c.key[1] == thesis.key[1]]
        kept = []
        for c in claims:
            if c.wrong and self.rng.random() < self.p["p_factcheck"]:
                continue                                          # 核查发现数字和查询结果不符
            kept.append(c)
        rises = any(c.key[0] == "prepost" and c.value and c.value > 0 and c.key[2] in VIOLENCE[1:] and not c.wrong for c in kept)
        for w in web_claims:
            if w.wrong and rises:
                continue                                          # 网上的说法被数据反驳了
            kept.append(w)
        return kept


# =============================================================================
# 第 4 部分：基线
# =============================================================================
def web_only(world, topic, seed=0):
    """
    基线 1：以网上的研究为主的深度研究智能体（类似 OpenAI Deep Research 拿到 CSV 时的表现：
    论文测得它的报告里只有 23.3% 的论断来自数据库）。

    【模拟】热启动的网上论断全部保留（包括误导性的那条）；只做 3 个数据库查询（随机挑，没有统计摘要、
        没有一致性）；不核查。
    """
    ds = DataSTORM(world, seed=seed, thesis=False, inductive=False, consistency=False)
    hints, web_claims = ds.warm_start(topic)
    run = Run()
    pool = candidates(topic["country"], hints)
    ds.rng.shuffle(pool)
    for q in pool[:3]:
        res = ds.executor.run(q, with_stats=False)
        run.queries += 1
        claim, _ = read_insight(res)
        if claim:
            run.bank.append(claim)
    run.report = list(run.bank) + web_claims
    return run


def react_agent(world, topic, seed=0):
    """
    基线 2：一个 ReAct 式的数据分析智能体（类似 InsightBench 上之前最好的 AgentPoirot）：
    同样多的问题，但没有论点、没有统计摘要、没有一致性模块。
    """
    return DataSTORM(world, seed=seed, thesis=False, inductive=False, consistency=False).research(topic)


# =============================================================================
# 第 5 部分：评判
# =============================================================================
def matches(claim, gold, world):
    """
    一条论断是否"对上了"一条标准洞见（程序化的评判；论文用 LLM 当评委）。

    【规则】结构化的描述相同，而且数值大致正确：
        前后对比：方向相同、相差不超过 15 个百分点；趋势：增长超过标准值的 80%；
        按省份：省相同、比例相差不超过 15 个百分点；季节：比例超过一半；突增：月份相同；
        行为者：数出的事件数和"把关联行为者也算上"的真实数相差不超过 10%。
    """
    if claim.source != "db" or claim.wrong or claim.key != gold.key:
        return False
    kind = gold.key[0]
    if kind == "prepost":
        return (claim.value > 0) == (gold.value > 0) and abs(claim.value - gold.value) <= 0.15
    if kind == "trend":
        return claim.value >= 0.8 * gold.value
    if kind == "concentration":
        return claim.value[0] == gold.value[0] and abs(claim.value[1] - gold.value[1]) <= 0.15
    if kind == "season":
        return claim.value >= 0.5
    if kind == "spike":
        return claim.value == gold.value
    if kind == "actor":
        true = world.query(f"SELECT COUNT(*) FROM events WHERE country = '{gold.key[1]}' AND (actor1 = \"{UNION}\" "
                           f"OR assoc_actor_1 = \"{UNION}\" OR actor2 = \"{UNION}\")")[0][0]
        return abs(claim.value - true) <= 0.1 * true
    return False


def consistency_score(run, detail=False):
    """
    查询一致性（论文附录 E.1 的模拟版）：报告用到的 SQL 里，指向同一个概念的几条查询，写法是否一致。

    【思路】论文的指标有六个维度（人群范围、事件类别、行为者、计数方式、地理、时间），由 GPT-5 先"分组"
        （找出指向同一个概念的两条以上的查询），再"判断"每组的写法是否等价。这里只有两个维度可能不一致：
        行为者（同一个行为者的查询，看的列是否相同）、时间（同一个国家的前后对比，前一段的起点是否相同）。
        一个维度的得分 = "有两条以上查询的概念组"里写法一致的比例；没有可比较的组时记 1.0（不可能不一致）。
        总分是各维度的平均。
    【参数】detail：True 时返回 {"actor": 得分, "temporal": 得分, "overall": 总分}
    【返回】总分（0~1），或上面的字典
    """
    results = [r for qs, rs in run.layers for r in rs]
    groups = {"actor": {}, "window": {}}
    for r in results:
        for dim in groups:
            if dim in r.choices:
                concept = (r.question.country, r.question.target if dim == "actor" else "")
                groups[dim].setdefault(concept, []).append(r.choices[dim])
    scores = {}
    for dim, by_concept in groups.items():
        multi = [choices for choices in by_concept.values() if len(choices) >= 2]
        scores["temporal" if dim == "window" else dim] = (sum(len(set(c)) == 1 for c in multi) / len(multi)
                                                          if multi else 1.0)
    overall = sum(scores.values()) / len(scores)
    return dict(scores, overall=overall) if detail else overall


def judge(run, topic, world):
    """
    评判一次研究。

    【返回】dict：
        recall      对上的标准洞见的比例（论文用 LLM 评委比较报告和专家文章；这里按结构化的描述和数值比较）
        db_ratio    报告里来自数据库的论断的比例（论文表 2 的 "DB Use Ratio"）
        focus       报告里和研究主题（国家）相关的论断的比例 —— 没有论点时，报告容易"散"
        wrong       报告里错误的论断条数（SQL 写错了没被发现；误导性的网上说法）
        consistency 查询一致性
    """
    gold = [g for g in world.gold if g.gid in topic["gold"]]
    hit = sum(any(matches(c, g, world) for c in run.report) for g in gold)
    n = max(1, len(run.report))
    return {"recall": hit / len(gold), "db_ratio": sum(c.source == "db" for c in run.report) / n,
            "focus": sum(c.key[1] == topic["country"] for c in run.report) / n,
            "wrong": sum(c.wrong for c in run.report), "consistency": consistency_score(run)}