src/utils/plotting.py

"""Plotly helpers tuned for the RHTP dashboard (dark theme)."""
from __future__ import annotations

from typing import Iterable, Optional

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio

from .styling import PALETTE


# Build a custom dark template once at import time so every plot
# inherits the same surface / type / grid styling.
_template = go.layout.Template()
_template.layout = go.Layout(
    paper_bgcolor=PALETTE["bg"],
    plot_bgcolor=PALETTE["bg"],
    font=dict(family="Inter, system-ui, sans-serif",
              size=12, color=PALETTE["text"]),
    xaxis=dict(
        gridcolor=PALETTE["border"], zeroline=False,
        linecolor=PALETTE["border"],
        tickfont=dict(color=PALETTE["text_dim"]),
        title=dict(font=dict(color=PALETTE["text_dim"])),
    ),
    yaxis=dict(
        gridcolor=PALETTE["border"], zeroline=False,
        linecolor=PALETTE["border"],
        tickfont=dict(color=PALETTE["text_dim"]),
        title=dict(font=dict(color=PALETTE["text_dim"])),
    ),
    legend=dict(
        bgcolor="rgba(0,0,0,0)",
        font=dict(color=PALETTE["text"], size=11),
    ),
    hoverlabel=dict(
        bgcolor=PALETTE["surface2"],
        bordercolor=PALETTE["border"],
        font_color=PALETTE["text"],
        font_size=12,
    ),
    colorway=[
        PALETTE["secondary"], PALETTE["primary"], PALETTE["good"],
        PALETTE["accent"], "#A487E0", "#D679A4", "#5BCFA9", "#E08F5B",
    ],
)
pio.templates["rhtp_dark"] = _template
PLOT_TEMPLATE = "rhtp_dark"


def _layout(fig: go.Figure, height: int = 380, **kwargs) -> go.Figure:
    default_title = dict(
        font=dict(size=14, color=PALETTE["text"], family="Inter, system-ui, sans-serif"),
        x=0.0, xanchor="left",
    )
    if "title" in kwargs:
        t = kwargs.pop("title")
        if isinstance(t, dict):
            default_title.update(t)
        else:
            default_title["text"] = t
    fig.update_layout(
        template=PLOT_TEMPLATE,
        height=height,
        margin=dict(t=46, l=12, r=12, b=42),
        title=default_title,
        legend=dict(orientation="h", yanchor="bottom", y=1.02,
                    xanchor="left", x=0.0, font=dict(size=11)),
        **kwargs,
    )
    return fig


def line_by_group(
    df: pd.DataFrame, x: str, y: str, group: str, *,
    title: str = "", y_label: Optional[str] = None,
    treatment_band: Optional[tuple[float, float]] = None,
    hover_extra: Iterable[str] = (),
    height: int = 380,
) -> go.Figure:
    fig = px.line(
        df, x=x, y=y, color=group,
        hover_data=list(hover_extra),
        markers=True,
    )
    if treatment_band:
        x0, x1 = treatment_band
        fig.add_vrect(
            x0=x0, x1=x1, fillcolor=PALETTE["accent"], opacity=0.10,
            line_width=0,
            annotation_text="RHTP rollout",
            annotation_position="top left",
            annotation_font_size=10,
            annotation_font_color=PALETTE["accent"])
    fig.update_yaxes(title=y_label or y)
    fig.update_xaxes(title=x)
    return _layout(fig, height=height, title=dict(text=title))


def scatter_with_trend(
    df: pd.DataFrame, x: str, y: str, *,
    color: Optional[str] = None,
    size: Optional[str] = None,
    title: str = "", height: int = 380,
    log_x: bool = False,
) -> go.Figure:
    fig = px.scatter(
        df, x=x, y=y, color=color, size=size,
        opacity=0.80, trendline="ols",
        trendline_color_override=PALETTE["accent"],
        log_x=log_x,
    )
    fig.update_traces(marker=dict(line=dict(width=0.4, color=PALETTE["bg"])))
    return _layout(fig, height=height, title=dict(text=title))


def event_study_plot(
    coefs: pd.DataFrame, *,
    title: str = "Event-study: dynamic treatment effect",
    y_label: str = "Estimated effect (vs. period -1)",
    height: int = 420,
) -> go.Figure:
    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=coefs["event_time"], y=coefs["coef"],
        error_y=dict(type="data",
                     array=coefs["ci_high"] - coefs["coef"],
                     arrayminus=coefs["coef"] - coefs["ci_low"],
                     color=PALETTE["secondary"], thickness=1.4, width=4),
        mode="markers+lines",
        line=dict(color=PALETTE["secondary"], width=2),
        marker=dict(size=8, color=PALETTE["secondary"],
                    line=dict(color=PALETTE["bg"], width=1.5)),
        name="point estimate (95% CI)",
    ))
    fig.add_hline(y=0, line=dict(color=PALETTE["text_dim"],
                                  dash="dash", width=1))
    fig.add_vline(
        x=-0.5, line=dict(color=PALETTE["primary"], dash="dot", width=1.6),
        annotation_text="implementation",
        annotation_position="top right",
        annotation_font_color=PALETTE["primary"])
    fig.update_xaxes(title="Event time (periods relative to implementation)",
                     dtick=1)
    fig.update_yaxes(title=y_label)
    return _layout(fig, height=height, title=dict(text=title))


