src/tools/analytics/charts.py

"""charts — shared chart-envelope kit + auto-chart builders (analytical tools round).

The house style (colorway, layout preset, JSON-safe coercion) lives here as the
single source of truth. `visualization.py`'s `render_chart` tool imports this kit,
and the analytical tools call the builders below to attach an auto-chart to their
output ("sepaket chart"): each builder returns the `dataeyond.chart.v1` envelope
(SPINE_V2_PLAN §4.2) the FE renders with plotly.js.

Deliberately NO plotly import — every builder returns a hand-built dict (locked
decision, DEV_PLAN deferred row #26), exactly like `render_chart`. Style is the
fixed module preset, never a planner/LLM choice.

Builders (envelope shape `{schema, chart_type, title, plotly:{data, layout}}`):
  - box_chart               distribution / spread (descriptive, anomaly context)
  - pareto_chart            80/20 contribution, bars + cumulative-% line (analyze_pareto)
  - cluster_scatter_chart   points colored per cluster (analyze_cluster)
  - anomaly_scatter_chart   points with anomalies highlighted (analyze_anomaly)

The wiring that carries these envelopes from a tool into the turn's persisted
charts (ToolOutput.meta -> the chat_handler save loop) is a spine change outside
this module and is added separately.
"""

from __future__ import annotations

import datetime as _dt
from decimal import Decimal
from typing import Any

import numpy as np
import pandas as pd

from src.tools.analytics.descriptive import ColumnNotFoundError

SCHEMA = "dataeyond.chart.v1"

# --------------------------------------------------------------------------- #
# House style preset — fixed by design, not a planner argument.
#
# Colorway = the dataviz reference categorical palette (8 slots, light mode).
# The slot ORDER is the colorblind-safety mechanism (validated: worst adjacent
# CVD deltaE 24.2) — series take slots in this fixed order, never re-picked.
# Backgrounds are transparent so the chart inherits the FE surface.
# --------------------------------------------------------------------------- #
COLORWAY = [
    "#2a78d6",  # blue
    "#1baf7a",  # aqua
    "#eda100",  # yellow
    "#008300",  # green
    "#4a3aa7",  # violet
    "#e34948",  # red
    "#e87ba4",  # magenta
    "#eb6834",  # orange
]

_INK_SECONDARY = "#52514e"  # body text
_INK_MUTED = "#898781"      # axis titles / tick labels (mode-neutral)
_GRIDLINE = "#e1e0d9"       # hairline grid
_BASELINE = "#c3c2b7"       # axis line
_HIGHLIGHT = COLORWAY[5]    # red — the "look here" ink for anomaly points

LAYOUT_PRESET: dict[str, object] = {
    "colorway": COLORWAY,
    "font": {
        "family": 'system-ui, -apple-system, "Segoe UI", sans-serif',
        "size": 13,
        "color": _INK_SECONDARY,
    },
    "paper_bgcolor": "rgba(0,0,0,0)",
    "plot_bgcolor": "rgba(0,0,0,0)",
    "margin": {"t": 48, "r": 16, "b": 48, "l": 56},
    "bargap": 0.25,
}

# Mark specs: thin marks — 2px lines, 8px markers.
_LINE_STYLE = {"width": 2}
_MARKER_STYLE = {"size": 8}


# --------------------------------------------------------------------------- #
# JSON-safe coercion + axis chrome (shared with render_chart).
# --------------------------------------------------------------------------- #
def _clean(value: object) -> object:
    """One JSON-safe scalar: numpy -> Python, NaN/NaT -> None, dates -> ISO strings.

    Chart specs are persisted as JSONB and shipped to the FE, so every value in a
    trace array must survive json.dumps.
    """
    if value is None or value is pd.NaT:
        return None
    if isinstance(value, float) and np.isnan(value):
        return None
    if hasattr(value, "item"):  # numpy scalar (incl. Timestamp via below)
        value = value.item()
        if isinstance(value, float) and np.isnan(value):
            return None
    if isinstance(value, Decimal):
        return float(value)
    if isinstance(value, pd.Timestamp):
        return value.isoformat()
    if isinstance(value, _dt.datetime | _dt.date):
        return value.isoformat()
    if isinstance(value, bool | int | float | str):
        return value
    return str(value)


