src/tools/analytics/relationship.py

"""analyze_correlation — correlation among numeric columns (KM-608).

An analytical "family" tool: in ONE call it measures how strongly numeric
columns move together. Returns the full correlation matrix plus a list of
column pairs ranked by strength. Answers questions like "does price relate to
units sold?".

STATUS: compute layer only — the function takes an already-materialized
DataFrame. The wrapper layer (fetching data from the catalog via source_id,
the ToolOutput envelope, ToolSpec registration) is added once the Planner
seam (KM-418) is settled. Keeping compute separate from data-fetching makes
this function easy to unit-test in isolation and stable when wrapped.
"""

from __future__ import annotations

import math

import pandas as pd

from src.tools.analytics.descriptive import ColumnNotFoundError

# Correlation methods supported by pandas .corr().
SUPPORTED_METHODS = ("pearson", "spearman", "kendall")


class InvalidMethodError(ValueError):
    """The requested method is not supported (maps to error_code INVALID_METHOD)."""


class NonNumericColumnError(ValueError):
    """A requested column is not numeric (maps to error_code NON_NUMERIC_COLUMN)."""


class NotEnoughColumnsError(ValueError):
    """Correlation needs at least two numeric columns (maps to NOT_ENOUGH_COLUMNS)."""


def _clean(value: object) -> float | None:
    """Cast to plain float; NaN (e.g. a zero-variance column) -> None."""
    if value is None:
        return None
    f = float(value)  # type: ignore[arg-type]
    return None if math.isnan(f) else f


# Prompt-style description read by the Planner to decide WHEN to pick this tool.
# Final destination is ToolSpec.description once the wrapper layer is built.
DESCRIPTION = """\
Summary: Pairwise correlation across numeric columns (pearson, spearman, or \
kendall). Returns a correlation matrix plus the strongest pairs ranked by \
absolute strength.

USE WHEN the question is about relationship or association between numeric \
variables. Trigger words: "correlation" (korelasi), "related/relationship" \
(hubungan/keterkaitan), "does X affect Y", "move together".

DON'T USE WHEN:
  - it implies causation — correlation is not causality; stay descriptive
  - it compares two groups of one metric -> analyze_comparison
  - it summarizes a single column -> analyze_descriptive

Example questions:
  - "is there a correlation between price and quantity sold?"
  - "which variables are most related to revenue?"
  - "do age and spending move together?"
  - "show the correlation matrix for the numeric columns"
"""


def analyze_correlation(
    df: pd.DataFrame,
    column_ids: list[str] | None = None,
    method: str = "pearson",
) -> dict[str, object]:
    """Pairwise correlation across numeric columns.

    Args:
        df: already-materialized data (in the real system the wrapper fetches
            this from a source_id).
        column_ids: numeric columns to correlate. If None, every numeric
            column in df is used.
        method: "pearson" (linear), "spearman" (rank), or "kendall".

    Returns:
        dict with:
            method   — echo of the chosen method
            columns  — the numeric columns actually correlated
            matrix   — { col: { col: corr|None } } full square matrix
            pairs    — [{"a", "b", "corr"}] unique pairs, strongest |corr| first

    Raises:
        InvalidMethodError: if method is unknown.
        ColumnNotFoundError: if an explicit column is absent.
        NonNumericColumnError: if an explicit column is not numeric.
        NotEnoughColumnsError: if fewer than two numeric columns remain.
    """
    if method not in SUPPORTED_METHODS:
        raise InvalidMethodError(
            f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}"
        )

    if column_ids is None:
        cols = [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]
    else:
        missing = [c for c in column_ids if c not in df.columns]
        if missing:
            raise ColumnNotFoundError(f"columns not found: {missing}")
        non_numeric = [
            c for c in column_ids if not pd.api.types.is_numeric_dtype(df[c])
        ]
        if non_numeric:
            raise NonNumericColumnError(f"columns are not numeric: {non_numeric}")
        cols = list(column_ids)

    if len(cols) < 2:
        raise NotEnoughColumnsError(
            f"need >= 2 numeric columns, got {len(cols)}: {cols}"
        )

    corr = df[cols].corr(method=method)
    matrix = {a: {b: _clean(corr.loc[a, b]) for b in cols} for a in cols}

    pairs = []
    for i in range(len(cols)):
        for j in range(i + 1, len(cols)):
            val = _clean(corr.iloc[i, j])
            if val is not None:
                pairs.append({"a": cols[i], "b": cols[j], "corr": val})
    pairs.sort(key=lambda p: abs(p["corr"]), reverse=True)

    return {"method": method, "columns": cols, "matrix": matrix, "pairs": pairs}