"""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
from typing import Any

import pandas as pd
from pydantic import BaseModel, Field

from src.tools.analytics.descriptive import ColumnNotFoundError

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

# A pair is weak evidence below this many complete rows or at/above this p-value.
# Audit P0-8: an r of 0.953 came from only 8 monthly totals and read as strong.
_WEAK_MIN_N = 10
_ALPHA = 0.05


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)
    return None if math.isnan(f) else f


def _p_value(a: pd.Series, b: pd.Series, method: str) -> float | None:
    """Two-sided p-value of the pair's correlation on its complete rows; None if < 3."""
    from scipy import stats  # lazy — keep module import light

    if a.shape[0] < 3:
        return None
    test = {"pearson": stats.pearsonr, "spearman": stats.spearmanr,
            "kendall": stats.kendalltau}[method]
    return _clean(test(a, b).pvalue)


def _is_weak(n: int, p_value: float | None) -> bool:
    return n < _WEAK_MIN_N or p_value is None or p_value >= _ALPHA


def _quality(top: dict[str, Any] | None) -> dict[str, object]:
    """Reserved `quality` block (conventions §3) for the strongest pair; CK7 surfaces it."""
    if top is None:
        return {"metric": "p_value", "value": None, "floor": _ALPHA, "weak": True,
                "note": "no pair had a defined correlation"}
    n, p = top["n"], top["p_value"]
    reasons = []
    if n < _WEAK_MIN_N:
        reasons.append(f"only {n} rows")
    if p is None or p >= _ALPHA:
        reasons.append("not significant at alpha=0.05")
    return {
        "metric": "p_value",
        "value": None if p is None else round(p, 4),
        "floor": _ALPHA,
        "weak": top["weak"],
        "note": (
            f"strongest pair {top['a']}–{top['b']} (r={top['corr']:.3f}, n={n}): "
            + ("; ".join(reasons) if reasons else f"significant at alpha={_ALPHA}")
        ),
    }


# 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.

class CorrelationInput(BaseModel):
    """Planner-facing inputs for `analyze_correlation` (spec_builder contract)."""

    data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).")
    column_ids: list[str] | None = Field(
        None,
        description="NUMERIC columns to correlate pairwise; omit to use every numeric column.",
    )
    method: str = Field(
        "pearson",
        description="Correlation method: pearson (linear), spearman or kendall (rank-based).",
    )

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_ab_test
  - 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 except `n_*` row-count columns.
        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", "n", "p_value", "weak"}] unique pairs,
                       strongest |corr| first; n = rows complete for both columns,
                       weak = n < 10 or p >= 0.05
            quality  — reserved block for the strongest pair (CK7 surfaces `weak`)
            caveats  — a note when the strongest pair is weak evidence

    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:
        # `n_*` columns are row counts the IR repair adds to rankings (n_periods), not
        # measures — correlating against them is noise.
        cols = [
            c for c in df.columns
            if pd.api.types.is_numeric_dtype(df[c]) and not str(c).startswith("n_")
        ]
    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: list[dict[str, Any]] = []
    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:
                both = df[[cols[i], cols[j]]].dropna()
                n = int(both.shape[0])
                p = _p_value(both.iloc[:, 0], both.iloc[:, 1], method)
                pairs.append({"a": cols[i], "b": cols[j], "corr": val, "n": n,
                              "p_value": p, "weak": _is_weak(n, p)})
    pairs.sort(key=lambda p: abs(p["corr"]), reverse=True)

    top = pairs[0] if pairs else None
    caveats: list[str] = []
    if top is not None and top["weak"]:
        if top["n"] < _WEAK_MIN_N:
            caveats.append(
                f"The strongest correlation ({top['a']}–{top['b']}, r={top['corr']:.3f}) "
                f"rests on {top['n']} rows — on so few points a single value can drive r, "
                "so treat it as a lead to check, not an established relationship."
            )
        else:
            caveats.append(
                f"The strongest correlation ({top['a']}–{top['b']}, r={top['corr']:.3f}) "
                f"is not significant (p >= {_ALPHA}) on {top['n']} rows — it may be noise."
            )

    return {
        "method": method,
        "columns": cols,
        "matrix": matrix,
        "pairs": pairs,
        "quality": _quality(top),
        "caveats": caveats,
    }