src/workflows/metadata_validation.py

"""
Metadata validation for sample metadata DataFrames and AnnData objects.

DESIGN PRINCIPLE
----------------
Manifests define column semantics (roles, allowed values, missing-value
meanings, reporting rules).  Tools compute counts dynamically from loaded
data.  This module validates semantics; it does not compare against
hardcoded expected counts.

1. Manifest-aware validators (plain-dict input, no pandas required)
   -------------------------------------------------------------------
   validate_condition_column   Check that a condition column exists and
                               contains only the declared allowed values.
   validate_metadata_semantics Check obs columns against metadata_columns
                               semantic definitions: roles, allowed values,
                               missing-value handling, encoded/decoded
                               consistency, and technical-column guards.
   validate_obs                Run all obs-level checks; return combined dict.

   These accept plain Python dicts derived from AnnData.obs so they can be
   used in lightweight test environments without importing scanpy.

       adata_obs_info = {
           "columns": list(adata.obs.columns),
           "values": {col: adata.obs[col].tolist() for col in adata.obs.columns},
           "n_obs": adata.n_obs,
       }

2. DataFrame-based helpers (pandas input, no AnnData or scanpy required)
   -----------------------------------------------------------------------
   list_valid_groups           Survey all columns for usability as grouping
                               variables (unique value counts, missingness,
                               minimum group size).
   validate_contrast           Validate a DE contrast (test vs control) against
                               a metadata DataFrame, with optional sample
                               subsetting via pandas query syntax.
   apply_subset                Apply a pandas query string to a DataFrame and
                               return clear error messages on failure.

   Typical call order:
       subset_result = apply_subset(obs_df, "tumor_subtype != ''")
       contrast_result = validate_contrast(
           subset_result["dataframe"],
           group_column="tumor_subtype",
           test_group="Classical",
           control_group="Basal",
       )
"""

from __future__ import annotations

from typing import Any

import pandas as pd


def validate_condition_column(
    obs_columns: list[str],
    obs_values: dict[str, list],
    manifest_obs: dict,
) -> dict[str, Any]:
    """
    Check that the condition column declared in the manifest exists in obs
    and contains only the expected values.

    Parameters
    ----------
    obs_columns:
        Column names present in AnnData.obs.
    obs_values:
        Mapping of column_name → list of unique values in that column.
    manifest_obs:
        The 'obs' section of the dataset manifest dict.

    Returns
    -------
    dict with keys:
        valid (bool), column_found (bool), unexpected_values (list),
        missing_values (list), message (str).
    """
    condition_col = manifest_obs.get("condition_column")
    expected_values = set(manifest_obs.get("condition_values", []))

    if not condition_col:
        return {
            "valid": False,
            "column_found": False,
            "unexpected_values": [],
            "missing_values": sorted(expected_values),
            "message": "manifest does not declare obs.condition_column",
        }

    if condition_col not in obs_columns:
        return {
            "valid": False,
            "column_found": False,
            "unexpected_values": [],
            "missing_values": sorted(expected_values),
            "message": f"condition column '{condition_col}' not found in obs",
        }

    actual_values = set(obs_values.get(condition_col, []))

    if expected_values:
        unexpected = sorted(actual_values - expected_values)
        missing = sorted(expected_values - actual_values)
        valid = len(unexpected) == 0
    else:
        # No expected values declared — column presence is enough
        unexpected, missing = [], []
        valid = True

    parts = []
    if unexpected:
        parts.append(f"unexpected values: {unexpected}")
    if missing:
        parts.append(f"missing values: {missing}")
    message = "; ".join(parts) if parts else f"condition column '{condition_col}' is valid"

    return {
        "valid": valid,
        "column_found": True,
        "unexpected_values": unexpected,
        "missing_values": missing,
        "message": message,
    }


