PTSRaiAgent/dashboard/components/charts.py

"""Plotly figures of the patient card. Pure builders: data in, `go.Figure` out (tested without Streamlit).

Chart rules: one y-axis per chart (scales of different range are separate charts, daily metrics are small
multiples on a shared time axis), a single neutral series colour, status colours only for thresholds.
"""

from __future__ import annotations

from collections.abc import Iterable
from datetime import date, timedelta

import plotly.graph_objects as go
from plotly.subplots import make_subplots

from app.api.schemas import SeriesPoint
from dashboard import theme
from dashboard.components.helpers import SCALE_MAX, metric_points, severity_bands, short_date

# Band shading: one warm hue with growing opacity (sequential), so "higher = worse" reads without a legend.
_BAND_ALPHA = [0.0, 0.05, 0.09, 0.14, 0.20]


def _day_ticks(days: list[date], max_ticks: int = 10) -> dict[str, object]:
    if not days:
        return {}
    first, last = min(days), max(days)
    span = (last - first).days + 1
    step = max(1, -(-span // max_ticks))
    ticks = [first + timedelta(days=i) for i in range(0, span, step)]
    return {"tickmode": "array", "tickvals": ticks, "ticktext": [short_date(d) for d in ticks]}


def _point_label(p: SeriesPoint) -> str:
    return f"{p.value:g}" if p.scheduled else f"{p.value:g} · внепл."


def _therapy_day(d: date, start: date | None) -> str:
    return f"день {(d - start).days}" if start else short_date(d)


def scale_figure(
    series: Iterable[SeriesPoint], scale_id: str, title: str, start: date | None = None
) -> go.Figure | None:
    """Total score over time with the severity gradations from config/scales; unscheduled sessions are hollow."""
    pts = metric_points(series, scale_id)
    if not pts:
        return None
    top = SCALE_MAX.get(scale_id, max(p.value for p in pts))
    fig = go.Figure()
    bands = severity_bands(scale_id)
    for i, b in enumerate(bands):
        alpha = _BAND_ALPHA[min(i, len(_BAND_ALPHA) - 1)] if len(bands) > 3 else [0.0, 0.08, 0.18][min(i, 2)]
        fig.add_hrect(y0=b.lo, y1=b.hi + 1 if b.hi < top else top, line_width=0, fillcolor=f"rgba(214,69,69,{alpha})")
        fig.add_annotation(
            x=0,
            xref="paper",
            xanchor="left",
            y=(b.lo + min(b.hi + 1, top)) / 2,
            text=b.label,
            showarrow=False,
            font={"size": 10, "color": "rgba(122,130,140,0.85)"},
        )
    xs = [p.day for p in pts]
    hover = [f"{short_date(p.day)} · {_therapy_day(p.day, start)}" for p in pts]
    fig.add_trace(
        go.Scatter(
            x=xs,
            y=[p.value for p in pts],
            mode="lines+markers+text",
            line={"color": theme.SERIES, "width": 2},
            marker={
                "size": [9 if p.scheduled else 13 for p in pts],
                "symbol": ["circle" if p.scheduled else "diamond-open" for p in pts],
                "color": theme.SERIES,
                "line": {"width": [2 if p.scheduled else 2.5 for p in pts], "color": "#FFFFFF"},
            },
            text=[_point_label(p) if i == len(pts) - 1 else "" for i, p in enumerate(pts)],
            textposition="top left",
            textfont={"color": theme.INK, "size": 12},
            customdata=[[h, "плановый" if p.scheduled else "внеплановый"] for h, p in zip(hover, pts, strict=True)],
            hovertemplate="%{customdata[0]}<br><b>%{y:g}</b> баллов · %{customdata[1]}<extra></extra>",
            name=title,
        )
    )
    pad = timedelta(days=1)
    fig.update_layout(
        theme.plotly_layout(title={"text": title, "font": {"size": 14, "color": theme.INK}, "x": 0}, height=260),
        hovermode="closest",
    )
    fig.update_xaxes(**theme.AXIS, **_day_ticks(xs, 4), tickangle=0, range=[min(xs) - pad, max(xs) + pad])
    fig.update_yaxes(**theme.AXIS, range=[0, top], fixedrange=True)
    return fig


def wellbeing_figure(series: Iterable[SeriesPoint], start: date | None = None) -> go.Figure | None:
    """Daily check-in as three small multiples on one time axis: mood 0–10, sleep hours, alcohol 0/1/2."""
    data = list(series)
    mood, sleep, alcohol = (metric_points(data, m) for m in ("mood", "sleep_hours", "alcohol"))
    if not (mood or sleep or alcohol):
        return None
    days = sorted({p.day for p in mood + sleep + alcohol})
    fig = make_subplots(
        rows=3,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.09,
        row_heights=[0.42, 0.33, 0.25],
        subplot_titles=("Настроение, 0–10", "Сон, часов", "Алкоголь"),
    )
    red_line = {"color": theme.TONES["red"].solid, "width": 1, "dash": "dot"}

    def hov(p: SeriesPoint) -> str:
        return f"{short_date(p.day)} · {_therapy_day(p.day, start)}"

    if mood:
        # Mood ≤ 3 for three days triggers an unscheduled PHQ-9 (config/rules.yaml) — shade that zone.
        fig.add_hrect(y0=0, y1=3, line_width=0, fillcolor="rgba(214,69,69,0.08)", row=1, col=1)
        fig.add_trace(
            go.Scatter(
                x=[p.day for p in mood],
                y=[p.value for p in mood],
                mode="lines+markers",
                line={"color": theme.SERIES, "width": 2},
                marker={"size": 8, "color": theme.SERIES, "line": {"width": 2, "color": "#FFFFFF"}},
                customdata=[hov(p) for p in mood],
                hovertemplate="%{customdata}<br>настроение <b>%{y:g}</b>/10<extra></extra>",
            ),
            row=1,
            col=1,
        )
    if sleep:
        fig.add_trace(
            go.Bar(
                x=[p.day for p in sleep],
                y=[p.value for p in sleep],
                marker={"color": [theme.TONES["red"].solid if p.value < 4 else "#9DBFEA" for p in sleep]},
                marker_cornerradius=4,
                customdata=[hov(p) for p in sleep],
                hovertemplate="%{customdata}<br>сон ≈ <b>%{y:g}</b> ч<extra></extra>",
            ),
            row=2,
            col=1,
        )
        fig.add_hline(y=4, line=red_line, row=2, col=1)
    if alcohol:
        alc_label = {0: "нет", 1: "1–2 порции", 2: "больше"}
        fig.add_trace(
            go.Bar(
                x=[p.day for p in alcohol],
                y=[max(p.value, 0.08) for p in alcohol],
                marker={
                    "color": [
                        {0: theme.LINE, 1: theme.TONES["yellow"].solid}.get(int(p.value), theme.TONES["red"].solid)
                        for p in alcohol
                    ]
                },
                marker_cornerradius=4,
                customdata=[[hov(p), alc_label.get(int(p.value), "?")] for p in alcohol],
                hovertemplate="%{customdata[0]}<br>алкоголь: <b>%{customdata[1]}</b><extra></extra>",
            ),
            row=3,
            col=1,
        )
    fig.update_layout(theme.plotly_layout(height=430, bargap=0.25, hovermode="x"))
    fig.update_annotations(font={"size": 12, "color": theme.INK_2}, x=0, xanchor="left")
    pad = timedelta(hours=14)
    fig.update_xaxes(**theme.AXIS, range=[days[0] - pad, days[-1] + pad])
    fig.update_xaxes(**_day_ticks(days), row=3, col=1)
    fig.update_yaxes(**theme.AXIS, fixedrange=True)
    fig.update_yaxes(range=[0, 10.5], tickvals=[0, 3, 5, 10], row=1, col=1)
    fig.update_yaxes(range=[0, 9], tickvals=[0, 4, 8], row=2, col=1)
    fig.update_yaxes(range=[0, 2.2], tickvals=[0, 1, 2], ticktext=["0", "1–2", "3+"], row=3, col=1)
    return fig