src/tools/dataset_tools/metadata.py

"""dataset planning + metadata-inspection tools."""
# ruff: noqa: F403, F405, F722  (star re-export from ._base; F722 = Annotated false-positive)

from ._base import *  # noqa: F401
from ._base import (  # noqa: F401
    _build_loading_plan,
    _classify_metadata_values,
    _col_semantics,
    _detect_intent,
    _flatten_single_column_summary,
    _load_metadata_df,
    _workflow_for_intent,
)


@dataset_mcp.tool
def dataset_plan_analysis(
    dataset_id: Annotated[
        str,
        "The dataset identifier (e.g. 'gse71729_moffitt'). "
        "Call dataset_list_available() first if unsure. "
        "Used to populate recommended tool arguments from the manifest.",
    ],
    user_question: Annotated[
        str,
        "The user's analysis question in plain language. Examples: "
        "'Compare Classical vs Basal tumors for TF activity', "
        "'Which samples have high TGFb pathway activity?', "
        "'Is HIF1A activity associated with survival?'",
    ],
) -> dict:
    """
    Return a recommended workflow for a dataset and a plain-language question.

    Call this when:
    - The user states a biological question and you need to choose which tools
      to run and in what order.
    - You want to surface required inputs, assumptions, and unimplemented steps
      before starting.
    - You need to flag that part of a requested analysis is not yet available.

    This tool does NOT run any analysis — it returns a plan only.
    Execute the plan by calling the listed tools in order.

    Detected intent classes
    -----------------------
    metadata_summary     Keywords: how many, count, distribution, breakdown,
                         value counts, what groups, what subtypes
                         → dataset_interpret_metadata →
                           dataset_count_metadata_values

    compare_groups       Keywords: compare, differ, vs, versus, higher/lower in,
                         upregulated, differential
                         → dataset_describe → load → validate_contrast →
                           decoupler_differential_expression → enrichment tools

    score_samples        Keywords: activity, score, which samples, sample-level,
                         per sample, pathway activity, TF activity
                         → dataset_score_bulk_samples →
                           dataset_compare_activity_by_group (optional)

    survival             Keywords: survival, prognosis, KM, Cox, time to event
                         → dataset_score_bulk_samples → survival (not implemented)

    correlate_continuous Keywords: correlate, associated with, continuous, covariate
                         → dataset_score_bulk_samples →
                           associate_activity_with_covariate (not implemented)

    unknown              No keyword match → generic orientation steps + warning
    """
    manifest: Any = None
    manifest_warning: str | None = None
    try:
        manifest = load_manifest(dataset_id)
    except KeyError:
        available = [d["dataset_id"] for d in list_available_datasets()]
        manifest_warning = (
            f"Dataset '{dataset_id}' not found in registry. "
            f"Available datasets: {available}. "
            "Plan uses generic placeholders (no manifest-specific args)."
        )

    intent, confidence, matched_kw = _detect_intent(user_question)
    workflow = _workflow_for_intent(intent, dataset_id, manifest)

    if manifest_warning:
        workflow["warnings"] = [manifest_warning] + workflow["warnings"]

    return {
        "dataset_id": dataset_id,
        "user_question": user_question,
        "detected_intent": intent,
        "confidence": confidence,
        "matched_keywords": matched_kw,
        "recommended_tools": workflow["steps"],
        "required_inputs": workflow["required_inputs"],
        "assumptions": workflow["assumptions"],
        "warnings": workflow["warnings"],
        "refusal_conditions": workflow["refusal_conditions"],
    }


# ---------------------------------------------------------------------------
# Metadata interpretation and dynamic counting tools
# ---------------------------------------------------------------------------


