Uni-Lab-OS/tests/compile/test_pump_separate_full_chain.py

"""
PumpTransfer 和 Separate 全链路测试

构建包含泵/阀门/分液漏斗的完整设备图，
输出完整的中间数据（最短路径、泵骨架、动作列表等）。
"""

import copy
import json
import pprint
import pytest
import networkx as nx

from unilabos.resources.resource_tracker import ResourceTreeSet
from unilabos.compile.utils.resource_helper import get_resource_id, get_resource_data
from unilabos.compile.utils.vessel_parser import get_vessel


def _make_raw_resource(id, uuid=None, name=None, klass="", type_="device",
                       parent=None, parent_uuid=None, data=None, config=None, extra=None):
    return {
        "id": id,
        "uuid": uuid or f"uuid-{id}",
        "name": name or id,
        "class": klass,
        "type": type_,
        "parent": parent,
        "parent_uuid": parent_uuid or "",
        "description": "",
        "config": config or {},
        "data": data or {},
        "extra": extra or {},
        "position": {"x": 0.0, "y": 0.0, "z": 0.0},
    }


def _simulate_enrichment(raw_data_list):
    tree_set = ResourceTreeSet.from_raw_dict_list(raw_data_list)
    root = tree_set.trees[0].root_node if tree_set.trees else None
    return root.get_plr_nested_dict() if root else {}


def _build_pump_transfer_graph():
    """
    构建带泵/阀门的设备图，用于测试 PumpTransfer:

    flask_water (container)
        ↓
    valve_1 (multiway_valve, pump_1 连接)
        ↓
    reactor_01 (device)

    同时有: stirrer_1, heatchill_1, separator_1
    """
    G = nx.DiGraph()

    # 源容器
    G.add_node("flask_water", **{
        "id": "flask_water", "name": "flask_water",
        "type": "container", "class": "",
        "data": {"reagent_name": "water", "liquid": [{"liquid_type": "water", "volume": 200.0}]},
        "config": {"reagent": "water"},
    })

    # 多通阀
    G.add_node("valve_1", **{
        "id": "valve_1", "name": "valve_1",
        "type": "device", "class": "multiway_valve",
        "data": {}, "config": {},
    })

    # 注射泵（连接到阀门）
    G.add_node("pump_1", **{
        "id": "pump_1", "name": "pump_1",
        "type": "device", "class": "virtual_pump",
        "data": {}, "config": {"max_volume": 25.0},
    })

    # 目标容器
    G.add_node("reactor_01", **{
        "id": "reactor_01", "name": "reactor_01",
        "type": "device", "class": "virtual_stirrer",
        "data": {"liquid": [{"liquid_type": "water", "volume": 50.0}]},
        "config": {},
    })

    # 搅拌器
    G.add_node("stirrer_1", **{
        "id": "stirrer_1", "name": "stirrer_1",
        "type": "device", "class": "virtual_stirrer",
        "data": {}, "config": {},
    })

    # 加热器
    G.add_node("heatchill_1", **{
        "id": "heatchill_1", "name": "heatchill_1",
        "type": "device", "class": "virtual_heatchill",
        "data": {}, "config": {},
    })

    # 分离器
    G.add_node("separator_1", **{
        "id": "separator_1", "name": "separator_1",
        "type": "device", "class": "separator_controller",
        "data": {}, "config": {},
    })

    # 废液容器
    G.add_node("waste_workup", **{
        "id": "waste_workup", "name": "waste_workup",
        "type": "container", "class": "",
        "data": {}, "config": {},
    })

    # 产物收集瓶
    G.add_node("product_flask", **{
        "id": "product_flask", "name": "product_flask",
        "type": "container", "class": "",
        "data": {}, "config": {},
    })

    # DCM溶剂瓶
    G.add_node("flask_dcm", **{
        "id": "flask_dcm", "name": "flask_dcm",
        "type": "container", "class": "",
        "data": {"reagent_name": "dcm", "liquid": [{"liquid_type": "dcm", "volume": 500.0}]},
        "config": {"reagent": "dcm"},
    })

    # 边连接 —— flask_water → valve_1 → reactor_01
    G.add_edge("flask_water", "valve_1", port={"valve_1": "port_1"})
    G.add_edge("valve_1", "reactor_01", port={"valve_1": "port_2"})
    # 阀门 → 泵
    G.add_edge("valve_1", "pump_1")
    G.add_edge("pump_1", "valve_1")
    # 搅拌器 ↔ reactor
    G.add_edge("stirrer_1", "reactor_01")
    # 加热器 ↔ reactor
    G.add_edge("heatchill_1", "reactor_01")
    # 分离器 ↔ reactor
    G.add_edge("separator_1", "reactor_01")
    G.add_edge("reactor_01", "separator_1")
    # DCM → valve → reactor (同一泵路)
    G.add_edge("flask_dcm", "valve_1", port={"valve_1": "port_3"})
    # reactor → valve → product/waste
    G.add_edge("valve_1", "product_flask", port={"valve_1": "port_4"})
    G.add_edge("valve_1", "waste_workup", port={"valve_1": "port_5"})

    return G