def coef_forest(
    coef_df: pd.DataFrame, *,
    title: str = "Coefficient estimates",
    height: int = 380,
) -> go.Figure:
    coef_df = coef_df.copy().iloc[::-1]
    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=coef_df["coef"],
        y=coef_df["term"],
        error_x=dict(
            type="data",
            array=coef_df["ci_high"] - coef_df["coef"],
            arrayminus=coef_df["coef"] - coef_df["ci_low"],
            color=PALETTE["secondary"], thickness=1.5, width=6,
        ),
        mode="markers",
        marker=dict(size=10, color=PALETTE["secondary"],
                    line=dict(color=PALETTE["bg"], width=1)),
        showlegend=False,
    ))
    fig.add_vline(x=0, line=dict(color=PALETTE["text_dim"],
                                  dash="dash", width=1))
    fig.update_xaxes(title="Estimated coefficient (95% CI)")
    fig.update_yaxes(title="")
    return _layout(fig, height=height, title=dict(text=title))


def montana_county_choropleth(
    counties: pd.DataFrame, *, color: str, color_label: str,
    title: str = "", height: int = 540,
) -> go.Figure:
    fig = px.choropleth_mapbox(
        counties,
        geojson="https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json",
        locations="fips",
        color=color,
        color_continuous_scale="Cividis",
        mapbox_style="carto-darkmatter",
        center={"lat": 47.0, "lon": -110.5},
        zoom=5.2,
        hover_name="county_name",
        hover_data={
            "fips": False, "rurality": True,
            "population_2020": ":,", color: ":.2f",
        },
        labels={color: color_label},
    )
    fig.update_layout(
        template=PLOT_TEMPLATE,
        margin=dict(l=0, r=0, t=42, b=0),
        height=height,
        paper_bgcolor=PALETTE["bg"],
        title=dict(text=title, x=0.0, xanchor="left",
                   font=dict(size=14, color=PALETTE["text"])),
    )
    return fig


def montana_with_hospitals(
    counties: pd.DataFrame, hospitals: pd.DataFrame, *,
    height: int = 620,
    facility_type_filter: Optional[list[str]] = None,
) -> go.Figure:
    cdf = counties.copy()
    cdf["rurality_code"] = cdf["rurality"].map({"Urban": 0, "Rural": 1, "Tribal": 2})

    if facility_type_filter:
        hdf = hospitals[hospitals["facility_type"].isin(facility_type_filter)].copy()
    else:
        hdf = hospitals.copy()

    hdf["marker_size"] = np.clip(hdf["staffed_beds"] / 4 + 6, 6, 32)
    type_color = {
        "Tertiary":       "#5BA3DA",
        "PPS":            "#7AB6E0",
        "Sole Community": "#9CC8E8",
        "CAH":            "#E0B458",
        "Tribal/IHS":     "#E15A63",
    }

    fig = go.Figure()

    # Choropleth base — subtle differentiation between rurality strata.
    fig.add_choroplethmapbox(
        geojson="https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json",
        locations=cdf["fips"],
        z=cdf["rurality_code"],
        colorscale=[[0, "#1A2030"], [0.5, "#2A2438"], [1, "#3A2030"]],
        marker_opacity=0.9,
        marker_line_width=0.5,
        marker_line_color="#0B0E13",
        showscale=False,
        text=cdf["county_name"] + " — pop " +
             cdf["population_2020"].astype(int).map("{:,}".format) +
             " (" + cdf["rurality"] + ")",
        hovertemplate="<b>%{text}</b><extra></extra>",
        name="Counties",
    )

    for ftype, sub in hdf.groupby("facility_type"):
        fig.add_scattermapbox(
            lat=sub["lat"], lon=sub["lon"],
            mode="markers",
            marker=dict(size=sub["marker_size"],
                        color=type_color.get(ftype, "#9CA3AF"),
                        opacity=0.92),
            text=sub["facility_name"]
                 + "<br>" + sub["county_name"] + " County"
                 + "<br>" + ftype + " — " + sub["staffed_beds"].astype(str)
                 + " staffed beds<br>CCN " + sub["ccn"].astype(str)
                 + "<br>" + sub["ownership"],
            hovertemplate="<b>%{text}</b><extra></extra>",
            name=ftype,
        )

    fig.update_layout(
        mapbox=dict(
            style="carto-darkmatter",
            center={"lat": 47.0, "lon": -110.0},
            zoom=5.4,
        ),
        paper_bgcolor=PALETTE["bg"],
        plot_bgcolor=PALETTE["bg"],
        font=dict(color=PALETTE["text"]),
        margin=dict(l=0, r=0, t=10, b=0),
        height=height,
        legend=dict(
            orientation="h", yanchor="bottom", y=0.0,
            xanchor="left", x=0.0,
            bgcolor="rgba(11,14,19,0.85)",
            bordercolor=PALETTE["border"], borderwidth=1,
            font=dict(color=PALETTE["text"]),
        ),
    )
    return fig


def stacked_bar(df: pd.DataFrame, x: str, ys: list[str], *,
                title: str = "", y_label: str = "", height: int = 360) -> go.Figure:
    fig = go.Figure()
    palette = [PALETTE["secondary"], "#7AB6E0", "#9CC8E8",
               PALETTE["primary"], PALETTE["accent"], "#A487E0"]
    for i, y in enumerate(ys):
        fig.add_bar(x=df[x], y=df[y], name=y,
                    marker_color=palette[i % len(palette)])
    fig.update_layout(barmode="stack")
    fig.update_yaxes(title=y_label)
    fig.update_xaxes(title=x)
    return _layout(fig, height=height, title=dict(text=title))