@dataset_mcp.tool
def dataset_interpret_metadata(
    adata_path: Annotated[
        str,
        "Path to a loaded h5ad file containing the dataset's obs metadata.",
    ],
    dataset_id: Annotated[
        str,
        "Dataset identifier from the registry (e.g. 'gse71729_moffitt').",
    ],
) -> dict:
    """
    Return the semantic interpretation guide for each metadata column declared
    in the manifest.

    Call this when:
    - A user asks what columns or groups are in the dataset.
    - You need to understand what values are biologically meaningful vs missing.
    - You need to know which columns may be used for grouping and which are
      technical-only.
    - You need to surface refusal rules before proceeding with analysis.

    This tool reads column semantics from the manifest — it does NOT compute
    counts. Use dataset_count_metadata_values to compute actual value counts
    from the loaded data.
    """
    try:
        manifest = load_manifest(dataset_id)
    except KeyError:
        return {"error": f"Dataset '{dataset_id}' not found in registry."}

    if not manifest.metadata_columns:
        return {
            "dataset_id": dataset_id,
            "message": "No metadata_columns defined in this manifest. "
            "Use dataset_list_valid_sample_groups to survey the data directly.",
            "metadata_columns": {},
        }

    columns_guide = {}
    for col_name, col_def in manifest.metadata_columns.items():
        if hasattr(col_def, "to_dict"):
            defn = col_def.to_dict()
        else:
            defn = dict(col_def)
        columns_guide[col_name] = {
            "role": defn.get("role"),
            "biological_grouping_allowed": defn.get("biological_grouping_allowed", True),
            "check_column": defn.get("decoded_column") or col_name,
            "source_column": defn.get("source_column"),
            "allowed_values": defn.get("allowed_values", []),
            "missing_values": defn.get("missing_values", []),
            "empty_value_meaning": defn.get("empty_value_meaning", ""),
            "code_map": defn.get("code_map", {}),
            "interpretation_warning": defn.get("interpretation_warning", ""),
            "reporting_rules": defn.get("reporting_rules", []),
            "refusal_rules": defn.get("refusal_rules", []),
        }

    return {
        "dataset_id": dataset_id,
        "adata_path": adata_path,
        "dataset_reporting_rules": manifest.reporting_rules,
        "dataset_refusal_rules": manifest.refusal_rules,
        "metadata_columns": columns_guide,
        "next_step": (
            "Call dataset_count_metadata_values to compute actual value counts "
            "from the loaded data applying these semantics."
        ),
    }


@dataset_mcp.tool
def dataset_count_metadata_values(
    adata_path: Annotated[
        str,
        "Path to a loaded h5ad file. Counts are computed from this file's obs.",
    ],
    dataset_id: Annotated[
        str,
        "Dataset identifier used to load manifest semantics.",
    ],
    column: Annotated[
        str | None,
        "Specific column name to count. If None, counts all metadata_columns "
        "declared in the manifest.",
    ] = None,
) -> dict:
    """
    Compute value counts dynamically from the loaded AnnData obs, applying the
    manifest's missing-value semantics to classify each value as annotated,
    missing, or unexpected.

    Call this when:
    - A user asks how many Classical or Basal samples there are.
    - You need to show a subtype breakdown or group distribution.
    - You need to know how many samples lack annotation before subsetting.

    Counts are computed from the actual data — they are NOT read from the
    manifest. The manifest provides semantics (which values are missing, which
    are allowed) to interpret the raw counts.
    """
    from collections import Counter as _Counter
    from pathlib import Path

    import scanpy as sc

    from src.core.data_io import resolve_to_local_path

    # Resolve local paths AND remote URLs (e.g. private pdac-research-data h5ads)
    # to a readable local file, authenticating via HF_TOKEN for private
    # huggingface.co /resolve/ URLs. Without this, an https URL was treated as a
    # local path (Path.exists() == False) and wrongly reported "File not found".
    try:
        local_path, is_temp = resolve_to_local_path(adata_path)
    except Exception as e:
        return {"error": f"Could not resolve {adata_path}: {e}"}

    if not is_temp and not Path(local_path).exists():
        return {"error": f"File not found: {adata_path}"}

    try:
        adata = sc.read_h5ad(local_path)
    except Exception as e:
        return {"error": f"Could not load h5ad: {e}"}
    finally:
        if is_temp:
            try:
                Path(local_path).unlink()
            except OSError:
                pass

    try:
        manifest = load_manifest(dataset_id)
    except KeyError:
        return {"error": f"Dataset '{dataset_id}' not found in registry."}

    # Determine which columns to count
    if manifest.metadata_columns:
        cols_to_count = (
            {column: manifest.metadata_columns[column]}
            if column and column in manifest.metadata_columns
            else manifest.metadata_columns
        )
    else:
        # No manifest semantics — fall back to counting all obs columns
        cols_to_count = {c: {} for c in adata.obs.columns}

    summaries = {}
    for col_name, col_def in cols_to_count.items():
        sem = _col_semantics(col_def, col_name)
        missing_set = sem["missing_set"]
        allowed_set = sem["allowed_set"]
        bio_ok = sem["bio_ok"]
        check_col = sem["check_col"]
        interp_warn = sem["interp_warn"]
        col_refusal_rules = sem["refusal_rules"]

        if check_col not in adata.obs.columns:
            summaries[col_name] = {
                "error": f"Column '{check_col}' not found in obs. "
                f"Available: {list(adata.obs.columns)}",
            }
            continue

        raw_counts = _Counter(adata.obs[check_col].astype(str).tolist())
        annotated, missing, unexpected = _classify_metadata_values(
            raw_counts, missing_set, allowed_set
        )

        total_annotated = sum(annotated.values())
        total_missing = sum(missing.values())
        total_unexpected = sum(unexpected.values())
        total = sum(raw_counts.values())

        summaries[col_name] = {
            "column": check_col,
            "biological_grouping_allowed": bio_ok,
            "total": total,
            "total_annotated": total_annotated,
            "total_missing": total_missing,
            "total_unexpected": total_unexpected,
            "annotated": dict(sorted(annotated.items(), key=lambda x: -x[1])),
            "missing": dict(sorted(missing.items(), key=lambda x: -x[1])),
            "unexpected": dict(sorted(unexpected.items(), key=lambda x: -x[1])),
            "interpretation_warning": interp_warn.strip() if interp_warn else None,
            "prohibited_inferences": col_refusal_rules,
        }

    result: dict = {
        "dataset_id": dataset_id,
        "adata_path": adata_path,
        "n_total_samples": adata.n_obs,
        "column_summaries": summaries,
        "dataset_refusal_rules": list(manifest.refusal_rules or []),
        "next_step": (
            "Use subset_query from the manifest's default_contrasts to isolate "
            "the relevant sample group before running differential expression."
        ),
    }

    # When a single column was requested, add a flattened top-level summary
    # so callers don't have to navigate column_summaries[col_name].
    if column and column in summaries:
        result.update(
            _flatten_single_column_summary(
                summaries[column],
                cols_to_count.get(column, {}),
                manifest,
                column,
                adata_path,
                dataset_id,
            )
        )

    return result