def validate_metadata_semantics(
    obs_columns: list[str],
    obs_values: dict[str, list],
    metadata_columns: dict,
) -> dict[str, Any]:
    """
    Validate obs data against the manifest's metadata_columns semantic definitions.

    Checks that declared columns exist, allowed values are present, technical
    columns are not used for grouping, and encoded/decoded column pairs are
    consistent.  Does not compare against hardcoded expected counts.

    Parameters
    ----------
    obs_columns:
        List of column names present in AnnData.obs.
    obs_values:
        Mapping of column_name → list of per-sample values.
    metadata_columns:
        Dict of column_name → MetadataColumnDef or plain dict (from manifest).

    Returns
    -------
    dict with keys:
        valid (bool), checks (dict of per-column results).
    """
    checks: dict[str, Any] = {}

    for col_name, col_def in metadata_columns.items():
        # Support both MetadataColumnDef objects and plain dicts
        if hasattr(col_def, "role"):
            role = col_def.role
            bio_ok = col_def.biological_grouping_allowed
            allowed = set(col_def.allowed_values)
            missing = set(col_def.missing_values)
            source_col = col_def.source_column
            decoded_col = col_def.decoded_column
            interp_warn = col_def.interpretation_warning
            refusal = col_def.refusal_rules
        else:
            role = col_def.get("role", "sample_origin")
            bio_ok = col_def.get("biological_grouping_allowed", True)
            allowed = set(col_def.get("allowed_values") or [])
            missing = set(col_def.get("missing_values") or [])
            source_col = col_def.get("source_column")
            decoded_col = col_def.get("decoded_column")
            interp_warn = col_def.get("interpretation_warning", "")
            refusal = col_def.get("refusal_rules") or []

        issues: list[str] = []
        warnings: list[str] = []

        # Which column to look for in obs (prefer decoded_column)
        check_col = decoded_col or col_name

        # 1. Column presence
        if check_col not in obs_columns:
            if source_col and source_col in obs_columns:
                warnings.append(
                    f"Source column '{source_col}' found but decoded column "
                    f"'{check_col}' is missing — numeric decoder may not have run."
                )
            elif source_col:
                issues.append(
                    f"Neither decoded column '{check_col}' nor source column "
                    f"'{source_col}' found in obs. Available: {obs_columns}"
                )
            else:
                issues.append(f"Column '{check_col}' not found in obs. Available: {obs_columns}")

        # 2. Allowed-value check (only when column is present and allowed_values declared)
        if check_col in obs_columns and allowed:
            actual = set(obs_values.get(check_col, []))
            non_missing = actual - missing
            unexpected = non_missing - allowed
            if unexpected:
                warnings.append(
                    f"Unexpected values in '{check_col}': "
                    f"{sorted(str(v) for v in unexpected)}. "
                    f"Declared allowed values: {sorted(allowed)}"
                )

        # 3. Technical-column guard
        if not bio_ok:
            warnings.append(
                f"Column '{col_name}' has biological_grouping_allowed=False "
                f"(role='{role}'). Do not use for DE or group comparison."
            )

        # 4. Interpretation and refusal rules surfaced as warnings
        if interp_warn:
            warnings.append(f"Interpretation: {interp_warn.strip()}")
        for rule in refusal:
            warnings.append(f"Refusal rule: {rule.strip()}")

        checks[col_name] = {
            "column": check_col,
            "role": role,
            "biological_grouping_allowed": bio_ok,
            "valid": len(issues) == 0,
            "issues": issues,
            "warnings": warnings,
        }

    all_valid = all(c["valid"] for c in checks.values())
    return {"valid": all_valid, "checks": checks}


def validate_obs(adata_obs_info: dict, manifest: dict | Any) -> dict[str, Any]:
    """
    Run obs-level validation checks against the manifest and return a summary.

    Checks condition column existence and metadata column semantics.
    Does NOT compare against hardcoded expected counts — counts are computed
    dynamically by dataset_count_metadata_values.

    Parameters
    ----------
    adata_obs_info:
        Dict describing the loaded AnnData.obs:
            columns (list[str])       — column names in obs
            values  (dict[str, list]) — column → per-sample value list
            n_obs   (int)             — total number of observations

    manifest:
        Parsed manifest dict or DatasetManifest instance.

    Returns
    -------
    dict with keys: valid (bool), checks (dict of check_name → result).
    """
    # Resolve manifest to usable dicts
    if hasattr(manifest, "metadata_columns"):
        # DatasetManifest instance
        manifest_obs = {
            "condition_column": manifest.group_columns[0] if manifest.group_columns else None,
            "condition_values": [],
        }
        metadata_cols = manifest.metadata_columns
    else:
        manifest_obs = manifest.get("obs", {})
        metadata_cols = manifest.get("metadata_columns", {})

    obs_columns = adata_obs_info.get("columns", [])
    obs_values = adata_obs_info.get("values", {})

    checks: dict[str, Any] = {}

    # Check 1: condition column presence and allowed values
    if manifest_obs.get("condition_column"):
        checks["condition_column"] = validate_condition_column(
            obs_columns, obs_values, manifest_obs
        )

    # Check 2: metadata semantics (replaces validate_sample_counts)
    if metadata_cols:
        checks["metadata_semantics"] = validate_metadata_semantics(
            obs_columns, obs_values, metadata_cols
        )

    all_valid = all(c.get("valid", False) for c in checks.values())
    return {"valid": all_valid, "checks": checks}


