src/tools/analytics/visualization.py

"""render_chart — declarative chart-spec builder (SPINE_V2_PLAN §4.1, S2).

First tool of the `render_*` family: turns an already-materialized DataFrame into
a Plotly-conformant JSON spec wrapped in the `dataeyond.chart.v1` envelope
(SPINE_V2_PLAN §4.2). The FE renders it with plotly.js
(`Plotly.newPlot(el, spec.plotly.data, spec.plotly.layout)`).

Deliberately NOT a code generator and NOT a plotly import (locked decision,
DEV_PLAN deferred row #26): the spec is a hand-built dict, so no generated code
ever executes and no new dependency lands. Style (colors, axis chrome) is the
fixed module preset below — never a planner/LLM decision (CoDA's "Design
Explorer" phase collapsed to a constant).

Pattern A, same as the `analyze_*` family: `data` is resolved to a DataFrame by
the invoker from an upstream table-kind output; this function never fetches.
"""

from __future__ import annotations

import pandas as pd

from src.tools.analytics.charts import (
    _LINE_STYLE,
    _MARKER_STYLE,
    LAYOUT_PRESET,
    _axis,
    _series_label,
    _values,
)
from src.tools.analytics.descriptive import ColumnNotFoundError

# The shared house-style kit (colorway, layout preset, JSON-safe coercion, axis
# chrome) now lives in `charts.py` so `render_chart` and the analytical tools'
# auto-charts stay visually identical. render_chart-specific bits stay here.

# v1 chart types (SPINE_V2_PLAN §4.1). `pie` maps x -> labels, y -> values.
CHART_TYPES = ("bar", "line", "pie", "scatter")


class UnsupportedChartTypeError(ValueError):
    """The requested chart_type is not in CHART_TYPES (error_code UNSUPPORTED_CHART_TYPE)."""


# Prompt-style description read by the Planner to decide WHEN to pick this tool.
DESCRIPTION = """\
Summary: Build a chart specification (bar, line, pie, scatter) from an \
already-retrieved table, for the app to render. Maps one x column and one y \
column (pie: x -> slice labels, y -> slice values); an optional `series` column \
splits bar/line/scatter into one trace per category. Produces a spec only — it \
computes no numbers.

USE WHEN the user EXPLICITLY asks to see a chart — "plot", "chart", "graph", \
"visualize", "diagram" (ID: "grafik", "diagram", "plot", "visualisasikan", \
"buatkan grafik/diagram"). ALWAYS a tail step: `data` consumes the output of an \
upstream task that yields a table (retrieve_data or a table-kind analyze_*).

DON'T USE WHEN:
  - the user did not explicitly ask for a chart -> answer with tables/stats only \
(never add a speculative chart)
  - the numbers still need computing -> run the analyze_*/retrieve step first, \
then chart its table output
  - the upstream output is stats- or series-kind -> feed it a table-kind task \
(e.g. a grouped retrieve or analyze_aggregate)

Example questions:
  - "plot revenue by region as a bar chart"
  - "buatkan grafik penjualan per kategori"
  - "show a pie chart of orders by sales channel"
  - "visualize monthly revenue" (chain a month-grouped table first, then line)
"""


# analyze_aggregate renames a measure column to "<col>_<fn>" (e.g. pa_percent ->
# pa_percent_mean). A chart planned before that rename references the original
# name, so resolve it to the aggregated column when the exact name is absent.
_AGG_SUFFIXES = ("mean", "sum", "count", "min", "max", "median", "nunique", "avg")


def _resolve_col(name: str, columns) -> str:
    if name in columns:
        return name
    hits = [c for c in columns if c in {f"{name}_{s}" for s in _AGG_SUFFIXES}]
    return hits[0] if len(hits) == 1 else name


def render_chart(
    df: pd.DataFrame,
    chart_type: str,
    x: str,
    y: str,
    series: str | None = None,
    title: str | None = None,
) -> dict[str, object]:
    """Build a `dataeyond.chart.v1` envelope from a table (SPINE_V2_PLAN §4.2).

    Args:
        df: already-materialized data (the invoker resolves `data` upstream).
        x: column for the x axis (pie: slice labels).
        y: column for the y axis (pie: slice values).
        series: optional column splitting bar/line/scatter into one trace per
            distinct value (fixed-order colorway slots). Ignored for pie.
        title: chart title; defaults to "<y> by <x>".

    Returns:
        The envelope dict: {schema, chart_type, title, plotly: {data, layout}}.

    Raises:
        UnsupportedChartTypeError: chart_type not in CHART_TYPES.
        ColumnNotFoundError: x, y, or series absent from the data.
    """
    if chart_type not in CHART_TYPES:
        raise UnsupportedChartTypeError(
            f"unsupported chart_type '{chart_type}'; supported: {list(CHART_TYPES)}"
        )
    # Resolve chart columns to their aggregated names (pa_percent -> pa_percent_mean)
    # so a retrieve_data -> analyze_aggregate -> render_chart chain works.
    x = _resolve_col(x, df.columns)
    y = _resolve_col(y, df.columns)
    if series is not None:
        series = _resolve_col(series, df.columns)
    needed = [x, y] if chart_type == "pie" or series is None else [x, y, series]
    missing = [c for c in needed if c not in df.columns]
    if missing:
        raise ColumnNotFoundError(f"columns not found: {missing}")

    chart_title = title or f"{y} by {x}"

    if chart_type == "pie":
        traces: list[dict[str, object]] = [
            {"type": "pie", "labels": _values(df[x]), "values": _values(df[y])}
        ]
    else:
        plotly_type = "bar" if chart_type == "bar" else "scatter"
        mode = {"line": "lines+markers", "scatter": "markers"}.get(chart_type)

        def _trace(frame: pd.DataFrame, name: str) -> dict[str, object]:
            t: dict[str, object] = {
                "type": plotly_type,
                "x": _values(frame[x]),
                "y": _values(frame[y]),
                "name": name,
            }
            if mode is not None:
                t["mode"] = mode
                t["line"] = dict(_LINE_STYLE)
                t["marker"] = dict(_MARKER_STYLE)
            return t

        if series is None:
            traces = [_trace(df, y)]
        else:
            # sort=True keeps trace order (and therefore colorway slot
            # assignment) deterministic across runs of the same data.
            traces = [
                _trace(group, _series_label(value))
                for value, group in df.groupby(series, dropna=False, sort=True)
            ]

    layout: dict[str, object] = {
        **LAYOUT_PRESET,
        "title": {"text": chart_title},
        # Legend only when there is more than one thing to identify: multi-series
        # traces, or pie slices (a single named series is titled, not legended).
        "showlegend": len(traces) > 1 or chart_type == "pie",
    }
    if chart_type != "pie":
        layout["xaxis"] = _axis(x)
        layout["yaxis"] = _axis(y)

    return {
        "schema": "dataeyond.chart.v1",
        "chart_type": chart_type,
        "title": chart_title,
        "plotly": {"data": traces, "layout": layout},
    }