def _format_action(action, indent=0):
    """格式化单个 action 为可读字符串"""
    prefix = "  " * indent
    if isinstance(action, list):
        # 并行动作
        lines = [f"{prefix}[PARALLEL]"]
        for sub in action:
            lines.append(_format_action(sub, indent + 1))
        return "\n".join(lines)

    name = action.get("action_name", "?")
    device = action.get("device_id", "")
    kwargs = action.get("action_kwargs", {})
    comment = action.get("_comment", "")
    meta = action.get("_transfer_meta", "")

    parts = [f"{prefix}→ {device}::{name}"]
    if kwargs:
        # 精简输出
        kw_str = ", ".join(f"{k}={v}" for k, v in kwargs.items()
                           if k not in ("progress_message",))
        if kw_str:
            parts.append(f"  kwargs: {{{kw_str}}}")
    if comment:
        parts.append(f"  # {comment}")
    if meta:
        parts.append(f"  meta: {meta}")
    return "\n".join(f"{prefix}{p}" if i > 0 else p for i, p in enumerate(parts))


def _dump_actions(actions, title=""):
    """打印完整动作列表"""
    print(f"\n{'='*70}")
    print(f"  {title}")
    print(f"  总动作数: {len(actions)}")
    print(f"{'='*70}")
    for i, action in enumerate(actions):
        print(f"\n  [{i:02d}] {_format_action(action, indent=2)}")
    print(f"\n{'='*70}\n")


# ==================== PumpTransfer 全链路 ====================

class TestPumpTransferFullChain:
    """PumpTransfer: 包含图路径查找、泵骨架构建、动作序列生成"""

    def test_pump_transfer_basic(self):
        """基础泵转移：flask_water → valve_1 → reactor_01"""
        from unilabos.compile.pump_protocol import generate_pump_protocol

        G = _build_pump_transfer_graph()

        # 检查最短路径
        path = nx.shortest_path(G, "flask_water", "reactor_01")
        print(f"\n最短路径: {path}")
        assert "valve_1" in path

        # 调用编译器
        actions = generate_pump_protocol(
            G=G,
            from_vessel_id="flask_water",
            to_vessel_id="reactor_01",
            volume=10.0,
            flowrate=2.5,
            transfer_flowrate=0.5,
        )

        _dump_actions(actions, "PumpTransfer: flask_water → reactor_01, 10mL")

        # 验证
        assert isinstance(actions, list)
        assert len(actions) > 0
        # 应该有 set_valve_position 和 set_position 动作
        flat = [a for a in actions if isinstance(a, dict)]
        action_names = [a.get("action_name") for a in flat]
        print(f"动作名称列表: {action_names}")
        assert "set_valve_position" in action_names
        assert "set_position" in action_names

    def test_pump_transfer_with_rinsing_enriched_vessel(self):
        """pump_with_rinsing 接收 enriched vessel dict"""
        from unilabos.compile.pump_protocol import generate_pump_protocol_with_rinsing

        G = _build_pump_transfer_graph()

        # 模拟 enrichment
        from_raw = [_make_raw_resource(
            id="flask_water", klass="", type_="container",
            data={"reagent_name": "water", "liquid": [{"liquid_type": "water", "volume": 200.0}]},
        )]
        to_raw = [_make_raw_resource(
            id="reactor_01", klass="virtual_stirrer", type_="device",
        )]

        from_enriched = _simulate_enrichment(from_raw)
        to_enriched = _simulate_enrichment(to_raw)

        print(f"\nfrom_vessel enriched: {json.dumps(from_enriched, indent=2, ensure_ascii=False)[:300]}...")
        print(f"to_vessel enriched: {json.dumps(to_enriched, indent=2, ensure_ascii=False)[:300]}...")

        # get_vessel 兼容
        fid, fdata = get_vessel(from_enriched)
        tid, tdata = get_vessel(to_enriched)
        print(f"from_vessel_id={fid}, to_vessel_id={tid}")
        assert fid == "flask_water"
        assert tid == "reactor_01"

        actions = generate_pump_protocol_with_rinsing(
            G=G,
            from_vessel=from_enriched,
            to_vessel=to_enriched,
            volume=15.0,
            flowrate=2.5,
            transfer_flowrate=0.5,
        )

        _dump_actions(actions, "PumpTransferWithRinsing: flask_water → reactor_01, 15mL (enriched)")

        assert isinstance(actions, list)
        assert len(actions) > 0

    def test_pump_transfer_multi_batch(self):
        """体积 > max_volume 时自动分批"""
        from unilabos.compile.pump_protocol import generate_pump_protocol

        G = _build_pump_transfer_graph()

        # pump_1 的 max_volume = 25mL，转 60mL 应该分 3 批
        actions = generate_pump_protocol(
            G=G,
            from_vessel_id="flask_water",
            to_vessel_id="reactor_01",
            volume=60.0,
            flowrate=2.5,
            transfer_flowrate=0.5,
        )

        _dump_actions(actions, "PumpTransfer 分批: 60mL (max_volume=25mL, 预期 3 批)")

        assert len(actions) > 0
        # 应该有多轮 set_position
        flat = [a for a in actions if isinstance(a, dict)]
        set_position_count = sum(1 for a in flat if a.get("action_name") == "set_position")
        print(f"set_position 动作数: {set_position_count}")
        # 3批 × 2次 (吸液 + 排液) = 6 次 set_position
        assert set_position_count >= 6

    def test_pump_transfer_no_path(self):
        """无路径时返回空"""
        from unilabos.compile.pump_protocol import generate_pump_protocol

        G = _build_pump_transfer_graph()
        G.add_node("isolated_flask", type="container")

        actions = generate_pump_protocol(
            G=G,
            from_vessel_id="isolated_flask",
            to_vessel_id="reactor_01",
            volume=10.0,
        )

        print(f"\n无路径时的动作列表: {actions}")
        assert actions == []

    def test_pump_backbone_filtering(self):
        """验证泵骨架过滤逻辑（电磁阀被跳过）"""
        from unilabos.compile.pump_protocol import generate_pump_protocol

        G = _build_pump_transfer_graph()
        # 添加电磁阀到路径中
        G.add_node("solenoid_valve_1", **{
            "type": "device", "class": "solenoid_valve",
            "data": {}, "config": {},
        })
        # flask_water → solenoid_valve_1 → valve_1 → reactor_01
        G.remove_edge("flask_water", "valve_1")
        G.add_edge("flask_water", "solenoid_valve_1")
        G.add_edge("solenoid_valve_1", "valve_1")

        path = nx.shortest_path(G, "flask_water", "reactor_01")
        print(f"\n含电磁阀的路径: {path}")
        assert "solenoid_valve_1" in path

        actions = generate_pump_protocol(
            G=G,
            from_vessel_id="flask_water",
            to_vessel_id="reactor_01",
            volume=10.0,
        )

        _dump_actions(actions, "PumpTransfer 含电磁阀: flask_water → solenoid → valve_1 → reactor_01")
        # 电磁阀应被跳过，泵骨架只有 valve_1
        assert len(actions) > 0