# ---------------------------------------------------------------------------
# DataFrame-based helpers
# ---------------------------------------------------------------------------


def list_valid_groups(
    metadata_df: pd.DataFrame,
    min_count: int = 3,
) -> dict[str, Any]:
    """
    Survey all columns of a metadata DataFrame for usability as grouping variables.

    A column is considered usable when it has at least 2 distinct non-null values
    and at least 2 groups each with ≥min_count samples.  Columns where every
    value is unique (likely sample IDs) are flagged explicitly.

    Parameters
    ----------
    metadata_df:
        Sample-level metadata DataFrame (e.g. AnnData.obs or a parsed GEO table).
        Rows are samples; columns are metadata fields.
    min_count:
        Minimum number of samples required per group for a column to be considered
        usable.  Default 3 (minimum for most statistical tests).

    Returns
    -------
    dict with keys:
        n_samples         (int)        — number of rows in the DataFrame.
        n_columns_checked (int)        — number of columns surveyed.
        usable_columns    (list[str])  — columns that pass the usability criteria.
        columns           (dict)       — per-column details:
            n_unique, n_missing, missing_rate, value_counts, unique_values_sample,
            values_truncated, is_usable, usability_reason.
    """
    n_rows = len(metadata_df)
    columns_info: dict[str, Any] = {}
    usable_columns: list[str] = []

    for col in metadata_df.columns:
        series = metadata_df[col]
        n_missing = int(series.isna().sum())
        non_null = series.dropna()
        n_unique = int(non_null.nunique())
        value_counts_raw = non_null.value_counts()
        value_counts = {str(k): int(v) for k, v in value_counts_raw.items()}

        # Unique values sample — cap at 20 to keep return value manageable
        unique_vals_sample = sorted([str(v) for v in non_null.unique()[:20]])
        values_truncated = n_unique > 20

        # Determine usability
        is_usable = False
        if n_rows == 0:
            usability_reason = "DataFrame is empty"
        elif n_unique == 0:
            usability_reason = "All values are missing"
        elif n_unique == 1:
            val = str(non_null.iloc[0])
            usability_reason = f"Only one unique value ('{val}') — cannot form a contrast"
        elif n_unique == n_rows and n_rows > 10:
            usability_reason = (
                f"All {n_unique} values are unique — likely a sample ID column, "
                "not usable for grouping"
            )
        else:
            qualifying = [v for v, c in value_counts_raw.items() if c >= min_count]
            small = [str(v) for v, c in value_counts_raw.items() if c < min_count]
            if len(qualifying) < 2:
                usability_reason = (
                    f"Fewer than 2 groups have ≥{min_count} samples "
                    f"(qualifying groups: {[str(q) for q in qualifying]})"
                )
            else:
                is_usable = True
                usability_reason = f"{len(qualifying)} group(s) with ≥{min_count} samples"
                if small:
                    usability_reason += (
                        f"; {len(small)} group(s) below minimum ({small}) "
                        "are present but excluded from contrast"
                    )

        columns_info[col] = {
            "n_unique": n_unique,
            "n_missing": n_missing,
            "missing_rate": round(n_missing / n_rows, 4) if n_rows > 0 else 0.0,
            "value_counts": value_counts,
            "unique_values_sample": unique_vals_sample,
            "values_truncated": values_truncated,
            "is_usable": is_usable,
            "usability_reason": usability_reason,
        }
        if is_usable:
            usable_columns.append(col)

    return {
        "n_samples": n_rows,
        "n_columns_checked": len(metadata_df.columns),
        "usable_columns": usable_columns,
        "columns": columns_info,
    }