@dataset_mcp.tool
def dataset_crosstab_metadata_values(
    adata_path: Annotated[
        str,
        "Path to a loaded h5ad file.",
    ],
    dataset_id: Annotated[
        str,
        "Dataset identifier used to load manifest semantics.",
    ],
    row_column: Annotated[
        str,
        "Obs column to use as crosstab rows (e.g. 'tumor_subtype').",
    ],
    col_column: Annotated[
        str,
        "Obs column to use as crosstab columns (e.g. 'stroma_subtype').",
    ],
    exclude_missing: Annotated[
        bool,
        "If True, drop rows/columns whose values are in the manifest's "
        "missing_values list. Default True.",
    ] = True,
) -> dict:
    """
    Cross-tabulate two metadata columns from the loaded AnnData, applying
    manifest missing-value semantics to optionally exclude unannotated samples.

    Call this when:
    - A user asks how tumor subtype and stroma subtype co-occur.
    - You need a two-way breakdown of sample groups.
    - You want to check whether two columns are confounded.
    """
    from pathlib import Path

    import pandas as pd
    import scanpy as sc

    from src.core.data_io import resolve_to_local_path

    # Resolve local paths AND remote URLs (incl. private HF repos via HF_TOKEN).
    try:
        local_path, is_temp = resolve_to_local_path(adata_path)
    except Exception as e:
        return {"error": f"Could not resolve {adata_path}: {e}"}

    if not is_temp and not Path(local_path).exists():
        return {"error": f"File not found: {adata_path}"}

    try:
        adata = sc.read_h5ad(local_path)
    except Exception as e:
        return {"error": f"Could not load h5ad: {e}"}
    finally:
        if is_temp:
            try:
                Path(local_path).unlink()
            except OSError:
                pass

    try:
        manifest = load_manifest(dataset_id)
    except KeyError:
        return {"error": f"Dataset '{dataset_id}' not found in registry."}

    def _resolve_col(col_name: str) -> tuple[str, set]:
        """Return (actual_obs_col, missing_values_set)."""
        col_def = manifest.metadata_columns.get(col_name, {})
        if hasattr(col_def, "decoded_column"):
            obs_col = col_def.decoded_column or col_name
            missing = set(col_def.missing_values)
        else:
            obs_col = col_def.get("decoded_column") or col_name
            missing = set(col_def.get("missing_values") or [])
        return obs_col, missing

    row_obs, row_missing = _resolve_col(row_column)
    col_obs, col_missing = _resolve_col(col_column)

    for obs_col in (row_obs, col_obs):
        if obs_col not in adata.obs.columns:
            return {
                "error": f"Column '{obs_col}' not found in obs. "
                f"Available: {list(adata.obs.columns)}"
            }

    obs_df = adata.obs[[row_obs, col_obs]].copy()
    obs_df[row_obs] = obs_df[row_obs].astype(str)
    obs_df[col_obs] = obs_df[col_obs].astype(str)

    if exclude_missing:
        all_missing = (row_missing | {"nan", "None", ""}) | (col_missing | {"nan", "None", ""})
        obs_df = obs_df[~obs_df[row_obs].isin(all_missing) & ~obs_df[col_obs].isin(all_missing)]

    crosstab = pd.crosstab(obs_df[row_obs], obs_df[col_obs])

    return {
        "dataset_id": dataset_id,
        "row_column": row_obs,
        "col_column": col_obs,
        "exclude_missing": exclude_missing,
        "n_samples_included": int(len(obs_df)),
        "n_samples_total": adata.n_obs,
        "crosstab": {
            str(row): {str(col): int(val) for col, val in row_data.items()}
            for row, row_data in crosstab.to_dict(orient="index").items()
        },
        "row_totals": {str(k): int(v) for k, v in crosstab.sum(axis=1).items()},
        "col_totals": {str(k): int(v) for k, v in crosstab.sum(axis=0).items()},
    }


