tests/test_loading_plan_url.py

"""
Tests for _build_loading_plan with url-type expression sources (Phase 3).

Verifies that url-type manifests produce the correct tool sequence:
  - decoupler_load_url_counts
  - decoupler_join_clinical_metadata  (when metadata_source has a URL)
  - decoupler_inspect_data
  - decoupler_load_and_filter_data + decoupler_preprocess_data  (Path A only)
  - dataset_validate_contrast
  - decoupler_differential_expression

No heavy dependencies, no file I/O, no network.
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from src.datasets.manifest_schema import DatasetManifest
from src.tools.dataset_tools import _build_loading_plan

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _url_manifest(
    data_level: str = "normalized",
    feature_id_type: str = "gene_symbol",
    expr_url: str = "https://xena.example.com/counts.tsv.gz",
    meta_url: str | None = None,
    meta_embedded: bool = False,
    survival_event_col: str | None = None,
    collapsed_url: str | None = None,
    feature_mapping: dict | None = None,
) -> DatasetManifest:
    """Build a minimal url-type DatasetManifest."""
    meta_source: dict = {"type": "supplementary_table", "embedded": meta_embedded}
    if meta_url:
        meta_source["url"] = meta_url
        meta_source["join_column"] = "sampleID"
        meta_source["truncate_to_patient"] = True

    expression_source = {"type": "url", "url": expr_url}
    if collapsed_url:
        expression_source["collapsed_url"] = collapsed_url

    return DatasetManifest(
        dataset_id="test_tcga",
        title="Test TCGA dataset",
        accession="TCGA-TEST",
        organism="human",
        modality="bulk_rnaseq",
        platform="Illumina HiSeq",
        data_level=data_level,
        feature_id_type=feature_id_type,
        expression_source=expression_source,
        metadata_source=meta_source,
        feature_mapping=feature_mapping or {},
        group_columns=["pathologic_stage"],
        valid_workflows=["activity_scoring", "survival"],
        limitations=["Test dataset only"],
        survival_columns={
            "event_column": survival_event_col,
            "time_column": "os_days" if survival_event_col else None,
        },
        default_contrasts=[
            {
                "design_factor": "pathologic_stage",
                "test_group": "Stage II",
                "control_group": "Stage I",
                "method": "ttest",
            }
        ],
    )


_PROBE_FEATURE_MAPPING = {
    "requires_collapse": True,
    "gene_symbol_column": "Gene Symbol",
    "collapse_method": "mean",
    "multi_gene_policy": "drop",
}


def _tool_sequence(steps: list[dict]) -> list[str]:
    return [s["tool"] for s in steps]


# ---------------------------------------------------------------------------
# Path B, no clinical join
# ---------------------------------------------------------------------------


class TestUrlPlanPathB:
    def test_normalized_no_clinical_url(self):
        """No metadata URL → no join step."""
        manifest = _url_manifest(data_level="normalized", meta_url=None)
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        assert tools[0] == "decoupler_load_url_counts"
        assert "decoupler_join_clinical_metadata" not in tools
        assert "decoupler_inspect_data" in tools
        assert "decoupler_load_and_filter_data" not in tools
        assert "decoupler_preprocess_data" not in tools
        assert tools[-1] == "decoupler_differential_expression"

    def test_normalized_with_clinical_url(self):
        """Metadata URL present → join step inserted after load."""
        manifest = _url_manifest(
            data_level="normalized",
            meta_url="https://xena.example.com/clinical.tsv",
        )
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        assert tools[0] == "decoupler_load_url_counts"
        assert tools[1] == "decoupler_join_clinical_metadata"
        assert "decoupler_inspect_data" in tools
        assert "decoupler_load_and_filter_data" not in tools
        assert tools[-1] == "decoupler_differential_expression"

    def test_clinical_join_step_has_correct_args(self):
        manifest = _url_manifest(
            data_level="normalized",
            meta_url="https://xena.example.com/clinical.tsv",
            survival_event_col="os_event",
        )
        steps = _build_loading_plan(manifest)
        join_step = next(s for s in steps if s["tool"] == "decoupler_join_clinical_metadata")

        assert (
            join_step["key_args"]["clinical_url_or_path"] == "https://xena.example.com/clinical.tsv"
        )
        assert join_step["key_args"]["barcode_column"] == "sampleID"
        assert join_step["key_args"]["truncate_to_patient"] is True
        assert join_step["key_args"]["add_survival_columns"] is True

    def test_embedded_metadata_no_join(self):
        """embedded=True → no join step even if type is supplementary_table."""
        manifest = _url_manifest(
            data_level="normalized",
            meta_embedded=True,
            meta_url=None,
        )
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)
        assert "decoupler_join_clinical_metadata" not in tools

    def test_load_url_counts_args(self):
        manifest = _url_manifest(
            data_level="normalized",
            feature_id_type="gene_symbol",
            expr_url="https://xena.example.com/HiSeqV2.gz",
        )
        steps = _build_loading_plan(manifest)
        load_step = steps[0]

        assert load_step["tool"] == "decoupler_load_url_counts"
        assert load_step["key_args"]["url_or_path"] == "https://xena.example.com/HiSeqV2.gz"
        assert load_step["key_args"]["feature_id_type"] == "gene_symbol"
        assert load_step["key_args"]["strip_ensembl_versions"] is False  # gene_symbol, not ensembl

    def test_ensembl_strip_versions_true(self):
        manifest = _url_manifest(
            data_level="normalized",
            feature_id_type="ensembl_gene_id",
        )
        steps = _build_loading_plan(manifest)
        load_step = steps[0]
        assert load_step["key_args"]["strip_ensembl_versions"] is True

    def test_step_numbers_sequential(self):
        manifest = _url_manifest(
            data_level="normalized",
            meta_url="https://xena.example.com/clinical.tsv",
        )
        steps = _build_loading_plan(manifest)
        for i, step in enumerate(steps, start=1):
            assert step["step"] == i, f"Step {i} has step number {step['step']}"

    def test_de_method_is_ttest_for_path_b(self):
        manifest = _url_manifest(data_level="normalized")
        steps = _build_loading_plan(manifest)
        de_step = steps[-1]
        assert de_step["key_args"]["method"] == "ttest"


# ---------------------------------------------------------------------------
# Path A (raw counts → DESeq2)
# ---------------------------------------------------------------------------


class TestUrlPlanPathA:
    def test_raw_counts_includes_filter_and_preprocess(self):
        """Path A: filter + preprocess steps before validate and DE."""
        manifest = _url_manifest(data_level="raw_counts")
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        assert tools[0] == "decoupler_load_url_counts"
        assert "decoupler_load_and_filter_data" in tools
        assert "decoupler_preprocess_data" in tools
        assert tools[-1] == "decoupler_differential_expression"

    def test_filter_before_preprocess(self):
        manifest = _url_manifest(data_level="raw_counts")
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        filter_idx = tools.index("decoupler_load_and_filter_data")
        preprocess_idx = tools.index("decoupler_preprocess_data")
        de_idx = tools.index("decoupler_differential_expression")
        assert filter_idx < preprocess_idx < de_idx

    def test_raw_counts_de_method_deseq2(self):
        """Path A: manifest default_contrast uses deseq2."""
        manifest = DatasetManifest(
            dataset_id="tcga_raw",
            title="TCGA raw",
            accession="TCGA-RAW",
            organism="human",
            modality="bulk_rnaseq",
            platform="Illumina HiSeq",
            data_level="raw_counts",
            feature_id_type="ensembl_gene_id",
            expression_source={"type": "url", "url": "https://example.com/counts.tsv"},
            metadata_source={"type": "supplementary_table", "embedded": True},
            group_columns=["condition"],
            valid_workflows=["activity_scoring"],
            limitations=["Test only"],
            default_contrasts=[
                {
                    "design_factor": "condition",
                    "test_group": "tumor",
                    "control_group": "normal",
                    "method": "deseq2",
                }
            ],
        )
        steps = _build_loading_plan(manifest)
        de_step = steps[-1]
        assert de_step["key_args"]["method"] == "deseq2"

    def test_clinical_join_before_filter_in_path_a(self):
        """Clinical join should come before filter/preprocess in Path A."""
        manifest = _url_manifest(
            data_level="raw_counts",
            meta_url="https://xena.example.com/clinical.tsv",
        )
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        join_idx = tools.index("decoupler_join_clinical_metadata")
        filter_idx = tools.index("decoupler_load_and_filter_data")
        assert join_idx < filter_idx


# ---------------------------------------------------------------------------
# GEO path still works (regression)
# ---------------------------------------------------------------------------


class TestGeoPathRegression:
    def test_geo_plan_unchanged(self):
        """Existing GEO datasets still produce the correct plan."""
        manifest = DatasetManifest(
            dataset_id="test_geo",
            title="Test GEO",
            accession="GSE99999",
            organism="human",
            modality="bulk_microarray",
            platform="Agilent",
            data_level="log_expression",
            feature_id_type="gene_symbol",
            expression_source={
                "type": "geo_series_matrix",
                "url": "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE99nnn/GSE99999/matrix/GSE99999_series_matrix.txt.gz",
            },
            metadata_source={"type": "geo_series_matrix", "embedded": True},
            group_columns=["condition"],
            valid_workflows=["microarray"],
            limitations=["Test only"],
            default_contrasts=[
                {
                    "design_factor": "condition",
                    "test_group": "tumor",
                    "control_group": "normal",
                    "method": "ttest",
                }
            ],
        )
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        assert tools[0] == "decoupler_load_geo_series_matrix"
        assert "decoupler_join_clinical_metadata" not in tools
        assert "decoupler_load_and_filter_data" not in tools
        assert tools[-1] == "decoupler_differential_expression"


# ---------------------------------------------------------------------------
# Precompute cache Phase 2: collapsed_url skips the live collapse step
# ---------------------------------------------------------------------------


class TestCollapsedUrlPrecompute:
    def test_probe_dataset_without_collapsed_url_has_collapse_step(self):
        """Baseline: probe-indexed dataset without collapsed_url still gets
        the live decoupler_collapse_probes_to_genes step."""
        manifest = _url_manifest(
            data_level="normalized",
            feature_id_type="probe_id",
            feature_mapping=_PROBE_FEATURE_MAPPING,
        )
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        assert "decoupler_collapse_probes_to_genes" in tools
        inspect_step = next(s for s in steps if s["tool"] == "decoupler_inspect_data")
        assert "features_look_like_probes=True" in inspect_step["note"]
        load_step = steps[0]
        assert load_step["key_args"]["url_or_path"] == manifest.expression_source["url"]

    def test_probe_dataset_with_collapsed_url_skips_collapse_step(self):
        """collapsed_url present → no live collapse step; load step points
        at the precomputed gene-collapsed h5ad instead."""
        collapsed_url = "https://huggingface.co/datasets/anne-voigt/pdac-research-data/resolve/main/test_tcga_collapsed.h5ad"
        manifest = _url_manifest(
            data_level="normalized",
            feature_id_type="probe_id",
            feature_mapping=_PROBE_FEATURE_MAPPING,
            collapsed_url=collapsed_url,
        )
        steps = _build_loading_plan(manifest)
        tools = _tool_sequence(steps)

        assert "decoupler_collapse_probes_to_genes" not in tools

        load_step = steps[0]
        assert load_step["tool"] == "decoupler_load_url_counts"
        assert load_step["key_args"]["url_or_path"] == collapsed_url
        assert "precomputed gene-collapsed" in load_step["note"]

        inspect_step = next(s for s in steps if s["tool"] == "decoupler_inspect_data")
        assert "features_look_like_probes=False" in inspect_step["note"]

    def test_non_probe_dataset_collapsed_url_is_noop(self):
        """collapsed_url on a non-probe dataset (requires_collapse=False)
        has no effect — load step still uses the regular url."""
        manifest = _url_manifest(
            data_level="normalized",
            feature_id_type="gene_symbol",
            collapsed_url="https://example.com/should_be_ignored.h5ad",
        )
        steps = _build_loading_plan(manifest)
        load_step = steps[0]
        assert load_step["key_args"]["url_or_path"] == manifest.expression_source["url"]
        assert "decoupler_collapse_probes_to_genes" not in _tool_sequence(steps)

    def test_dataset_describe_reports_precomputed_flag(self, monkeypatch):
        """dataset_describe's analysis_guidance.probe_collapse_precomputed
        reflects whether collapsed_url is set for a probe-indexed dataset."""
        import src.tools.dataset_tools as dataset_tools

        manifest = _url_manifest(
            data_level="normalized",
            feature_id_type="probe_id",
            feature_mapping=_PROBE_FEATURE_MAPPING,
            collapsed_url="https://huggingface.co/datasets/anne-voigt/pdac-research-data/resolve/main/test_tcga_collapsed.h5ad",
        )
        monkeypatch.setattr(dataset_tools.catalog, "load_manifest", lambda dataset_id: manifest)

        desc = dataset_tools.dataset_describe("test_tcga")
        guidance = desc["analysis_guidance"]
        assert guidance["probe_collapse_required"] is True
        assert guidance["probe_collapse_precomputed"] is True