def _values(col: pd.Series) -> list[object]:
    return [_clean(v) for v in col.tolist()]


def _series_label(value: object) -> str:
    """Trace name for one series group; null groups get an explicit label."""
    cleaned = _clean(value)
    return "(missing)" if cleaned is None else str(cleaned)


def _axis(label: str) -> dict[str, object]:
    """Recessive axis chrome + the column name as the axis title."""
    return {
        "title": {"text": label, "font": {"color": _INK_MUTED}},
        "tickfont": {"color": _INK_MUTED},
        "gridcolor": _GRIDLINE,
        "linecolor": _BASELINE,
        "automargin": True,
    }


# --------------------------------------------------------------------------- #
# Envelope + layout assembly.
# --------------------------------------------------------------------------- #
def _require(df: pd.DataFrame, cols: list[str]) -> None:
    missing = [c for c in cols if c not in df.columns]
    if missing:
        raise ColumnNotFoundError(f"columns not found: {missing}")


def _layout(
    title: str,
    *,
    x_label: str | None = None,
    y_label: str | None = None,
    showlegend: bool = False,
    extra: dict[str, Any] | None = None,
) -> dict[str, Any]:
    layout: dict[str, Any] = {**LAYOUT_PRESET, "title": {"text": title}, "showlegend": showlegend}
    if x_label is not None:
        layout["xaxis"] = _axis(x_label)
    if y_label is not None:
        layout["yaxis"] = _axis(y_label)
    if extra:
        layout.update(extra)
    return layout


def envelope(
    chart_type: str, title: str, traces: list[dict[str, Any]], layout: dict[str, Any]
) -> dict[str, Any]:
    """Wrap traces + layout in the `dataeyond.chart.v1` envelope (SPINE_V2_PLAN §4.2)."""
    return {
        "schema": SCHEMA,
        "chart_type": chart_type,
        "title": title,
        "plotly": {"data": traces, "layout": layout},
    }


# --------------------------------------------------------------------------- #
# Auto-chart builders.
# --------------------------------------------------------------------------- #
def box_chart(
    df: pd.DataFrame,
    value_column: str,
    group_column: str | None = None,
    title: str | None = None,
) -> dict[str, Any]:
    """Box plot of `value_column`, optionally one box per `group_column` value."""
    _require(df, [value_column, *([group_column] if group_column else [])])
    chart_title = title or (
        f"Distribution of {value_column}" + (f" by {group_column}" if group_column else "")
    )
    trace: dict[str, Any] = {
        "type": "box",
        "y": _values(df[value_column]),
        "name": value_column,
        "marker": {"color": COLORWAY[0]},
    }
    if group_column:
        trace["x"] = _values(df[group_column])
    layout = _layout(chart_title, x_label=group_column, y_label=value_column)
    return envelope("box", chart_title, [trace], layout)


