tests/test_microarray_workflow.py

"""
Tests for src/workflows/microarray.py — expression-level functions.

All tests use small in-memory toy DataFrames or temporary CSV files.
No AnnData, scanpy, or heavy bioinformatics dependencies required.

Toy dataset
-----------
10 samples (S001–S010), 5 genes (GENE_A – GENE_E).
Group A: S001–S005 (GENE_A intentionally up-regulated by +2 units).
Group B: S006–S010.
Values are in a log2-like range (8–12) with small Gaussian noise.
"""

from __future__ import annotations

import sys
import tempfile
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

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

from src.workflows.microarray import (
    collapse_duplicate_genes,
    detect_log_scale,
    harmonize_expression_and_metadata,
    load_expression_matrix,
    prepare_gene_level_statistics,
)

# ---------------------------------------------------------------------------
# Shared toy data
# ---------------------------------------------------------------------------

_RNG = np.random.default_rng(42)
_SAMPLE_IDS = [f"S{i:03d}" for i in range(1, 11)]
_GENES = ["GENE_A", "GENE_B", "GENE_C", "GENE_D", "GENE_E"]

# Baseline expression around 10 (log2-like range)
_BASE = _RNG.normal(loc=10.0, scale=0.4, size=(10, 5))
# GENE_A is +2 in Group A to create a detectable signal
_BASE[:5, 0] += 2.0

EXPR = pd.DataFrame(_BASE, index=_SAMPLE_IDS, columns=_GENES)
META = pd.DataFrame(
    {"group": ["A"] * 5 + ["B"] * 5, "tissue": ["tumor"] * 10},
    index=_SAMPLE_IDS,
)


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


def _write_csv(df: pd.DataFrame, suffix: str = ".csv") -> str:
    """Write a DataFrame to a temp file and return the path."""
    tmp = tempfile.NamedTemporaryFile(suffix=suffix, mode="w", delete=False, encoding="utf-8")
    sep = "\t" if suffix in (".tsv", ".txt") else ","
    df.to_csv(tmp.name, sep=sep)
    tmp.close()
    return tmp.name


# ===========================================================================
# load_expression_matrix
# ===========================================================================


class TestLoadExpressionMatrix:
    def test_loads_csv(self):
        path = _write_csv(EXPR, ".csv")
        result = load_expression_matrix(path)
        assert result["n_samples"] == 10
        assert result["n_features"] == 5

    def test_loads_tsv(self):
        path = _write_csv(EXPR, ".tsv")
        result = load_expression_matrix(path)
        assert result["n_samples"] == 10
        assert result["n_features"] == 5

    def test_returns_dataframe(self):
        path = _write_csv(EXPR)
        result = load_expression_matrix(path)
        assert isinstance(result["dataframe"], pd.DataFrame)

    def test_sample_ids_preserved(self):
        path = _write_csv(EXPR)
        result = load_expression_matrix(path)
        assert list(result["dataframe"].index) == _SAMPLE_IDS

    def test_gene_names_preserved(self):
        path = _write_csv(EXPR)
        result = load_expression_matrix(path)
        assert list(result["dataframe"].columns) == _GENES

    def test_sample_id_sample_field(self):
        path = _write_csv(EXPR)
        result = load_expression_matrix(path)
        assert len(result["sample_id_sample"]) <= 5
        assert all(isinstance(s, str) for s in result["sample_id_sample"])

    def test_non_numeric_columns_dropped_with_warning(self):
        df_mixed = EXPR.copy()
        df_mixed["notes"] = "text"
        path = _write_csv(df_mixed)
        result = load_expression_matrix(path)
        assert "notes" not in result["dataframe"].columns
        assert any("non-numeric" in w.lower() for w in result["warnings"])

    def test_file_not_found_raises(self):
        with pytest.raises(FileNotFoundError):
            load_expression_matrix("/nonexistent/path/file.csv")

    def test_transposed_orientation_warns(self):
        # Create a matrix with far more rows than columns (genes as rows)
        # Make it 200 rows (genes) × 5 columns (samples) — should warn
        big_df = pd.DataFrame(
            _RNG.normal(10, 0.5, size=(200, 5)),
            index=[f"GENE_{i}" for i in range(200)],
            columns=[f"S{i:03d}" for i in range(5)],
        )
        path = _write_csv(big_df)
        result = load_expression_matrix(path)
        assert any("transpos" in w.lower() for w in result["warnings"])

    def test_no_warning_for_normal_orientation(self):
        path = _write_csv(EXPR)
        result = load_expression_matrix(path)
        # 10 samples × 5 genes → no orientation warning
        assert not any("transpos" in w.lower() for w in result["warnings"])