def apply_subset(
    metadata_df: pd.DataFrame,
    subset_query: str | None,
) -> dict[str, Any]:
    """
    Apply a pandas query string to a metadata DataFrame.

    Returns a dict rather than raising so that callers can handle failures
    gracefully and surface a clear message to the user or agent.

    Parameters
    ----------
    metadata_df:
        The metadata DataFrame to filter.
    subset_query:
        A pandas query string, e.g. ``"tissue == 'pancreas' and age > 50"``.
        None or empty string returns the DataFrame unchanged.

    Returns
    -------
    dict with keys:
        success    (bool)                  — True if query applied without error.
        query      (str | None)            — the query string used.
        n_before   (int)                   — rows before filtering.
        n_after    (int | None)            — rows after filtering; None on failure.
        n_dropped  (int | None)            — rows removed; None on failure.
        dataframe  (pd.DataFrame | None)   — filtered DataFrame; None on failure.
                                             Note: not JSON-serialisable.
        error      (str | None)            — error description; None on success.
    """
    n_before = len(metadata_df)

    if not subset_query or not str(subset_query).strip():
        return {
            "success": True,
            "query": subset_query,
            "n_before": n_before,
            "n_after": n_before,
            "n_dropped": 0,
            "dataframe": metadata_df,
            "error": None,
        }

    try:
        subset = metadata_df.query(subset_query)
        n_after = len(subset)
        return {
            "success": True,
            "query": subset_query,
            "n_before": n_before,
            "n_after": n_after,
            "n_dropped": n_before - n_after,
            "dataframe": subset,
            "error": None,
        }
    except Exception as exc:
        return {
            "success": False,
            "query": subset_query,
            "n_before": n_before,
            "n_after": None,
            "n_dropped": None,
            "dataframe": None,
            "error": (
                f"Query failed: {exc}. "
                "Use pandas query syntax, e.g. "
                "\"tumor_subtype == 'Classical'\" or "
                "\"age > 50 and tissue == 'pancreas'\"."
            ),
        }


