"""analyze_importance — rank which columns most explain a target (§0.9).

Third ML analytical tool. Fits a standardised linear model (Ridge for a numeric
target, Logistic for a categorical one) on a train split, scores it on a holdout, and
ranks the features by the magnitude of their standardised coefficient — "what drives
downtime?".

Named `importance`, not `driver`: "driver" was rejected at the 2026-09-15 discussion as
ambiguous, and "feature importance" as too technical. The thing the user wants is
*"kolom mana yang paling berdampak"*.

**Standardised coefficients, not permutation importance.** For a linear model the
standardised coefficient IS the importance, and it comes free with the fit. Permutation
importance re-scores the model once per feature, which is the latency tall pole
identified in the plan (§0.9 A8 / E4) — it buys nothing here because the model is
linear by design. If a non-linear model is ever added, permutation becomes necessary
and belongs behind an explicit argument.

**Leakage-safe by construction.** The scaler is fitted on the TRAIN split only and then
applied to the holdout, so no test-set statistic reaches the fit — the one thing the
plan (R-3) and the survey literature both single out. It also reports its holdout
metric rather than a training score: a ranking from a model that cannot predict is not
a finding, and the checkpoint's CK7 needs that number to flag the case.

Follows docs/tools/ANALYTICAL_TOOLS_CONVENTIONS.md: compute is synchronous except for
the `await run_fit` offload; `data` is never self-fetched (Pattern A); errors escape to
the invoker's never-throw seam; output is deterministic (fixed random_state). Mirrors
analyze_anomaly's shape so the ML tools stay uniform.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from pydantic import BaseModel, Field

from src.tools.analytics.charts import pareto_chart
from src.tools.analytics.descriptive import ColumnNotFoundError
from src.tools.ml.executor import run_fit
from src.tools.spec_builder import build_description

# A supervised fit with a holdout needs a sane floor; below this the ranking is noise.
_MIN_ROWS = 20
# Feature cap (§0.9 R-5): bounds the fit AND keeps the ranking readable.
_MAX_FEATURES = 20
# A target with at most this many distinct values is treated as categorical.
_MAX_CLASSES = 20
# Fixed so the same data always ranks the same (family-B determinism promise).
_RANDOM_STATE = 42
_TEST_SIZE = 0.25
# Below this holdout score the ranking is noise, not a finding — surfaced as the
# reserved `quality` block the checkpoint's CK7 reads.
_WEAK_FIT = 0.3
# Chart column names (internal; never collide with user data).
_FEATURE_COL = "feature"
_WEIGHT_COL = "importance"


class NotEnoughDataError(ValueError):
    """Too few complete rows to fit an importance model (error_code NOT_ENOUGH_DATA)."""


class NoUsableFeaturesError(ValueError):
    """No usable feature column survived cleaning (error_code NO_USABLE_FEATURES)."""


class ConstantTargetError(ValueError):
    """The target never varies, so nothing can explain it (error_code CONSTANT_TARGET)."""


class ImportanceInput(BaseModel):
    # `data` FIRST — property order follows field order (planner-prompt readability).
    data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).")
    target_column: str = Field(
        ..., description="The column to be explained — the outcome the user cares about."
    )
    feature_columns: list[str] | None = Field(
        None,
        description="Candidate explanatory NUMERIC columns; omit to use every other "
        "numeric column in the table.",
    )
    top_n: int | None = Field(
        None, description="Return only the top N features by importance; omit for all."
    )


DESCRIPTION = build_description(
    summary="Rank which columns most explain a target, with a holdout score.",
    use_when="the user asks what drives, explains, influences or affects a metric — "
    "'apa yang paling mempengaruhi downtime', 'what drives churn', "
    "'faktor apa yang bikin PA turun'.",
    dont_use_when=[
        "the user asks how two specific columns move together -> analyze_correlation "
        "(that is a pairwise relationship, this ranks many columns against one target)",
        "the user wants which CATEGORIES make up most of a total -> analyze_pareto "
        "(that is arithmetic share, this is statistical explanation)",
        "the user wants a future value -> analyze_forecast",
        "the target column does not exist in the data -> report it as infeasible",
    ],
    output="a ranked list of features with signed standardised coefficients and "
    "normalised importance, the model's holdout score, and a pareto-style bar chart.",
    examples=[
        "apa yang paling mempengaruhi downtime?",
        "what drives low availability?",
        "faktor apa yang berpengaruh ke MTTR?",
    ],
)


def _fit_linear(x: np.ndarray, y: np.ndarray, classify: bool, random_state: int):
    """Split → scale on TRAIN only → fit → score on the holdout.

    Imported lazily so app startup doesn't pay for scikit-learn unless an importance
    analysis runs. Runs inside run_fit's thread.

    Returns (coefficients, metric_name, metric_value). Fitting the scaler inside the
    train split is what keeps the holdout score honest — see the module docstring.
    """
    from sklearn.linear_model import LogisticRegression, Ridge
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler

    stratify = y if classify and np.min(np.bincount(y.astype(int))) >= 2 else None
    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=_TEST_SIZE, random_state=random_state, stratify=stratify
    )

    scaler = StandardScaler().fit(x_train)  # TRAIN ONLY — no test statistic leaks in
    x_train_s, x_test_s = scaler.transform(x_train), scaler.transform(x_test)

    if classify:
        model = LogisticRegression(max_iter=1000, random_state=random_state)
        model.fit(x_train_s, y_train)
        # Multi-class gives one row of coefficients per class; average their magnitude.
        coefs = model.coef_[0] if model.coef_.shape[0] == 1 else model.coef_.mean(axis=0)
        return coefs, "accuracy", float(model.score(x_test_s, y_test))

    model = Ridge(random_state=random_state)
    model.fit(x_train_s, y_train)
    return model.coef_, "r2", float(model.score(x_test_s, y_test))


def _importance_chart(rows: list[dict], target_column: str) -> dict[str, object]:
    """Pareto of the ranked features.

    Reuses `pareto_chart` deliberately: bars descending plus a cumulative line answers
    "which few columns account for most of the explanatory weight", which is exactly how
    an importance ranking is read. The cumulative share is share of *explanatory
    weight*, not a causal contribution — the caveat says so.
    """
    chart_df = pd.DataFrame(
        {
            _FEATURE_COL: [r["feature"] for r in rows],
            _WEIGHT_COL: [r["importance"] for r in rows],
        }
    )
    return pareto_chart(
        chart_df, _FEATURE_COL, _WEIGHT_COL, title=f"What explains {target_column}"
    )


async def analyze_importance(
    df: pd.DataFrame,
    target_column: str,
    feature_columns: list[str] | None = None,
    top_n: int | None = None,
) -> dict[str, object]:
    """Rank `feature_columns` by how much they explain `target_column`.

    Async ONLY so the model fit can be offloaded via `run_fit`; everything else is
    synchronous compute. `data` is materialised upstream (Pattern A) and arrives as
    `df`.

    Returns a dict with: task ("regression"/"classification"), target_column,
    n_rows, n_features, `quality` (holdout metric + floor + weak flag), `features`
    (ranked, each with a signed coefficient and a normalised importance share),
    `caveats`, and `charts` (the reserved auto-chart key).

    Raises (all wrapped by the invoker's never-throw seam):
        ColumnNotFoundError, NotEnoughDataError, NoUsableFeaturesError,
        ConstantTargetError.
    """
    if target_column not in df.columns:
        raise ColumnNotFoundError(f"target column not found: {target_column!r}")

    if feature_columns:
        missing = [c for c in feature_columns if c not in df.columns]
        if missing:
            raise ColumnNotFoundError(f"feature columns not found: {missing}")
        candidates = [c for c in feature_columns if c != target_column]
    else:
        candidates = [c for c in df.columns if c != target_column]

    # Numeric features only; a column that is entirely non-numeric is dropped, not an
    # error — the planner picking a text column shouldn't fail the whole analysis.
    features = df[candidates].apply(pd.to_numeric, errors="coerce")
    usable = [c for c in features.columns if features[c].notna().any()]
    dropped = [c for c in candidates if c not in usable]

    # Drop zero-variance features BEFORE the cap. A column that never changes cannot
    # explain one that does — its standardised coefficient is zero by construction —
    # so keeping it is not merely useless: `capped` takes the first _MAX_FEATURES in
    # COLUMN ORDER, so every constant kept silently evicts a real candidate further
    # right. The demo file has 8 constant columns among 36 and lost 3 genuine ones
    # that way ("only the first 20 of 23"). Same blindness as A28, one layer down.
    constant = [c for c in usable if features[c].nunique(dropna=True) <= 1]
    usable = [c for c in usable if c not in constant]
    if not usable:
        raise NoUsableFeaturesError(
            f"no numeric feature column among {candidates} varies; nothing to rank"
        )
    capped = usable[:_MAX_FEATURES]

    target = df[target_column]
    numeric_target = pd.to_numeric(target, errors="coerce")
    # Categorical when it isn't numeric at all, or is numeric with few distinct values
    # (a 0/1 flag is a classification target, not a regression one).
    classify = numeric_target.isna().all() or numeric_target.nunique() <= _MAX_CLASSES

    frame = features[capped].join(target.rename("__target__"))
    frame = frame.dropna()
    if frame.shape[0] < _MIN_ROWS:
        raise NotEnoughDataError(
            f"need >= {_MIN_ROWS} rows with a value in {target_column!r} and every "
            f"feature to rank importance, got {frame.shape[0]}"
        )

    y_raw = frame["__target__"]
    if y_raw.nunique() <= 1:
        raise ConstantTargetError(
            f"{target_column!r} never varies in these rows, so nothing can explain it"
        )

    x = frame[capped].to_numpy(dtype=float)
    if classify:
        y = pd.factorize(y_raw)[0].astype(float)
    else:
        y = pd.to_numeric(y_raw, errors="coerce").to_numpy(dtype=float)

    coefs, metric, metric_value = await run_fit(
        _fit_linear, x, y, classify, _RANDOM_STATE, label="analyze_importance"
    )

    magnitudes = np.abs(coefs)
    total = float(magnitudes.sum())
    ranked = sorted(
        (
            {
                "feature": col,
                "coefficient": round(float(coef), 4),
                "direction": "increases" if coef > 0 else "decreases",
                "importance": round(float(mag / total) if total else 0.0, 4),
            }
            for col, coef, mag in zip(capped, coefs, magnitudes, strict=True)
        ),
        key=lambda r: r["importance"],
        reverse=True,
    )
    if top_n:
        ranked = ranked[:top_n]

    caveats = [
        "Association, not causation — a high-importance column moves with the target, "
        "which is not evidence that changing it would change the target.",
        f"Ranking is from a linear model scored on a {int(_TEST_SIZE * 100)}% holdout; "
        f"{metric} = {metric_value:.3f}.",
    ]
    if metric_value < _WEAK_FIT:
        caveats.append(
            f"The model barely explains {target_column!r} ({metric} = {metric_value:.3f}), "
            "so this ranking is weak evidence — treat the order as tentative."
        )
    if dropped:
        caveats.append(f"Non-numeric columns skipped: {dropped}.")
    if constant:
        caveats.append(
            f"Columns with the same value in every row were skipped (they cannot "
            f"explain anything): {constant}."
        )
    if len(usable) > _MAX_FEATURES:
        caveats.append(
            f"Only the first {_MAX_FEATURES} of {len(usable)} numeric columns were used."
        )

    return {
        "task": "classification" if classify else "regression",
        "target_column": target_column,
        "n_rows": int(frame.shape[0]),
        "n_features": len(capped),
        # Reserved `quality` block (conventions §3): every ML tool reports its own
        # score and whether it clears its own floor. The tool owns that judgement —
        # it knows what "good" means for its method — and the checkpoint's CK7 just
        # surfaces `weak` deterministically, so a new tool needs no checkpoint change.
        "quality": {
            "metric": metric,
            "value": round(metric_value, 4),
            "floor": _WEAK_FIT,
            "weak": metric_value < _WEAK_FIT,
        },
        "features": ranked,
        "caveats": caveats,
        "charts": [_importance_chart(ranked, target_column)],
    }