# ===========================================================================
# harmonize_expression_and_metadata
# ===========================================================================


class TestHarmonizeExpressionAndMetadata:
    def test_fully_aligned_returns_same_size(self):
        result = harmonize_expression_and_metadata(EXPR, META)
        assert result["n_aligned_samples"] == 10
        assert result["n_expression_only"] == 0
        assert result["n_metadata_only"] == 0

    def test_valid_is_true_when_aligned(self):
        result = harmonize_expression_and_metadata(EXPR, META)
        assert result["valid"] is True

    def test_expression_only_samples_reported(self):
        # Extra samples in expression not in metadata
        expr_extra = pd.concat([EXPR, pd.DataFrame([[9.0] * 5], index=["S999"], columns=_GENES)])
        result = harmonize_expression_and_metadata(expr_extra, META)
        assert result["n_expression_only"] == 1
        assert "S999" in result["expression_only_samples"]
        assert any("no metadata" in w for w in result["warnings"])

    def test_metadata_only_samples_reported(self):
        # Extra samples in metadata not in expression
        meta_extra = pd.concat(
            [META, pd.DataFrame([["A", "tumor"]], index=["S999"], columns=META.columns)]
        )
        result = harmonize_expression_and_metadata(EXPR, meta_extra)
        assert result["n_metadata_only"] == 1
        assert "S999" in result["metadata_only_samples"]

    def test_no_common_samples_invalid(self):
        meta_disjoint = META.copy()
        meta_disjoint.index = [f"X{i:03d}" for i in range(10)]
        result = harmonize_expression_and_metadata(EXPR, meta_disjoint)
        assert result["n_aligned_samples"] == 0
        assert result["valid"] is False
        assert any("no common" in w.lower() for w in result["warnings"])

    def test_sample_id_column_parameter(self):
        # Metadata with sample IDs in a column rather than the index
        meta_col = META.reset_index().rename(columns={"index": "sample_id"})
        result = harmonize_expression_and_metadata(EXPR, meta_col, sample_id_column="sample_id")
        assert result["n_aligned_samples"] == 10

    def test_invalid_sample_id_column_raises(self):
        with pytest.raises(ValueError, match="not found"):
            harmonize_expression_and_metadata(EXPR, META, sample_id_column="bad_col")

    def test_aligned_dataframes_have_same_index(self):
        result = harmonize_expression_and_metadata(EXPR, META)
        assert list(result["expression_df"].index) == list(result["metadata_df"].index)

    def test_partial_overlap_aligns_correctly(self):
        expr_half = EXPR.iloc[:6]  # S001–S006
        meta_half = META.iloc[4:]  # S005–S010 → common: S005, S006
        result = harmonize_expression_and_metadata(expr_half, meta_half)
        assert result["n_aligned_samples"] == 2
        assert result["valid"] is True


# ===========================================================================
# detect_log_scale
# ===========================================================================


