src/customer_intelligence/app/components/charts.py

"""Reusable Plotly chart builders."""

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go

_COLORS = [
    "#3B82F6", "#10B981", "#F59E0B", "#EF4444",
    "#8B5CF6", "#06B6D4", "#F97316", "#84CC16",
]
_TEMPLATE = "plotly_white"


def make_line_chart(
    df: pd.DataFrame,
    x: str,
    y: str,
    title: str = "",
    color: str = "#3B82F6",
    add_ma: int | None = 30,
) -> go.Figure:
    fig = px.line(df, x=x, y=y, title=title, color_discrete_sequence=[color])
    if add_ma and len(df) > add_ma:
        ma = df[y].rolling(add_ma).mean()
        fig.add_scatter(x=df[x], y=ma, mode="lines", name=f"{add_ma}-day MA",
                        line=dict(dash="dash", color="#F59E0B"))
    fig.update_layout(height=350, margin=dict(t=40, b=20), template=_TEMPLATE)
    return fig


def make_bar_chart(
    df: pd.DataFrame,
    x: str,
    y: str,
    title: str = "",
    orientation: str = "v",
    color_col: str | None = None,
) -> go.Figure:
    fig = px.bar(df, x=x, y=y, title=title, orientation=orientation,
                 color=color_col, text_auto=".2s",
                 color_discrete_sequence=_COLORS)
    fig.update_layout(height=350, margin=dict(t=40, b=20), template=_TEMPLATE)
    return fig


def make_treemap(
    df: pd.DataFrame,
    path: list[str],
    values: str,
    title: str = "",
) -> go.Figure:
    fig = px.treemap(df, path=path, values=values, title=title,
                     color_discrete_sequence=_COLORS)
    fig.update_layout(height=400, margin=dict(t=40, b=20), template=_TEMPLATE)
    return fig


def make_donut(df: pd.DataFrame, names: str, values: str, title: str = "") -> go.Figure:
    fig = px.pie(df, names=names, values=values, title=title, hole=0.4,
                 color_discrete_sequence=_COLORS)
    fig.update_layout(height=350, margin=dict(t=40, b=20), template=_TEMPLATE)
    return fig


def make_histogram(df: pd.DataFrame, x: str, title: str = "", nbins: int = 40) -> go.Figure:
    fig = px.histogram(df, x=x, title=title, nbins=nbins,
                       color_discrete_sequence=_COLORS)
    fig.update_layout(height=350, margin=dict(t=40, b=20), template=_TEMPLATE)
    return fig


def make_scatter(
    df: pd.DataFrame,
    x: str,
    y: str,
    color: str | None = None,
    title: str = "",
    opacity: float = 0.5,
) -> go.Figure:
    fig = px.scatter(df, x=x, y=y, color=color, title=title, opacity=opacity,
                     color_discrete_sequence=_COLORS)
    fig.update_layout(height=400, margin=dict(t=40, b=20), template=_TEMPLATE)
    return fig


def make_map_scatter(
    df: pd.DataFrame,
    lat: str,
    lon: str,
    color: str | None = None,
    size: str | None = None,
    title: str = "",
    hover_name: str | None = None,
) -> go.Figure:
    fig = px.scatter_mapbox(
        df, lat=lat, lon=lon, color=color, size=size, title=title,
        hover_name=hover_name, mapbox_style="open-street-map", zoom=3,
        center={"lat": -14.2, "lon": -51.9}, opacity=0.6,
        color_discrete_sequence=_COLORS,
    )
    fig.update_layout(height=450, margin=dict(t=40, b=0))
    return fig


def make_gauge(
    value: float,
    title: str = "",
    max_val: float = 1.0,
    good_high: bool = False,
) -> go.Figure:
    lo, hi = ("salmon", "lightgreen") if good_high else ("lightgreen", "salmon")
    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=value * 100,
        title={"text": title, "font": {"size": 16}},
        gauge={
            "axis": {"range": [0, max_val * 100]},
            "bar": {"color": "#3B82F6"},
            "steps": [
                {"range": [0, 30], "color": lo},
                {"range": [30, 60], "color": "khaki"},
                {"range": [60, 100], "color": hi},
            ],
            "threshold": {
                "line": {"color": "#1e293b", "width": 3},
                "thickness": 0.8,
                "value": value * 100,
            },
        },
        number={"suffix": "%", "font": {"size": 36}},
    ))
    fig.update_layout(
        height=320, margin=dict(t=50, b=30, l=30, r=30), template=_TEMPLATE
    )
    return fig


def make_heatmap(df: pd.DataFrame, title: str = "") -> go.Figure:
    """Render a retention heatmap. Rows = cohort, columns = month offset.

    NaN cells (cohort too new to have data for that month) are rendered in
    light grey so they're visually distinct from real 0 % retention cells.
    The colorscale uses -1 as the NaN sentinel: values in [-1, 0) → grey,
    values in [0, 100] → Blues gradient.
    """
    import math

    raw = df.values.tolist()

    # Replace NaN with -1 (sentinel for "no data yet")
    z = [
        [-1 if (v is None or (isinstance(v, float) and math.isnan(v))) else v for v in row]
        for row in raw
    ]
    text_vals = [
        [f"{v:.1f}%" if v >= 0 else "" for v in row]
        for row in z
    ]

    # Custom colorscale: grey for -1 (no data), Blues for 0–100
    # Positions are normalised over the range [-1, 100] → span of 101
    grey = "#D1D5DB"
    colorscale = [
        [0.0,        grey],        # -1  → grey
        [0.99 / 101, grey],        # ~0  → grey boundary
        [1.0  / 101, "#EFF6FF"],   # 0 % → lightest blue
        [1.0,        "#1E3A8A"],   # 100 % → darkest blue
    ]

    fig = go.Figure(go.Heatmap(
        z=z,
        x=df.columns.tolist(),
        y=df.index.tolist(),
        text=text_vals,
        texttemplate="%{text}",
        textfont={"size": 11},
        colorscale=colorscale,
        zmin=-1,
        zmax=100,
        colorbar={
            "title": "Retention %",
            "ticksuffix": "%",
            "tickvals": [0, 25, 50, 75, 100],
        },
    ))
    fig.update_layout(
        title=title,
        height=max(300, len(df) * 28 + 80),
        margin=dict(t=50, b=20, l=80, r=40),
        template=_TEMPLATE,
        xaxis_title="Months Since First Purchase",
        yaxis_title="Acquisition Cohort",
        yaxis={"autorange": "reversed"},
    )
    return fig


def make_sunburst(df: pd.DataFrame, path: list[str], values: str, title: str = "") -> go.Figure:
    fig = px.sunburst(df, path=path, values=values, title=title,
                      color_discrete_sequence=_COLORS)
    fig.update_layout(height=420, margin=dict(t=40, b=0), template=_TEMPLATE)
    return fig