# ==================== Separate 全链路 ====================

class TestSeparateProtocolFullChain:
    """Separate: 包含 bug 确认和正常路径测试"""

    def test_separate_bug_line_128_fixed(self):
        """验证 separate_protocol.py:128 的 bug 已修复（不再 crash）"""
        from unilabos.compile.separate_protocol import generate_separate_protocol

        G = _build_pump_transfer_graph()

        raw_data = [_make_raw_resource(
            id="reactor_01", klass="virtual_stirrer",
            data={"liquid": [{"liquid_type": "water", "volume": 100.0}]},
        )]
        enriched = _simulate_enrichment(raw_data)

        # 修复前：final_vessel_id, _ = vessel_id 会 crash（字符串解包）
        # 修复后：final_vessel_id = vessel_id，正常返回 action 列表
        result = generate_separate_protocol(
            G=G,
            vessel=enriched,
            purpose="extract",
            product_phase="top",
            product_vessel="product_flask",
            waste_vessel="waste_workup",
            solvent="dcm",
            volume="100 mL",
        )
        assert isinstance(result, list)
        assert len(result) > 0

    def test_separate_manual_workaround(self):
        """
        绕过 line 128 bug，手动测试分离编译器中可以工作的子函数
        """
        from unilabos.compile.separate_protocol import (
            find_separator_device,
            find_separation_vessel_bottom,
        )
        from unilabos.compile.utils.vessel_parser import (
            find_connected_stirrer,
            find_solvent_vessel,
        )
        from unilabos.compile.utils.unit_parser import parse_volume_input
        from unilabos.compile.utils.resource_helper import get_resource_liquid_volume as get_vessel_liquid_volume

        G = _build_pump_transfer_graph()

        # 1. get_vessel 解析 enriched dict
        raw_data = [_make_raw_resource(
            id="reactor_01", klass="virtual_stirrer",
            data={"liquid": [{"liquid_type": "water", "volume": 100.0}]},
        )]
        enriched = _simulate_enrichment(raw_data)
        vessel_id, vessel_data = get_vessel(enriched)
        print(f"\nvessel_id: {vessel_id}")
        print(f"vessel_data: {vessel_data}")
        assert vessel_id == "reactor_01"
        assert vessel_data["liquid"][0]["volume"] == 100.0

        # 2. find_separator_device
        sep = find_separator_device(G, vessel_id)
        print(f"分离器设备: {sep}")
        assert sep == "separator_1"

        # 3. find_connected_stirrer
        stirrer = find_connected_stirrer(G, vessel_id)
        print(f"搅拌器设备: {stirrer}")
        assert stirrer == "stirrer_1"

        # 4. find_solvent_vessel
        solvent_v = find_solvent_vessel(G, "dcm")
        print(f"DCM溶剂容器: {solvent_v}")
        assert solvent_v == "flask_dcm"

        # 5. parse_volume_input
        vol = parse_volume_input("200 mL")
        print(f"体积解析: '200 mL' → {vol}")
        assert vol == 200.0

        vol2 = parse_volume_input("1.5 L")
        print(f"体积解析: '1.5 L' → {vol2}")
        assert vol2 == 1500.0

        # 6. get_vessel_liquid_volume
        liq_vol = get_vessel_liquid_volume(enriched)
        print(f"液体体积 (enriched dict): {liq_vol}")
        assert liq_vol == 100.0

        # 7. find_separation_vessel_bottom
        bottom = find_separation_vessel_bottom(G, vessel_id)
        print(f"分离容器底部: {bottom}")
        # 当前图中没有命名匹配的底部容器

    def test_pump_transfer_for_separate_subflow(self):
        """测试 separate 中调用的 pump 子流程（溶剂添加 → 分液漏斗）"""
        from unilabos.compile.pump_protocol import generate_pump_protocol_with_rinsing

        G = _build_pump_transfer_graph()

        # 模拟分离前的溶剂添加步骤
        actions = generate_pump_protocol_with_rinsing(
            G=G,
            from_vessel="flask_dcm",
            to_vessel="reactor_01",
            volume=100.0,
            flowrate=2.5,
            transfer_flowrate=0.5,
        )

        _dump_actions(actions, "Separate 子流程: flask_dcm → reactor_01, 100mL DCM")

        assert isinstance(actions, list)
        assert len(actions) > 0

        # 模拟分离后产物转移
        actions2 = generate_pump_protocol_with_rinsing(
            G=G,
            from_vessel="reactor_01",
            to_vessel="product_flask",
            volume=50.0,
            flowrate=2.5,
            transfer_flowrate=0.5,
        )

        _dump_actions(actions2, "Separate 子流程: reactor_01 → product_flask, 50mL 产物")

        assert len(actions2) > 0

        # 废液转移
        actions3 = generate_pump_protocol_with_rinsing(
            G=G,
            from_vessel="reactor_01",
            to_vessel="waste_workup",
            volume=50.0,
            flowrate=2.5,
            transfer_flowrate=0.5,
        )

        _dump_actions(actions3, "Separate 子流程: reactor_01 → waste_workup, 50mL 废液")

        assert len(actions3) > 0