def validate_contrast(
    metadata_df: pd.DataFrame,
    group_column: str,
    test_group: str,
    control_group: str,
    subset_query: str | None = None,
    min_samples_per_group: int = 3,
) -> dict[str, Any]:
    """
    Validate a differential expression contrast against a metadata DataFrame.

    Optionally filters the DataFrame first via a pandas query (subset_query)
    before checking group membership and sample counts.

    This function ONLY validates metadata and sample selection.  It does not
    run any expression analysis.  Use the returned test_sample_ids and
    control_sample_ids to slice your expression matrix.

    Parameters
    ----------
    metadata_df:
        Sample-level metadata DataFrame (e.g. AnnData.obs).
        Index values are used as sample IDs in the return value.
    group_column:
        Column in metadata_df containing the group labels.
    test_group:
        Label of the test / foreground group.
    control_group:
        Label of the reference / background group.
    subset_query:
        Optional pandas query string applied before contrast validation,
        e.g. ``"tissue == 'tumor'"`` to restrict to tumour samples only.
    min_samples_per_group:
        Minimum number of samples required in each group.  Default 3.

    Returns
    -------
    dict with keys:
        valid               (bool)       — True if contrast is usable for DE.
        reason              (str | None) — explanation if valid=False, else None.
        group_column        (str)
        test_group          (str)
        control_group       (str)
        n_test              (int | None) — samples in test group; None on early failure.
        n_control           (int | None) — samples in control group; None on early failure.
        n_total_selected    (int | None) — n_test + n_control; None on early failure.
        test_sample_ids     (list[str])  — index values of test-group samples.
        control_sample_ids  (list[str])  — index values of control-group samples.
        subset_query        (str | None) — the subset_query used, if any.
        warnings            (list[str])  — non-fatal issues (e.g. small group sizes,
                                           excluded samples).
    """
    warnings: list[str] = []

    def _early_return(reason: str) -> dict[str, Any]:
        return {
            "valid": False,
            "reason": reason,
            "group_column": group_column,
            "test_group": test_group,
            "control_group": control_group,
            "n_test": None,
            "n_control": None,
            "n_total_selected": None,
            "test_sample_ids": [],
            "control_sample_ids": [],
            "subset_query": subset_query,
            "warnings": warnings,
        }

    # ── Step 1: Apply optional subset ────────────────────────────────────
    working_df = metadata_df
    if subset_query:
        sub = apply_subset(metadata_df, subset_query)
        if not sub["success"]:
            return _early_return(f"Subset query failed: {sub['error']}")
        working_df = sub["dataframe"]
        if len(working_df) == 0:
            return _early_return(f"Subset query '{subset_query}' selected 0 samples")

    # ── Step 2: group_column existence ───────────────────────────────────
    if group_column not in working_df.columns:
        available = sorted(working_df.columns.tolist())
        return _early_return(
            f"group_column '{group_column}' not found. Available columns: {available}"
        )

    # ── Step 3: Group existence ──────────────────────────────────────────
    available_groups = working_df[group_column].dropna().unique().tolist()
    available_str = sorted(str(g) for g in available_groups)

    if test_group == control_group:
        return _early_return("test_group and control_group are the same")

    if test_group not in available_groups:
        return _early_return(
            f"test_group '{test_group}' not found in column '{group_column}'. "
            f"Available groups: {available_str}"
        )
    if control_group not in available_groups:
        return _early_return(
            f"control_group '{control_group}' not found in column '{group_column}'. "
            f"Available groups: {available_str}"
        )

    # ── Step 4: Count and collect sample IDs ─────────────────────────────
    test_mask = working_df[group_column] == test_group
    ctrl_mask = working_df[group_column] == control_group

    test_ids = [str(i) for i in working_df[test_mask].index.tolist()]
    ctrl_ids = [str(i) for i in working_df[ctrl_mask].index.tolist()]
    n_test = len(test_ids)
    n_ctrl = len(ctrl_ids)

    # ── Step 5: Minimum sample count ─────────────────────────────────────
    if n_test < min_samples_per_group:
        return _early_return(
            f"test_group '{test_group}' has {n_test} sample(s), "
            f"below minimum {min_samples_per_group}"
        )
    if n_ctrl < min_samples_per_group:
        return _early_return(
            f"control_group '{control_group}' has {n_ctrl} sample(s), "
            f"below minimum {min_samples_per_group}"
        )

    # ── Step 6: Non-fatal warnings ───────────────────────────────────────
    if n_test < 5:
        warnings.append(
            f"test_group '{test_group}' has only {n_test} samples — "
            "statistical power may be limited"
        )
    if n_ctrl < 5:
        warnings.append(
            f"control_group '{control_group}' has only {n_ctrl} samples — "
            "statistical power may be limited"
        )

    n_excluded = len(working_df) - n_test - n_ctrl
    if n_excluded > 0:
        other = [str(g) for g in available_groups if g not in (test_group, control_group)]
        warnings.append(
            f"{n_excluded} sample(s) in other group(s) will be excluded from DE "
            f"(groups: {other[:5]}{'...' if len(other) > 5 else ''})"
        )

    return {
        "valid": True,
        "reason": None,
        "group_column": group_column,
        "test_group": test_group,
        "control_group": control_group,
        "n_test": n_test,
        "n_control": n_ctrl,
        "n_total_selected": n_test + n_ctrl,
        "test_sample_ids": test_ids,
        "control_sample_ids": ctrl_ids,
        "subset_query": subset_query,
        "warnings": warnings,
    }


def subset_and_require_group(
    metadata_df: pd.DataFrame, subset_query: str | None, group_column: str
) -> pd.DataFrame:
    """Apply an optional subset query and require group_column to be present.

    Shared by the activity-stats and microarray workflows. Raises ValueError on a
    failed or empty query, or a missing column; returns the (possibly subset) frame.
    """
    working_meta = metadata_df
    if subset_query:
        try:
            working_meta = metadata_df.query(subset_query)
        except Exception as exc:
            raise ValueError(f"subset_query '{subset_query}' failed: {exc}") from exc
        if len(working_meta) == 0:
            raise ValueError(f"subset_query '{subset_query}' selected 0 samples.")
    if group_column not in working_meta.columns:
        raise ValueError(
            f"group_column '{group_column}' not found in metadata. "
            f"Available: {list(working_meta.columns)}"
        )
    return working_meta