@dataset_mcp.tool
def dataset_validate_manifest_against_data(
    dataset_id: Annotated[
        str,
        "The registered dataset_id to validate (e.g. 'gse71729_moffitt'). "
        "Call dataset_list_available() first if unsure.",
    ],
    adata_path: Annotated[
        str,
        "Path to a loaded .h5ad file for this dataset. "
        "Load with decoupler_load_geo_series_matrix (or equivalent) first.",
    ],
) -> dict:
    """
    Validate a loaded dataset against its manifest declaration.

    Call this tool:
    - After loading a dataset for the first time, before running any analysis.
    - When adding a new dataset to confirm the manifest is correct.
    - When you see unexpected metadata values or data type errors during analysis.

    Five checks are run:
      data_level        — detected expression type vs. manifest declaration.
      feature_id_type   — var.index format (gene symbols, Ensembl, probe IDs).
      metadata_columns  — declared columns present; values within allowed set.
      group_columns     — each group column exists with ≥2 usable groups.
      default_contrasts — each contrast has ≥3 samples per group after subset.

    Interpreting the result:
    - overall_valid=True, n_warnings=0  → proceed directly to analysis.
    - overall_valid=True, n_warnings>0  → safe to proceed; review warnings first.
    - overall_valid=False               → fix errors before running DE or enrichment.

    Next steps:
    - If valid: call dataset_plan_analysis(dataset_id, user_question) to start.
    - If errors: read the errors list and update the manifest or re-load the data.
    - If metadata decoder warning: re-run decoupler_load_geo_series_matrix
      with correct condition_column to trigger numeric decoding.
    """
    from pathlib import Path as _Path

    import numpy as np
    import scanpy as sc

    from src.core.data_io import resolve_to_local_path
    from src.workflows.manifest_data_validation import validate_manifest_against_data

    # Resolve local paths AND remote URLs (incl. private HF repos via HF_TOKEN).
    try:
        local_path, _is_temp = resolve_to_local_path(adata_path)
    except Exception as e:
        return {"error": f"Could not resolve {adata_path}: {e}"}

    try:
        adata = sc.read_h5ad(local_path)
    except Exception as e:
        return {"error": f"Could not load h5ad at '{adata_path}': {e}"}
    finally:
        if _is_temp:
            try:
                _Path(local_path).unlink()
            except OSError:
                pass

    try:
        manifest = load_manifest(dataset_id)
    except KeyError as e:
        return {"error": str(e)}

    # Sample expression matrix (cap at 50 000 values for speed)
    X = adata.X
    if hasattr(X, "toarray"):
        X = X.toarray()
    X_flat = X.flatten()
    if len(X_flat) > 50_000:
        rng = np.random.default_rng(0)
        X_flat = rng.choice(X_flat, size=50_000, replace=False)

    var_index = list(adata.var.index)
    obs_df = adata.obs.copy()

    return validate_manifest_against_data(X_flat, var_index, obs_df, manifest)