def pareto_chart(
    df: pd.DataFrame,
    category_column: str,
    value_column: str,
    title: str | None = None,
    top_n: int | None = None,
) -> dict[str, Any]:
    """Pareto (80/20): bars of `value_column` per category, descending, plus a
    cumulative-% line on a secondary axis and an 80% reference line. Rows are
    summed per category; the tail past `top_n` is folded into an 'Other' bar."""
    _require(df, [category_column, value_column])
    values = pd.to_numeric(df[value_column], errors="coerce").fillna(0.0)
    agg = values.groupby(df[category_column], dropna=False).sum().sort_values(ascending=False)
    if top_n is not None and len(agg) > top_n:
        head = agg.iloc[:top_n]
        agg = pd.concat([head, pd.Series({"Other": float(agg.iloc[top_n:].sum())})])
    total = float(agg.sum()) or 1.0

    categories = [_series_label(c) for c in agg.index]
    bar_values = [float(v) for v in agg.to_numpy()]
    cum_pct = [round(v, 2) for v in (agg.cumsum() / total * 100.0).to_numpy()]
    chart_title = title or f"Pareto of {value_column} by {category_column}"

    bar = {"type": "bar", "x": categories, "y": bar_values, "name": value_column}
    line = {
        "type": "scatter",
        "mode": "lines+markers",
        "x": categories,
        "y": cum_pct,
        "name": "Cumulative %",
        "yaxis": "y2",
        "line": {**_LINE_STYLE, "color": _HIGHLIGHT},
        "marker": {**_MARKER_STYLE, "color": _HIGHLIGHT},
    }
    layout = _layout(
        chart_title,
        x_label=category_column,
        y_label=value_column,
        showlegend=True,
        extra={
            "yaxis2": {
                "title": {"text": "Cumulative %", "font": {"color": _INK_MUTED}},
                "tickfont": {"color": _INK_MUTED},
                "overlaying": "y",
                "side": "right",
                "range": [0, 105],
                "showgrid": False,
            },
            "shapes": [
                {
                    "type": "line",
                    "xref": "paper",
                    "x0": 0,
                    "x1": 1,
                    "yref": "y2",
                    "y0": 80,
                    "y1": 80,
                    "line": {"color": _BASELINE, "width": 1, "dash": "dot"},
                }
            ],
        },
    )
    return envelope("pareto", chart_title, [bar, line], layout)


def cluster_scatter_chart(
    df: pd.DataFrame,
    x: str,
    y: str,
    cluster_column: str,
    title: str | None = None,
) -> dict[str, Any]:
    """Scatter of `x` vs `y`, one colorway-slotted trace per `cluster_column` value."""
    _require(df, [x, y, cluster_column])
    chart_title = title or f"{y} vs {x} by cluster"
    # sort=True keeps trace order (and colorway slots) deterministic across runs.
    traces = [
        {
            "type": "scatter",
            "mode": "markers",
            "x": _values(group[x]),
            "y": _values(group[y]),
            "name": _series_label(value),
            "marker": dict(_MARKER_STYLE),
        }
        for value, group in df.groupby(cluster_column, dropna=False, sort=True)
    ]
    layout = _layout(chart_title, x_label=x, y_label=y, showlegend=True)
    return envelope("cluster_scatter", chart_title, traces, layout)


def anomaly_scatter_chart(
    df: pd.DataFrame,
    x: str,
    y: str,
    flag_column: str,
    title: str | None = None,
) -> dict[str, Any]:
    """Scatter of `x` vs `y` with rows where `flag_column` is truthy drawn as
    larger red 'anomaly' markers over the normal points."""
    _require(df, [x, y, flag_column])
    flags = df[flag_column].astype(bool)
    normal, anomalies = df[~flags], df[flags]
    chart_title = title or f"{y} vs {x} (anomalies highlighted)"

    traces: list[dict[str, Any]] = [
        {
            "type": "scatter",
            "mode": "markers",
            "x": _values(normal[x]),
            "y": _values(normal[y]),
            "name": "normal",
            "marker": {**_MARKER_STYLE, "color": COLORWAY[0]},
        }
    ]
    if len(anomalies):
        traces.append(
            {
                "type": "scatter",
                "mode": "markers",
                "x": _values(anomalies[x]),
                "y": _values(anomalies[y]),
                "name": "anomaly",
                "marker": {
                    "size": 11,
                    "color": _HIGHLIGHT,
                    "line": {"width": 1, "color": "#ffffff"},
                },
            }
        )
    layout = _layout(chart_title, x_label=x, y_label=y, showlegend=True)
    return envelope("scatter", chart_title, traces, layout)