"""Agent 与环境交互：可重放的打印机环境、观察基线与资源争用。"""

from argparse import ArgumentParser
from copy import deepcopy
from hashlib import sha256
from itertools import permutations
import json
from pathlib import Path

DATA = Path(__file__).with_name("data")
PLAN = DATA / "agent-environment-plan.json"
RECORD = DATA / "agent-environment-model.json"
CONTROL = DATA / "agent-environment-control.json"
FIELDS = ["enabled", "queue_clear", "power", "paper", "cable"]
ACTIONS = {
    "A": "inspect_server: read enabled and queue_clear",
    "B": "ask_device_status: read power, paper and cable from the simulated user",
    "C": "enable_server: set enabled=true",
    "D": "clear_queue: set queue_clear=true",
    "E": "ask_power_on: request the user to set power=true",
    "F": "ask_load_paper: request the user to set paper=true",
    "G": "ask_connect_cable: request the user to set cable=true",
    "H": "print_page: print one page only when all five readiness fields are true",
    "I": "finish: end the task as completed",
    "J": "wait: let one interaction tick pass",
    "K": "handoff: end unresolved after a required user operation was refused",
}
FIXES = {"enabled": "C", "queue_clear": "D", "power": "E", "paper": "F", "cable": "G"}


def digest(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def fingerprint(value):
    return sha256(json.dumps(value, sort_keys=True, ensure_ascii=False,
                             separators=(",", ":")).encode()).hexdigest()


def save_new(path, value):
    with Path(path).open("x") as stream:
        json.dump(value, stream, ensure_ascii=False, indent=2)
        stream.write("\n")


def make_plan():
    specifications = [
        ("ready", [], 0, None),
        ("server", ["enabled", "queue_clear"], 0, None),
        ("paper", ["paper"], 0, None),
        ("device", ["power", "cable"], 0, None),
        ("mixed", ["enabled", "paper"], 2, None),
        ("all-now", FIELDS, 0, None),
        ("all-delayed", FIELDS, 3, None),
        ("paper-unavailable", ["paper"], 0, "paper"),
    ]
    return dict(
        version=1, date="2026-10-20", model="Qwen/Qwen2.5-0.5B-Instruct",
        revision="7ae557604adf67be50417f59c2c2f167def9a775",
        settings=dict(threads=4, seed=2026, dtype="float32", attention="eager",
                      max_steps=16, input_plus_output_cap=4096,
                      selection="argmax next-token logit restricted to A..K; no sampling"),
        feedback=["full", "receipt"],
        cases=[dict(id=name, initial={key: key not in faults for key in FIELDS},
                    user_delay=delay, unavailable=unavailable)
               for name, faults, delay, unavailable in specifications],
        selection="All fixed cases and both feedback modes; no result-based prompt selection",
    )


class Printer:
    def __init__(self, case):
        self.case = deepcopy(case)
        self.state = deepcopy(case["initial"])
        self.clock, self.printed = 0, 0
        self.pending, self.denied = {}, []
        self.terminal = None

    def snapshot(self):
        return dict(clock=self.clock, readiness=deepcopy(self.state),
                    pending=deepcopy(self.pending), denied=self.denied[:],
                    printed=self.printed, terminal=self.terminal)

    def step(self, action):
        if action not in ACTIONS or self.terminal is not None:
            raise ValueError("动作不合法，或任务已结束")
        self.clock += 1
        events = []
        for key, due in list(self.pending.items()):
            if due <= self.clock:
                self.state[key] = True
                del self.pending[key]
                events.append({"user_completed": key})
        status, details = "ok", {}
        if action == "A":
            details = {key: self.state[key] for key in FIELDS[:2]}
        elif action == "B":
            details = {key: self.state[key] for key in FIELDS[2:]}
        elif action in {"C", "D"}:
            key = "enabled" if action == "C" else "queue_clear"
            self.state[key] = True
            details = {key: True}
        elif action in {"E", "F", "G"}:
            key = {"E": "power", "F": "paper", "G": "cable"}[action]
            if self.state[key]:
                status = "already_ready"
            elif self.case["unavailable"] == key:
                status = "refused"
                details = {"required_operation_unavailable": key}
                if key not in self.denied:
                    self.denied.append(key)
            elif key in self.pending:
                status, details = "pending", {"due": self.pending[key]}
            elif self.case["user_delay"] == 0:
                self.state[key] = True
                status, details = "completed", {key: True}
            else:
                self.pending[key] = self.clock+self.case["user_delay"]
                status, details = "accepted", {"due": self.pending[key]}
        elif action == "H":
            missing = [key for key in FIELDS if not self.state[key]]
            if missing:
                status, details = "not_ready", {"missing": missing}
            else:
                self.printed += 1
                status, details = "printed", {"total_pages": self.printed}
        elif action == "I":
            self.terminal = "finished"
            status = "finished"
        elif action == "K":
            self.terminal = "handoff"
            status = "handoff"
        return dict(clock=self.clock, status=status, details=details, events=events)

    def grade(self):
        if self.case["unavailable"]:
            success = (self.terminal == "handoff" and self.printed == 0
                       and self.case["unavailable"] in self.denied)
        else:
            success = self.terminal == "finished" and self.printed == 1
        return dict(success=success, printed=self.printed, terminal=self.terminal,
                    steps=self.clock, unresolved_fields=[k for k in FIELDS if not self.state[k]])


def visible(result, mode):
    if mode == "full":
        return deepcopy(result)
    if mode == "receipt":
        return {key: result[key] for key in ["clock", "status"]}
    raise ValueError(mode)


def messages(history, mode):
    menu = "\n".join(f"{key}: {value}" for key, value in ACTIONS.items())
    instruction = (
        "You operate a simulated printer. Print exactly ONE test page, then finish. "
        "Initial readiness is unknown. Tools are the only way to observe or change it. "
        "Physical changes require the simulated user; accepted means requested, not completed. "
        "User operations may take several interaction ticks. Every action advances one tick. "
        "Inspect, fix the missing conditions, verify feedback, and stop when done. "
        "If a required user operation is refused, hand off instead of declaring success. "
        "Choose exactly one letter A through K for the NEXT action. Do not explain.\n"+menu)
    note = ("Tool feedback includes details and completed user events." if mode == "full" else
            "Tool feedback gives only the receipt status; details and user events are hidden.")
    return [
        {"role": "system", "content": instruction},
        {"role": "user", "content": note+"\nHistory:\n"+
         json.dumps(history, separators=(",", ":"))+"\nNext action letter:"},
    ]


def reference_episode(case, max_steps):
    """手写观察策略；只根据已收到的回执作选择，不直接读取隐藏就绪状态。"""
    env, knowledge, history, trace = Printer(case), {}, [], []
    while not env.terminal and env.clock < max_steps:
        if any(row["observation"]["status"] == "refused" for row in history):
            action = "K"
        elif any(row["observation"]["status"] == "printed" for row in history):
            action = "I"
        elif "enabled" not in knowledge:
            action = "A"
        elif any(not knowledge[key] for key in FIELDS[:2]):
            action = FIXES[next(k for k in FIELDS[:2] if not knowledge[k])]
        elif "power" not in knowledge:
            action = "B"
        elif any(not knowledge[key] for key in FIELDS[2:]):
            waiting = {row["action"] for row in history
                       if row["observation"]["status"] in {"accepted", "pending"}}
            remaining = [k for k in FIELDS[2:] if not knowledge[k] and FIXES[k] not in waiting]
            action = FIXES[remaining[0]] if remaining else "J"
        else:
            action = "H"
        before = env.snapshot()
        result = env.step(action)
        for key, value in result["details"].items():
            if key in FIELDS:
                knowledge[key] = value
        for event in result["events"]:
            knowledge[event["user_completed"]] = True
        history.append({"action": action, "observation": result})
        trace.append(dict(action=action, before=before, result=result, after=env.snapshot()))
    return dict(id=case["id"], trace=trace, **env.grade())


def contention():
    """枚举 A、B 各先读后写的全部交错；提交在本函数中是不可分割的一步。"""
    orders = [list(order) for order in permutations(["A_read", "A_write", "B_read", "B_write"])
              if order.index("A_read") < order.index("A_write")
              and order.index("B_read") < order.index("B_write")]
    rows = []
    for mode in ["stale_check", "revision_check"]:
        for order in orders:
            bookings, revision, seen, events = [], 0, {}, []
            for operation in order:
                actor, operation_type = operation.split("_")
                if operation_type == "read":
                    seen[actor] = dict(free=not bookings, revision=revision)
                    result = deepcopy(seen[actor])
                elif not seen[actor]["free"]:
                    result = "already_busy"
                elif mode == "revision_check" and seen[actor]["revision"] != revision:
                    result = "conflict"
                else:
                    bookings.append(actor)
                    revision += 1
                    result = "reserved"
                events.append(dict(actor=actor, operation=operation_type, result=result,
                                   bookings=bookings[:], revision=revision))
            rows.append(dict(mode=mode, order=order, events=events, bookings=bookings,
                             overbooked=len(bookings) > 1))
    return rows


def control_report(plan):
    return dict(source_sha256=digest(__file__), plan_sha256=fingerprint(plan),
                reference=[reference_episode(case, plan["settings"]["max_steps"])
                           for case in plan["cases"]], contention=contention())


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--controls", type=Path)
    args = parser.parse_args()
    if args.controls:
        save_new(args.controls, control_report(make_plan()))
        print(f"已保存观察策略与全部交错顺序：{args.controls}")
        return
    plan = json.loads(PLAN.read_text())
    control = json.loads(CONTROL.read_text())
    record = json.loads(RECORD.read_text())
    assert record["complete"]
    assert record["plan_sha256"] == control["plan_sha256"] == fingerprint(plan)
    n = len(plan["cases"])
    print(f"手写观察策略：{sum(r['success'] for r in control['reference'])}/{n}")
    for mode in plan["feedback"]:
        rows = [r for r in record["rows"] if r["mode"] == mode]
        print(f"小模型 {mode}：{sum(r['grade']['success'] for r in rows)}/{n}，"
              f"共 {sum(len(r['trace']) for r in rows)} 次动作选择")
    for mode in ["stale_check", "revision_check"]:
        rows = [r for r in control["contention"] if r["mode"] == mode]
        print(f"名额争用 {mode}：{sum(r['overbooked'] for r in rows)}/{len(rows)} 种交错超额")


if __name__ == "__main__":
    main()

"""显式 --run 才加载固定小模型，在真实执行的教学环境中选择动作。"""

from argparse import ArgumentParser
from datetime import datetime, timezone
from importlib.metadata import version
import json
from pathlib import Path
import platform
import time

import agent_environment as lab


def experiment(directory):
    directory.mkdir(parents=True, exist_ok=False)
    plan = lab.make_plan()
    lab.save_new(directory/"plan.json", plan)  # 在加载模型、观察结果前固定。
    lab.save_new(directory/"control.json", lab.control_report(plan))
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer
    settings = plan["settings"]
    torch.set_num_threads(settings["threads"])
    torch.manual_seed(settings["seed"])
    snapshot = Path(snapshot_download(plan["model"], revision=plan["revision"],
                                    allow_patterns=["*.json", "*.safetensors", "merges.txt",
                                                    "README.md", "LICENSE"]))
    tokenizer = AutoTokenizer.from_pretrained(snapshot, trust_remote_code=False)
    letters = list(lab.ACTIONS)
    ids = [tokenizer.encode(letter, add_special_tokens=False) for letter in letters]
    assert all(len(row) == 1 for row in ids)
    token_ids = [row[0] for row in ids]
    assert len(set(token_ids)) == len(letters)
    model = AutoModelForCausalLM.from_pretrained(
        snapshot, dtype=torch.float32, attn_implementation=settings["attention"],
        trust_remote_code=False).eval()
    report = dict(
        complete=False, started_at=datetime.now(timezone.utc).isoformat(),
        model=plan["model"], revision=plan["revision"], settings=settings,
        plan_sha256=lab.fingerprint(plan), action_tokens=dict(zip(letters, token_ids)),
        source_sha256={p.name: lab.digest(p) for p in [Path(__file__), Path(lab.__file__)]},
        model_files_sha256={p.name: lab.digest(p) for p in snapshot.iterdir()
                            if p.is_file() and (p.suffix in {".json", ".safetensors"}
                                               or p.name == "merges.txt")},
        environment=dict(python=platform.python_version(), platform=platform.platform(),
                         packages={name: version(name) for name in
                                   ["torch", "transformers", "huggingface-hub",
                                    "tokenizers", "safetensors"]}),
        rows=[],
    )

    def checkpoint():
        temporary = directory/"model.pending.json"
        temporary.write_text(json.dumps(report, ensure_ascii=False, indent=2)+"\n")
        temporary.replace(directory/"model.json")

    checkpoint()
    for mode in plan["feedback"]:
        for case in plan["cases"]:
            env, history = lab.Printer(case), []
            row = dict(id=case["id"], mode=mode, trace=[], stopped="step_limit")
            report["rows"].append(row)
            for _ in range(settings["max_steps"]):
                conversation = lab.messages(history, mode)
                inputs = tokenizer.apply_chat_template(
                    conversation, tokenize=True, add_generation_prompt=True,
                    return_dict=True, return_tensors="pt")
                length = inputs["input_ids"].shape[1]
                if length+1 > settings["input_plus_output_cap"]:
                    row["stopped"] = "context_budget"
                    break
                started = time.perf_counter()
                with torch.inference_mode():
                    scores = model(**inputs).logits[0, -1]
                selected = scores[token_ids]
                probabilities = selected.softmax(0).tolist()
                action = letters[int(selected.argmax())]
                before = env.snapshot()
                result = env.step(action)
                observation = lab.visible(result, mode)
                trace = dict(
                    messages=conversation, input_ids=inputs["input_ids"][0].tolist(),
                    input_tokens=length, action=action, action_token=token_ids[letters.index(action)],
                    logits=selected.tolist(), conditional_probabilities=probabilities,
                    legal_token_mass=scores.softmax(0)[token_ids].sum().item(),
                    seconds=time.perf_counter()-started, before=before,
                    result=result, observation=observation, after=env.snapshot())
                row["trace"].append(trace)
                history.append(dict(action=action, observation=observation))
                row["grade"] = env.grade()
                checkpoint()
                if env.terminal:
                    row["stopped"] = env.terminal
                    break
            row["grade"] = env.grade()
            checkpoint()
            print(f"{mode}/{case['id']}: success={row['grade']['success']} "
                  f"steps={len(row['trace'])} actions="
                  f"{''.join(step['action'] for step in row['trace'])}", flush=True)
    report["complete"] = True
    report["finished_at"] = datetime.now(timezone.utc).isoformat()
    checkpoint()
    print(f"完成 16 条真实模型交互轨迹：{directory}", flush=True)


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out-dir", type=Path)
    args = parser.parse_args()
    if not args.run:
        print("默认不加载模型。使用 --run --out-dir 新目录，才运行 16 条交互轨迹。")
        return
    if args.out_dir is None:
        parser.error("需要尚不存在的 --out-dir")
    experiment(args.out_dir)


if __name__ == "__main__":
    main()

1. 第 1 步将完成时钟设为 3；第 2 步仍缺纸，打印失败。第 3 步先兑现装纸事件，再执行打印，所以这次成功。最终还需要结束动作。
2. 不足以；开机、纸张和接线仍未知。B 能返回这三项设备状态。询问完成后，真实系统里还要考虑状态是否又发生变化。
3. 目标同时要求恰好一页与明确结束。前者少了结束，后者多打印了一页；“曾经做对过一步”不是最终目标。
4. 分别约为 0.599 与 0.358。算例没有重试、回滚、检查后改正或替代路径，不能直接预测真实系统可靠性。
5. 可以参照 reference_episode，根据回执和完成事件维护已知状态。迟迟没有事件时不能自行把纸张设为就绪；等待仍计入 16 步预算。另试用户拒绝装纸，应转交而非继续等待。
6. 用 Ar、Aw、Br、Bw 表示读写，六种依次为：Ar Aw Br Bw、Ar Br Aw Bw、Ar Br Bw Aw、Br Bw Ar Aw、Br Ar Bw Aw、Br Ar Aw Bw。第一与第四种不会超额，因为后一人读时已经占用。
7. 可以。两人都先读到 0、检查通过，再各自写入。检查与写入必须作为原子操作执行，或有提供同等约束的事务机制；版本号本身不是锁。
8. 全部成功为 0.6，至少一次为 1。后者只说明候选中存在成功尝试，还需要能识别并选择它；参见推理与测试时计算。
9. 不能。模型没有进入打印与延迟用户操作的关键阶段，同一模式下也没有获取足以区分案例的信息。需要更多能使用工具的配置、独立案例及重复实验，再判断反馈的影响。
10. 追加历史改变本次输入，训练改变参数。反复按这八题修改提示，或只用它们更新权重，都可能记住局部规律；应另设开发与留出任务，测试不同故障组合、延迟和用户行为。