class TestDetectLogScale:
    def test_log2_like_data_detected(self):
        # EXPR values are in range ~8–14 (log2-like), non-integer
        result = detect_log_scale(EXPR)
        assert result["likely_log_scale"] is True
        assert result["likely_log2"] is True

    def test_raw_counts_not_log(self):
        counts = pd.DataFrame(
            _RNG.integers(0, 50000, size=(10, 5)).astype(float),
            index=_SAMPLE_IDS,
            columns=_GENES,
        )
        result = detect_log_scale(counts)
        assert result["likely_log_scale"] is False

    def test_negative_values_flagged(self):
        log_ratio = EXPR - 10.0  # centre around 0, creates negatives
        result = detect_log_scale(log_ratio)
        assert result["has_negative_values"] is True
        assert result["likely_log_scale"] is True

    def test_stats_reported(self):
        result = detect_log_scale(EXPR)
        for key in ["value_min", "value_max", "value_median", "value_mean", "fraction_integer"]:
            assert key in result
            assert result[key] is not None

    def test_diagnostic_notes_present(self):
        result = detect_log_scale(EXPR)
        assert len(result["diagnostic_notes"]) > 0

    def test_warnings_always_present(self):
        result = detect_log_scale(EXPR)
        assert isinstance(result["warnings"], list)

    def test_empty_dataframe(self):
        empty = pd.DataFrame()
        result = detect_log_scale(empty)
        assert result["likely_log_scale"] is False
        assert result["value_min"] is None

    def test_fraction_integer_near_zero_for_log(self):
        result = detect_log_scale(EXPR)
        assert result["fraction_integer"] < 0.1

    def test_fraction_integer_near_one_for_counts(self):
        counts = pd.DataFrame(
            _RNG.integers(0, 10000, size=(10, 5)).astype(float), columns=_GENES, index=_SAMPLE_IDS
        )
        result = detect_log_scale(counts)
        assert result["fraction_integer"] > 0.9


# ===========================================================================
# collapse_duplicate_genes
# ===========================================================================


class TestCollapseDuplicateGenes:
    def test_no_duplicates_returns_unchanged(self):
        result = collapse_duplicate_genes(EXPR)
        assert result["n_features_before"] == result["n_features_after"]
        assert result["n_duplicated_genes"] == 0
        assert result["dataframe"].equals(EXPR)

    def test_no_duplicate_warning_message(self):
        result = collapse_duplicate_genes(EXPR)
        assert any("unchanged" in w.lower() for w in result["warnings"])

    def test_mean_collapse(self):
        # Add a duplicate column
        df_dup = pd.concat([EXPR, EXPR[["GENE_A"]].rename(columns={"GENE_A": "GENE_A"})], axis=1)
        result = collapse_duplicate_genes(df_dup, method="mean")
        assert result["n_features_after"] == 5
        assert result["n_duplicated_genes"] == 1
        assert "GENE_A" in result["duplicated_gene_sample"]

    def test_max_collapse(self):
        df_dup = pd.concat([EXPR, EXPR[["GENE_B"]].rename(columns={"GENE_B": "GENE_B"})], axis=1)
        result = collapse_duplicate_genes(df_dup, method="max")
        assert result["n_features_after"] == 5

    def test_most_variable_collapse(self):
        df_dup = pd.concat([EXPR, EXPR[["GENE_C"]]], axis=1)
        result = collapse_duplicate_genes(df_dup, method="most_variable")
        assert result["n_features_after"] == 5

    def test_multiple_duplicated_genes(self):
        df_dup = pd.concat([EXPR, EXPR[["GENE_A", "GENE_B"]]], axis=1)
        result = collapse_duplicate_genes(df_dup, method="mean")
        assert result["n_duplicated_genes"] == 2
        assert result["n_features_after"] == 5

    def test_invalid_method_raises(self):
        with pytest.raises(ValueError, match="method"):
            collapse_duplicate_genes(EXPR, method="invalid")

    def test_output_has_no_duplicates(self):
        df_dup = pd.concat([EXPR, EXPR[["GENE_A", "GENE_B"]]], axis=1)
        result = collapse_duplicate_genes(df_dup, method="mean")
        out_counts = result["dataframe"].columns.value_counts()
        assert (out_counts > 1).sum() == 0

    def test_sample_count_preserved(self):
        df_dup = pd.concat([EXPR, EXPR[["GENE_A"]]], axis=1)
        result = collapse_duplicate_genes(df_dup, method="mean")
        assert result["dataframe"].shape[0] == EXPR.shape[0]


