palladium/core/notebook_helpers/charts.py

"""
Plotly charts for TEI analyses.

Each function returns a ``plotly.graph_objects.Figure`` so callers can
``.show()`` (notebook), pass to ``st.plotly_chart`` (Streamlit), or write to
HTML / image. No styling is hard-coded beyond a neutral default palette.
"""

from __future__ import annotations

from collections.abc import Iterable

import plotly.graph_objects as go

PALETTE = {
    "benefits": "#2E7D32",  # green
    "costs": "#C62828",  # red
    "net_positive": "#1565C0",  # blue
    "net_negative": "#C62828",
    "cumulative": "#616161",  # grey
}


def cashflow_chart(
    yearly_breakdown: list[dict],
    *,
    title: str = "Cash Flow Analysis (Risk-Adjusted)",
    initial_cost: float = 0.0,
) -> go.Figure:
    """
    Stacked bars of benefits & costs by year + cumulative net line.

    Mirrors the chart on page 25 of the Forrester Amazon Connect TEI study.
    """
    if not yearly_breakdown:
        return go.Figure(layout={"title": title})

    years = ["Initial"] + [f"Year {row['year']}" for row in yearly_breakdown]
    benefits = [0.0] + [float(row.get("benefits", 0)) for row in yearly_breakdown]
    costs = [-float(initial_cost)] + [
        -float(row.get("costs", 0)) for row in yearly_breakdown
    ]
    # cumulative_net assumes initial cost has already been deducted
    cumulative = [-float(initial_cost)] + [
        float(row.get("cumulative_net", 0)) for row in yearly_breakdown
    ]

    fig = go.Figure()
    fig.add_bar(
        name="Total benefits",
        x=years,
        y=benefits,
        marker_color=PALETTE["benefits"],
    )
    fig.add_bar(
        name="Total costs",
        x=years,
        y=costs,
        marker_color=PALETTE["costs"],
    )
    fig.add_scatter(
        name="Cumulative net benefits",
        x=years,
        y=cumulative,
        mode="lines+markers",
        line={"color": PALETTE["cumulative"], "width": 3},
    )
    fig.update_layout(
        title=title,
        barmode="relative",
        yaxis_tickformat="$,.0f",
        legend={"orientation": "h", "y": -0.15},
        margin={"l": 40, "r": 20, "t": 60, "b": 40},
    )
    return fig


def benefits_bar(items: list[dict], *, title: str = "Benefits (Three-Year)") -> go.Figure:
    """Horizontal bars of risk-adjusted three-year totals per benefit."""
    labels: list[str] = []
    totals: list[float] = []
    for it in items:
        rf = float(it.get("risk_adjustment") or 0.0)
        yv = it.get("year_values") or {}
        ra_total = sum(float(v or 0) * (1.0 - rf) for v in yv.values())
        labels.append(it.get("label", "") or it.get("field_key", ""))
        totals.append(ra_total)

    fig = go.Figure(
        go.Bar(
            x=totals,
            y=labels,
            orientation="h",
            marker_color=PALETTE["benefits"],
            text=[f"${t/1_000_000:,.1f}M" for t in totals],
            textposition="auto",
        )
    )
    fig.update_layout(
        title=title,
        xaxis_tickformat="$,.0f",
        yaxis={"autorange": "reversed"},
        margin={"l": 40, "r": 20, "t": 60, "b": 40},
    )
    return fig


def cost_breakdown_pie(
    items: list[dict], *, title: str = "Cost Breakdown (Three-Year, Risk-Adjusted)"
) -> go.Figure:
    """Pie chart of risk-adjusted costs by category/label."""
    labels: list[str] = []
    values: list[float] = []
    for it in items:
        rf = float(it.get("risk_adjustment") or 0.0)
        yv = it.get("year_values") or {}
        initial = float(it.get("initial") or 0.0)
        ra_total = (
            initial * (1.0 + rf)
            + sum(float(v or 0) * (1.0 + rf) for v in yv.values())
        )
        labels.append(it.get("label", "") or it.get("field_key", ""))
        values.append(ra_total)

    fig = go.Figure(go.Pie(labels=labels, values=values, hole=0.35))
    fig.update_layout(title=title, margin={"l": 40, "r": 20, "t": 60, "b": 40})
    return fig


def scenario_comparison(scenarios: dict) -> go.Figure:
    """Grouped bars comparing NPV and Costs PV across scenarios."""
    keys: list[str] = list(scenarios.keys())
    if not keys:
        return go.Figure()
    benefits = [float(scenarios[k].get("total_benefits_pv") or 0) for k in keys]
    costs = [float(scenarios[k].get("total_costs_pv") or 0) for k in keys]
    npvs = [float(scenarios[k].get("npv") or 0) for k in keys]

    fig = go.Figure()
    fig.add_bar(name="Benefits PV", x=keys, y=benefits, marker_color=PALETTE["benefits"])
    fig.add_bar(name="Costs PV", x=keys, y=costs, marker_color=PALETTE["costs"])
    fig.add_bar(name="NPV", x=keys, y=npvs, marker_color=PALETTE["net_positive"])
    fig.update_layout(
        title="Scenario Comparison",
        barmode="group",
        yaxis_tickformat="$,.0f",
        legend={"orientation": "h", "y": -0.15},
    )
    return fig


def cumulative_benefits_chart(
    yearly_breakdown: list[dict],
    *,
    title: str = "Cumulative Net Benefits",
) -> go.Figure:
    """Single-line cumulative net benefits trajectory."""
    if not yearly_breakdown:
        return go.Figure(layout={"title": title})
    years = [f"Year {row['year']}" for row in yearly_breakdown]
    cumulative = [float(row.get("cumulative_net", 0)) for row in yearly_breakdown]
    fig = go.Figure(
        go.Scatter(
            x=years,
            y=cumulative,
            mode="lines+markers",
            fill="tozeroy",
            line={"color": PALETTE["net_positive"], "width": 3},
        )
    )
    fig.update_layout(title=title, yaxis_tickformat="$,.0f")
    return fig


def waterfall(values: Iterable[tuple[str, float]], *, title: str = "TEI Waterfall") -> go.Figure:
    """
    Generic waterfall (pass tuples of (label, value)).

    Used by 03_business_case to show: Benefits PV → Costs PV → NPV.
    """
    labels, amounts = zip(*values, strict=True) if values else ([], [])
    measures = ["relative"] * (len(labels) - 1) + ["total"] if labels else []
    fig = go.Figure(
        go.Waterfall(
            x=list(labels),
            y=list(amounts),
            measure=measures,
            text=[f"${v/1_000_000:,.1f}M" for v in amounts],
            textposition="outside",
        )
    )
    fig.update_layout(title=title, yaxis_tickformat="$,.0f")
    return fig