PlantAI-dev/tests/unit/test_subfamily_prediction_plan.py

import importlib.util
import json
from pathlib import Path
from types import SimpleNamespace

import pytest

from gfa_core.gene_family_stage_context import (
    enrich_stage_params_with_context,
    infer_analysis_route,
)
from gfa_core.mas.artifact_resolver import ArtifactResolver


def _load_vue3_main():
    module_path = Path(__file__).resolve().parents[2] / "apps" / "web" / "vue3" / "main.py"
    spec = importlib.util.spec_from_file_location("bbav5_vue3_main_plan_test", module_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_subfamily_prediction_plan_uses_per_target_evidence_route():
    main = _load_vue3_main()
    plan = {"stages": ["identify", "reference_classification"], "stages_detail": []}
    intent = {"stages": ["identify", "reference_classification"]}

    stage_ids, details = main._normalize_subfamily_prediction_plan(
        plan,
        intent,
        plan["stages"],
        plan["stages_detail"],
    )

    assert stage_ids == [
        "identify",
        "phylogeny",
        "motif",
        "domain",
        "structure",
        "combo",
        "reference_classification",
    ]
    detail_by_id = {item["id"]: item for item in details}
    assert detail_by_id["phylogeny"]["dependencies"] == ["identify"]
    assert detail_by_id["combo"]["dependencies"] == ["motif", "domain", "structure"]
    assert detail_by_id["reference_classification"]["dependencies"] == ["combo"]
    assert detail_by_id["reference_classification"]["params"]["subfamilyPipeline"] is True
    assert detail_by_id["reference_classification"]["params"]["perTargetEvidenceOnly"] is True
    assert detail_by_id["reference_classification"]["params"]["analysisRoute"] == "subfamily_prediction"
    assert plan["total_stages"] == 7
    assert intent["primary_stage"] == "reference_classification"


def test_subfamily_request_roles_do_not_promote_reference_to_target(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_configured_species_scope",
        lambda _job_id: {
            "targetSpecies": ["Os"],
            "referenceSpecies": ["At"],
            "comparativeSpecies": [],
        },
    )
    monkeypatch.setattr(
        main,
        "_extract_species_codes_from_text_for_job",
        lambda _job_id, _message: ["Os", "At"],
    )

    roles = main._resolve_subfamily_request_species_roles(
        "job-real",
        "执行目标物种 Os 和参考物种 At 的亚家族分类",
        {
            "target_species": ["Os"],
            "reference_species": ["At"],
            "params": {"targetSpecies": ["Os"], "referenceSpecies": ["At"]},
        },
    )

    assert roles == {
        "target_species": ["Os"],
        "reference_species": ["At"],
        "comparative_species": [],
        "classification_species": ["Os"],
    }


def test_contextual_gene_family_request_handles_explicit_multi_target_scope(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_canonicalize_species_codes_for_job",
        lambda _job_id, species: list(species),
    )
    monkeypatch.setattr(
        main,
        "_reconcile_confirmed_gene_family_species",
        lambda *_args, **_kwargs: {"applicable": False, "conflicts": []},
    )
    monkeypatch.setattr(main, "_load_identify_species_selection", lambda _job_id: {})
    monkeypatch.setattr(main, "_gene_family_from_config", lambda _path: "SPL")
    monkeypatch.setattr(main, "_confirmed_subfamily_species", lambda *_args: [])
    monkeypatch.setattr(
        main,
        "_extract_species_codes_from_text_for_job",
        lambda *_args: [],
    )
    monkeypatch.setattr(
        main,
        "_subfamily_state_for_targets",
        lambda *_args: {
            "all_confirmed": False,
            "by_species": {
                "Ft": {"status": "not_classified"},
                "At": {"status": "not_classified"},
            },
            "issues": [],
        },
    )
    monkeypatch.setattr(
        main,
        "_validate_gene_family_stage_context",
        lambda **_kwargs: {"valid": True, "issues": [], "evidence": {}},
    )

    contextual = main._resolve_contextual_gene_family_request(
        "job-multi-target",
        "执行基因家族分析",
        requested_stages=["phylogeny"],
        request_intent={
            "gene_family": "SPL",
            "target_species": ["Ft", "At"],
            "reference_species": [],
            "params": {
                "targetSpecies": ["Ft", "At"],
                "referenceSpecies": [],
            },
        },
    )

    assert contextual["targetSpecies"] == ["Ft", "At"]
    assert contextual["needs_species_confirmation"] is False


def test_generic_target_reference_scope_does_not_guess_subfamily_from_stage_name():
    scope = {
        "scope_mode": "target_plus_reference",
        "target_species": ["Ft"],
        "reference_species": ["At"],
        "comparative_species": [],
    }

    assert infer_analysis_route("phylogeny", scope) == "multi_species"
    assert infer_analysis_route("motif", scope) == "multi_species"


def test_newick_manifest_artifact_satisfies_tree_requirement():
    assert ArtifactResolver._kind_matches("tree", "newick") is True


def test_semantic_preflight_materializes_selected_tree(monkeypatch):
    main = _load_vue3_main()

    class Selected:
        path = "/workspace/final/Target.PEBP.subfamily_evidence.tree.nwk"

    class Decision:
        resolved = True
        selected = Selected()
        candidates = []

        def to_dict(self):
            return {
                "resolved": True,
                "selected": {
                    "path": self.selected.path,
                    "kind": "newick",
                    "source": "artifact:phylogeny:tree",
                },
                "candidates": [],
            }

    class Resolver:
        def __init__(self, _context):
            pass

        def resolve_stage(self, stage, species="", family=""):
            return {
                "family_members_pep": Decision(),
                "tree": Decision(),
            }

    context = SimpleNamespace(
        family="PEBP",
        target_species=["Os"],
        reference_species=["At"],
        comparative_species=[],
        config_path="/workspace/input/config.xlsx",
    )
    monkeypatch.setattr(main, "build_project_context", lambda **_kwargs: context)
    monkeypatch.setattr(main, "ArtifactResolver", Resolver)
    monkeypatch.setattr(
        main,
        "_validate_gene_family_stage_context",
        lambda **_kwargs: {"valid": True, "issues": [], "evidence": {}},
    )

    result = main._api_semantic_preflight(
        "job-direct-motif",
        "motif",
        "Os.PEBP",
        {
            "geneFamily": "PEBP",
            "targetSpecies": ["Os"],
            "referenceSpecies": ["At"],
            "scopeMode": "target_plus_reference",
            "analysisRoute": "subfamily_evidence",
        },
    )

    assert result["valid"] is True
    assert result["materialized_params"]["treeFile"] == Selected.path


def test_scoped_tree_resolution_prefers_preflight_selected_tree(tmp_path):
    from apps.web import pipeline_backend

    workspace = tmp_path / "workspace"
    final_dir = workspace / "final"
    final_dir.mkdir(parents=True)
    scope_file = final_dir / "Os.PEBP.scope.json"
    scope_file.write_text(
        json.dumps({
            "scope_id": "PEBP.target-Os.ref-At.target_plus_reference",
            "scope_mode": "target_plus_reference",
            "target_species": ["Os"],
            "reference_species": ["At"],
        }),
        encoding="utf-8",
    )
    selected_tree = final_dir / "Target.PEBP.subfamily_evidence.tree.nwk"
    selected_tree.write_text("(Os1,At1);\n", encoding="utf-8")

    resolved, _scope = pipeline_backend.resolve_scoped_tree_from_phylogeny(
        workspace,
        final_dir,
        "Os.PEBP",
        scope_file,
        explicit_tree_file=str(selected_tree),
    )

    assert resolved == selected_tree.resolve()


def test_explicit_subfamily_and_regular_routes_remain_separate():
    shared = {
        "geneFamily": "SPL",
        "targetSpecies": ["Ft", "Os"],
        "referenceSpecies": ["At"],
        "scopeMode": "target_plus_reference",
        "analysisMode": "joint_multi_species",
    }

    evidence = enrich_stage_params_with_context(
        "motif",
        "Target.SPL",
        {**shared, "analysisRoute": "subfamily_evidence"},
    )
    prediction = enrich_stage_params_with_context(
        "reference_classification",
        "Target.SPL",
        {**shared, "analysisRoute": "subfamily_prediction"},
    )
    regular = enrich_stage_params_with_context("motif", "Target.SPL", shared)

    assert evidence["analysisRoute"] == "subfamily_evidence"
    assert evidence["memberSetId"].endswith(":subfamily_evidence:domain_refined")
    assert evidence["analysisMode"] == "joint_multi_species"
    assert prediction["analysisRoute"] == "subfamily_prediction"
    assert regular["analysisRoute"] == "multi_species"
    assert regular["memberSetId"].endswith(":multi_species:all_members")


def test_direct_motif_request_uses_job_subfamily_state(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_load_identify_species_selection",
        lambda _job_id: {
            "groups": {
                "target": [{"species": "Ft"}],
                "reference": [{"species": "At"}],
                "comparative": [],
            }
        },
    )
    monkeypatch.setattr(main, "_gene_family_from_config", lambda _path: "SPL")
    monkeypatch.setattr(main, "_confirmed_subfamily_species", lambda _job_id, _family: [])
    monkeypatch.setattr(main, "_extract_species_codes_from_text_for_job", lambda _job_id, _text: [])
    monkeypatch.setattr(
        main,
        "_subfamily_state_for_targets",
        lambda _job_id, _species, _family: {
            "all_confirmed": False,
            "by_species": {"Ft": {"status": "not_classified"}},
            "issues": [{"code": "subfamily_classification_missing"}],
        },
    )
    monkeypatch.setattr(
        main,
        "_validate_gene_family_stage_context",
        lambda **_kwargs: {"valid": True, "issues": [], "evidence": {}},
    )

    contextual = main._resolve_contextual_gene_family_request(
        "job-direct-motif",
        "执行 Motif 分析",
        "Ft.SPL",
        requested_stages=["motif"],
    )

    assert contextual["stages"] == ["motif"]
    assert contextual["targetSpecies"] == ["Ft"]
    assert contextual["referenceSpecies"] == ["At"]
    assert contextual["contextGate"]["needs_subfamily_decision"] is True


def test_direct_stage_preserves_species_roles_from_intent_when_selector_is_empty(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(main, "_load_identify_species_selection", lambda _job_id: {})
    monkeypatch.setattr(main, "_gene_family_from_config", lambda _path: "PEBP")
    monkeypatch.setattr(main, "_confirmed_subfamily_species", lambda _job_id, _family: [])
    monkeypatch.setattr(main, "_extract_species_codes_from_text_for_job", lambda _job_id, _text: [])
    monkeypatch.setattr(
        main,
        "_subfamily_state_for_targets",
        lambda _job_id, _species, _family: {
            "all_confirmed": False,
            "by_species": {"Os": {"status": "not_classified"}},
            "issues": [],
        },
    )
    monkeypatch.setattr(
        main,
        "_validate_gene_family_stage_context",
        lambda **_kwargs: {"valid": True, "issues": [], "evidence": {}},
    )

    contextual = main._resolve_contextual_gene_family_request(
        "job-direct-motif",
        "执行 Motif 分析",
        "Os.PEBP",
        requested_stages=["motif"],
        request_intent={
            "target_species": ["Os"],
            "reference_species": ["At"],
            "params": {
                "targetSpecies": ["Os"],
                "referenceSpecies": ["At"],
            },
        },
    )

    assert contextual["targetSpecies"] == ["Os"]
    assert contextual["referenceSpecies"] == ["At"]
    assert contextual["scopeMode"] == "target_plus_reference"

    request = SimpleNamespace(
        job_id="job-direct-motif",
        message="执行 Motif 分析",
        prefix="Os.PEBP",
        intent={
            "stages": ["motif"],
            "params": {
                "targetSpecies": ["Os"],
                "referenceSpecies": ["At"],
            },
        },
    )
    response = main._build_contextual_gene_family_plan_response(
        request,
        "session-direct-motif",
        contextual,
    )
    details = {item["id"]: item for item in response["plan"]["all_stages_detail"]}

    assert details["phylogeny"]["params"]["analysisRoute"] == "subfamily_evidence"
    assert details["motif"]["params"]["analysisRoute"] == "subfamily_evidence"
    assert details["phylogeny"]["params"]["memberSetId"] == details["motif"]["params"]["memberSetId"]


def test_contextual_handoff_uses_confirmed_config_roles_when_selector_is_empty(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(main, "_load_identify_species_selection", lambda _job_id: {})
    monkeypatch.setattr(
        main,
        "_confirmed_config_species_scope",
        lambda _job_id: {
            "targetSpecies": ["Os"],
            "referenceSpecies": ["At"],
            "comparativeSpecies": [],
        },
    )
    monkeypatch.setattr(main, "_gene_family_from_config", lambda _path: "PEBP")
    monkeypatch.setattr(main, "_confirmed_subfamily_species", lambda _job_id, _family: [])
    monkeypatch.setattr(
        main,
        "_extract_species_codes_from_text_for_job",
        lambda _job_id, _text: ["Os", "At"],
    )
    monkeypatch.setattr(
        main,
        "_subfamily_state_for_targets",
        lambda _job_id, species, _family: {
            "all_confirmed": species == ["Os"],
            "by_species": {"Os": {"status": "confirmed"}},
            "issues": [],
        },
    )
    monkeypatch.setattr(
        main,
        "_validate_gene_family_stage_context",
        lambda **_kwargs: {"valid": True, "issues": [], "evidence": {}},
    )

    contextual = main._resolve_contextual_gene_family_request(
        "job-real",
        (
            "执行水稻PEBP基因家族分析。基因家族 PEBP；"
            "目标物种 Os；参考物种 At；亚家族分类已经确认。"
        ),
        "Os.PEBP",
        requested_stages=["phylogeny", "motif", "domain", "structure", "combo"],
        request_intent={
            "gene_family": "PEBP",
            "target_species": ["Os"],
            "reference_species": ["At"],
            "params": {
                "targetSpecies": ["Os"],
                "referenceSpecies": ["At"],
            },
        },
    )

    assert contextual["targetSpecies"] == ["Os"]
    assert contextual["referenceSpecies"] == ["At"]
    assert contextual["prefix"] == "Os.PEBP"
    assert contextual["context_route"] == "post_subfamily_gene_family_analysis"


def test_confirmed_multi_species_roles_remain_separate(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_confirmed_config_species_scope",
        lambda _job_id: {
            "targetSpecies": ["Ft"],
            "referenceSpecies": ["At"],
            "comparativeSpecies": ["Ppy"],
        },
    )

    resolved = main._reconcile_confirmed_gene_family_species(
        "job-multi",
        "Ft是目标物种，Ppy是比较物种，At是参考物种，执行联合亚家族分类",
        {
            "gene_family": "SPL",
            "target_species": ["Ft", "Ppy"],
            "reference_species": ["At"],
            "params": {
                "targetSpecies": ["Ft", "Ppy"],
                "referenceSpecies": ["At"],
            },
        },
    )

    assert resolved["conflicts"] == []
    assert resolved["target_species"] == ["Ft"]
    assert resolved["comparative_species"] == ["Ppy"]
    assert resolved["reference_species"] == ["At"]
    assert resolved["classification_species"] == ["Ft", "Ppy"]


def test_confirmation_ack_normalizes_structured_role_drift(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_confirmed_config_species_scope",
        lambda _job_id: {
            "targetSpecies": ["Ft"],
            "referenceSpecies": ["At"],
            "comparativeSpecies": [],
        },
    )

    resolved = main._reconcile_confirmed_gene_family_species(
        "job-confirmed-ack",
        "是",
        {
            "gene_family": "SPL",
            "stages": ["identify"],
            "target_species": ["Ft", "At"],
            "params": {"targetSpecies": ["Ft", "At"]},
        },
    )

    assert resolved["conflicts"] == []
    assert resolved["target_species"] == ["Ft"]
    assert resolved["reference_species"] == ["At"]


def test_unassigned_reference_species_analysis_request_remains_blocked(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_confirmed_config_species_scope",
        lambda _job_id: {
            "targetSpecies": ["Ft"],
            "referenceSpecies": ["At"],
            "comparativeSpecies": [],
        },
    )

    resolved = main._reconcile_confirmed_gene_family_species(
        "job-unassigned-reference",
        "分析 At 的 SPL 基因家族",
        {
            "gene_family": "SPL",
            "stages": ["identify"],
            "target_species": ["At"],
            "params": {"targetSpecies": ["At"]},
        },
    )

    assert any(
        conflict["code"] == "analysis_target_config_role_conflict"
        and conflict["species"] == "At"
        for conflict in resolved["conflicts"]
    )


def test_confirmed_config_duplicate_roles_are_blocked(monkeypatch):
    main = _load_vue3_main()
    monkeypatch.setattr(
        main,
        "_confirmed_config_species_scope",
        lambda _job_id: {
            "targetSpecies": ["Ft"],
            "referenceSpecies": ["Ft", "At"],
            "comparativeSpecies": [],
        },
    )

    resolved = main._reconcile_confirmed_gene_family_species(
        "job-invalid-config",
        "执行SPL基因家族分析",
        {"gene_family": "SPL", "stages": ["identify"]},
    )

    assert any(
        conflict["code"] == "config_species_role_conflict"
        and conflict["species"] == "Ft"
        for conflict in resolved["conflicts"]
    )
    assert "打开 config.xlsx" in main._species_role_conflict_message(resolved["conflicts"])


@pytest.mark.asyncio
async def test_plan_api_blocks_confirmed_config_role_conflict(monkeypatch, tmp_path):
    import pandas as pd
    from gfa_core.mas import confirmation_snapshot

    main = _load_vue3_main()
    input_root = tmp_path / "input"
    job_dir = input_root / "job-os"
    job_dir.mkdir(parents=True)
    pd.DataFrame(
        [
            {"Species_Code": "Os", "Species_Role": "target", "Ref": False},
            {"Species_Code": "At", "Species_Role": "reference", "Ref": True},
        ]
    ).to_excel(job_dir / "config.xlsx", index=False)
    monkeypatch.setattr(main, "INPUT_ROOT", input_root)
    monkeypatch.setattr(
        confirmation_snapshot,
        "build_confirmation_snapshot",
        lambda *_args, **_kwargs: {
            "config_confirmed": True,
            "samples_confirmed": False,
            "subfamily_confirmed": False,
            "config": {"artifact_path": str(job_dir / "config.xlsx")},
            "samples": None,
            "subfamily": None,
            "errors": (),
        },
    )

    result = await main.agent_plan(
        main.AgentPlanRequest(
            message="At是目标物种，Os是参考物种，请执行PEBP基因家族分析",
            session_id="role-conflict",
            job_id="job-os",
            prefix="Os.PEBP",
            intent={
                "gene_family": "PEBP",
                "stages": ["identify"],
                "target_species": ["At"],
                "reference_species": ["Os"],
                "params": {
                    "geneFamily": "PEBP",
                    "targetSpecies": ["At"],
                    "referenceSpecies": ["Os"],
                },
            },
        )
    )

    assert result["success"] is False
    assert result["status"] == "species_role_conflict"
    assert "打开 config.xlsx" in result["message"]
    assert "plan" not in result