# ==================== 图路径可视化 ====================

class TestGraphPathVisualization:
    """输出图中关键路径信息"""

    def test_all_shortest_paths(self):
        """输出所有容器之间的最短路径"""
        G = _build_pump_transfer_graph()

        containers = [n for n in G.nodes() if G.nodes[n].get("type") == "container"]
        devices = [n for n in G.nodes() if G.nodes[n].get("type") == "device"]

        print(f"\n{'='*70}")
        print(f"  设备图概览")
        print(f"{'='*70}")
        print(f"  容器节点 ({len(containers)}): {containers}")
        print(f"  设备节点 ({len(devices)}): {devices}")
        print(f"  边数: {G.number_of_edges()}")
        print(f"  边列表:")
        for u, v, data in G.edges(data=True):
            port_info = data.get("port", "")
            print(f"    {u} → {v}  {port_info if port_info else ''}")

        print(f"\n  关键路径:")
        pairs = [
            ("flask_water", "reactor_01"),
            ("flask_dcm", "reactor_01"),
            ("reactor_01", "product_flask"),
            ("reactor_01", "waste_workup"),
            ("flask_water", "product_flask"),
        ]
        for src, dst in pairs:
            try:
                path = nx.shortest_path(G, src, dst)
                length = len(path) - 1
                # 标注路径上的节点类型
                annotated = []
                for n in path:
                    ntype = G.nodes[n].get("type", "?")
                    nclass = G.nodes[n].get("class", "")
                    annotated.append(f"{n}({ntype}{'/' + nclass if nclass else ''})")
                print(f"    {src} → {dst}: 距离={length}")
                print(f"      路径: {' → '.join(annotated)}")
            except nx.NetworkXNoPath:
                print(f"    {src} → {dst}: 无路径!")

        print(f"{'='*70}\n")