# ===========================================================================
# prepare_gene_level_statistics
# ===========================================================================


class TestPrepareGeneLevelStatistics:
    def test_returns_result_dict(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        assert isinstance(result, dict)
        assert "dataframe" in result

    def test_output_columns(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        df = result["dataframe"]
        for col in ["statistic", "pvalue", "padj", "mean_test", "mean_control", "log2fc_like"]:
            assert col in df.columns, f"Missing column: {col}"

    def test_index_named_gene(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        assert result["dataframe"].index.name == "gene"

    def test_all_genes_in_output(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        assert result["n_genes"] == 5
        assert len(result["dataframe"]) == 5

    def test_sample_counts(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        assert result["n_test_samples"] == 5
        assert result["n_control_samples"] == 5

    def test_gene_a_has_largest_effect(self):
        # GENE_A has +2 units in Group A — should have largest |log2fc_like|
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        df = result["dataframe"]
        gene_a_lfc = abs(df.loc["GENE_A", "log2fc_like"])
        others = df.drop("GENE_A")["log2fc_like"].abs()
        assert gene_a_lfc > others.max()

    def test_gene_a_is_significant(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        df = result["dataframe"]
        assert df.loc["GENE_A", "padj"] < 0.05

    def test_sorted_by_padj(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        df = result["dataframe"]
        assert list(df["padj"]) == sorted(df["padj"].tolist())

    def test_significant_gene_counts(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        assert isinstance(result["significant_genes_05"], int)
        assert result["significant_genes_05"] >= 1  # GENE_A should pass

    def test_warnings_include_log2fc_caveat(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        combined = " ".join(result["warnings"])
        assert "log2fc_like" in combined or "log2" in combined.lower()

    def test_pvalues_in_valid_range(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        df = result["dataframe"]
        assert (df["pvalue"] >= 0).all() and (df["pvalue"] <= 1).all()
        assert (df["padj"] >= 0).all() and (df["padj"] <= 1).all()

    def test_missing_group_column_raises(self):
        with pytest.raises(ValueError, match="group_column"):
            prepare_gene_level_statistics(EXPR, META, "bad_col", "A", "B")

    def test_unknown_test_group_raises(self):
        with pytest.raises(ValueError, match="test_group"):
            prepare_gene_level_statistics(EXPR, META, "group", "X", "B")

    def test_unknown_control_group_raises(self):
        with pytest.raises(ValueError, match="control_group"):
            prepare_gene_level_statistics(EXPR, META, "group", "A", "X")

    def test_too_few_samples_raises(self):
        tiny_expr = EXPR.iloc[:3]  # S001–S003 → only 3 A, 0 B
        tiny_meta = META.iloc[:3]
        with pytest.raises(ValueError):
            prepare_gene_level_statistics(tiny_expr, tiny_meta, "group", "A", "B")

    def test_subset_query_restricts_samples(self):
        meta_tissue = META.copy()
        meta_tissue["tissue"] = ["tumor"] * 8 + ["normal"] * 2
        result = prepare_gene_level_statistics(
            EXPR, meta_tissue, "group", "A", "B", subset_query="tissue == 'tumor'"
        )
        # Only 8 tumor samples, but group A has 5 tumor, group B has 3 tumor
        assert result["n_test_samples"] == 5
        assert result["n_control_samples"] == 3

    def test_invalid_subset_query_raises(self):
        with pytest.raises(ValueError, match="subset_query"):
            prepare_gene_level_statistics(
                EXPR, META, "group", "A", "B", subset_query="not_a_column === broken"
            )

    def test_invalid_method_raises(self):
        with pytest.raises(ValueError, match="method"):
            prepare_gene_level_statistics(EXPR, META, "group", "A", "B", method="deseq2")

    def test_metadata_fields_in_result(self):
        result = prepare_gene_level_statistics(EXPR, META, "group", "A", "B")
        assert result["group_column"] == "group"
        assert result["test_group"] == "A"
        assert result["control_group"] == "B"
        assert result["subset_query"] is None