sync_pilot/dashboard/plots.py

"""Plotly figure builders for the sync_pilot dashboard.

All chart styling lives here so pages stay declarative. Plotly only (no
matplotlib, no Altair). We use ``plotly.graph_objects.Figure`` for full
control over chrome — anything the brief mandates (no axis titles,
threshold lines, fixed [0,1] x-range on tag bars) is enforced once
in the helpers below.
"""

from __future__ import annotations

from collections import Counter
from typing import Any, Iterable

import plotly.graph_objects as go

from sync_pilot.dashboard.styles import (
    CATEGORY_COLORS,
    CLAP_COLOR,
    MAEST_COLOR,
    MUQ_COLOR,
    MUQ_REVIEW_COLOR,
    PASST_COLOR,
    STEM_COLOR,
    TAXONOMY_ADAPTER_COLOR,
    THEME_COLOR,
)

_MAEST_DISPLAY_MIN_SCORE = 0.15
_CLAP_DISPLAY_TOP_K = 3
_CLAP_MOOD_DISPLAY_TOP_K = 2
_SOURCE_BUCKET_ORDER = [
    "maest",
    "clap-zs",
    "muq-probe-v1",
    "muq-probe-expanded-v1",
    "passt-probe-expanded-v1",
    "taxonomy-adapter-v1",
    "muq-probe-review-v1",
    "lyrics-llm",
    "stem-window-v1",
]
_SOURCE_LABELS = {
    "maest": "MAEST",
    "clap-zs": "CLAP",
    "muq-probe-v1": "MuQ probe",
    "muq-probe-expanded-v1": "MuQ expanded",
    "passt-probe-expanded-v1": "PaSST expanded",
    "taxonomy-adapter-v1": "taxonomy adapter",
    "muq-probe-review-v1": "review probe",
    "lyrics-llm": "lyrics LLM",
    "stem-window-v1": "stem windows",
}
_SOURCE_COLORS = {
    "maest": MAEST_COLOR,
    "clap-zs": CLAP_COLOR,
    "muq-probe-v1": MUQ_COLOR,
    "muq-probe-expanded-v1": MUQ_COLOR,
    "passt-probe-expanded-v1": PASST_COLOR,
    "taxonomy-adapter-v1": TAXONOMY_ADAPTER_COLOR,
    "muq-probe-review-v1": MUQ_REVIEW_COLOR,
    "lyrics-llm": THEME_COLOR,
    "stem-window-v1": STEM_COLOR,
}


def _empty_figure(message: str) -> go.Figure:
    """Placeholder figure for the "no data" case — keeps page layout stable."""
    fig = go.Figure()
    fig.add_annotation(
        text=message,
        xref="paper",
        yref="paper",
        x=0.5,
        y=0.5,
        showarrow=False,
        font={"color": "#9CA3AF", "size": 13},
    )
    fig.update_layout(
        height=220,
        margin={"l": 10, "r": 10, "t": 10, "b": 10},
        plot_bgcolor="white",
        paper_bgcolor="white",
        xaxis={"visible": False},
        yaxis={"visible": False},
    )
    return fig


def _source_color(source: str) -> str:
    """Map a tag source to its canonical color."""
    bucket = _source_bucket(source) or source
    return _SOURCE_COLORS.get(bucket, CATEGORY_COLORS["other"])


def _source_bucket(source: str) -> str | None:
    if source.startswith("clap"):
        return "clap-zs"
    if source.startswith("maest"):
        return "maest"
    if source.startswith("muq-probe-review"):
        return "muq-probe-review-v1"
    if source.startswith("muq-probe-expanded"):
        return "muq-probe-expanded-v1"
    if source.startswith("muq-probe"):
        return "muq-probe-v1"
    if source.startswith("passt-probe-expanded"):
        return "passt-probe-expanded-v1"
    if source.startswith("taxonomy-adapter"):
        return "taxonomy-adapter-v1"
    if source.startswith("lyrics"):
        return "lyrics-llm"
    if source.startswith("muq"):
        return "muq-mulan-zs"
    if source.startswith("stem-window"):
        return "stem-window-v1"
    return None


def _source_label(source: str) -> str:
    return _SOURCE_LABELS.get(source, source)


def _catalog_display_tags(tags: list[dict[str, Any]], category: str) -> list[dict[str, Any]]:
    rows = [t for t in tags if t.get("category") == category]
    filtered: list[dict[str, Any]] = []
    for tag in rows:
        source = str(tag.get("source", ""))
        score = float(tag.get("score") or 0.0)
        if category == "genre" and source.startswith("maest") and score < _MAEST_DISPLAY_MIN_SCORE:
            continue
        filtered.append(tag)

    capped: list[dict[str, Any]] = []
    clap_rows = [t for t in filtered if str(t.get("source", "")).startswith("clap")]
    other_rows = [t for t in filtered if not str(t.get("source", "")).startswith("clap")]
    clap_rows.sort(key=lambda t: float(t.get("score") or 0.0), reverse=True)
    clap_limit = _CLAP_MOOD_DISPLAY_TOP_K if category == "mood" else _CLAP_DISPLAY_TOP_K
    capped.extend(clap_rows[:clap_limit])
    capped.extend(other_rows)
    return capped


def tag_bar_chart(
    tags: list[dict[str, Any]],
    *,
    title: str,
    top_k: int = 10,
    threshold_lines: Iterable[tuple[float, str]] = (
        (0.30, "CLAP τ"),
        (0.15, "MAEST τ"),
    ),
    height: int = 280,
) -> go.Figure:
    """Horizontal bar chart of a single category's tags, source-colored.

    ``tags`` is a list of ``{name, score, source}`` dicts (already filtered
    to the category of interest by the caller). Score range is **always**
    locked to [0, 1] so the eye can compare confidence across small-
    multiples without re-anchoring.

    Threshold lines are dashed verticals at the relevant source cutoffs; pass
    ``threshold_lines=()`` to suppress.
    """
    if not tags:
        return _empty_figure("no tags above threshold")

    # Sort by descending score so the most-confident tag sits at the top
    # of the horizontal bar chart (Plotly's y-axis grows upward).
    sorted_tags = sorted(tags, key=lambda t: t.get("score", 0.0), reverse=True)[:top_k]
    sorted_tags = list(reversed(sorted_tags))  # Plotly draws bottom-up

    names = [t.get("name", "?") for t in sorted_tags]
    scores = [float(t.get("score", 0.0)) for t in sorted_tags]
    sources = [t.get("source", "") for t in sorted_tags]
    evidence = [t.get("evidence", "") for t in sorted_tags]
    colors = [_source_color(s) for s in sources]

    # customdata stacks (source, evidence) per bar so the hovertemplate can
    # surface the LLM-extracted evidence quote on theme tags without
    # cluttering the genre/mood/instrument/vocal hovers (where evidence is
    # always empty string and the conditional below skips its line).
    customdata = [[src, ev] for src, ev in zip(sources, evidence)]
    has_any_evidence = any(ev for ev in evidence)
    hovertemplate = (
        "<b>%{y}</b><br>score %{x:.3f}<br>source %{customdata[0]}"
        + ("<br>evidence: %{customdata[1]}" if has_any_evidence else "")
        + "<extra></extra>"
    )

    fig = go.Figure(
        go.Bar(
            x=scores,
            y=names,
            orientation="h",
            marker={"color": colors, "line": {"width": 0}},
            text=[f"{s:.2f}" for s in scores],
            textposition="outside",
            cliponaxis=False,
            hovertemplate=hovertemplate,
            customdata=customdata,
        )
    )

    for x, _label in threshold_lines:
        fig.add_vline(
            x=x,
            line_dash="dash",
            line_color="#9CA3AF",
            opacity=0.6,
            line_width=1,
        )

    fig.update_layout(
        title={
            "text": title,
            "x": 0.0,
            "xanchor": "left",
            "font": {"size": 14, "color": "#1F2933"},
        },
        height=height,
        margin={"l": 10, "r": 50, "t": 40, "b": 20},
        plot_bgcolor="white",
        paper_bgcolor="white",
        showlegend=False,
        font={"size": 12},
        xaxis={
            "range": [0, 1.05],
            "showgrid": False,
            "zeroline": False,
            "showline": False,
            "ticks": "",
            "tickvals": [0.0, 0.25, 0.5, 0.75, 1.0],
            "tickfont": {"color": "#6B7280", "size": 10},
        },
        yaxis={
            "showgrid": False,
            "zeroline": False,
            "showline": False,
            "ticks": "",
        },
    )
    return fig


def catalog_top_tags_chart(
    tracks: list[dict[str, Any]],
    *,
    category: str,
    title: str,
    top_k: int = 10,
    height: int = 300,
) -> go.Figure:
    """Top-K tags across the whole catalog for one category, source-split.

    Aggregates per-source counts after applying the same display precision
    policy as the track-level charts, so stale low-score tags do not dominate
    the catalog view.

    A grouped horizontal bar with two stacks per term (MAEST + CLAP) would
    be the strictest reading, but for top-10 readability we instead emit
    one bar per (term, source) pair so terms with both-source presence
    appear twice — visually showing the cross-source agreement (or lack
    thereof) without forcing same-axis alignment of disagreeing terms.
    """
    per_source: dict[str, Counter[str]] = {
        source: Counter() for source in _SOURCE_BUCKET_ORDER
    }
    for tr in tracks:
        for tag in _catalog_display_tags(tr.get("tags", []) or [], category):
            src = tag.get("source", "")
            name = tag.get("name", "")
            bucket = _source_bucket(str(src))
            if bucket and name:
                if bucket not in per_source:
                    per_source[bucket] = Counter()
                per_source[bucket][name] += 1

    # Build a single combined list of (count, name, source) and keep the
    # top_k overall.
    combined: list[tuple[int, str, str]] = []
    for src, counter in per_source.items():
        for name, count in counter.items():
            combined.append((count, name, src))
    if not combined:
        return _empty_figure(f"no {category} tags found")

    combined.sort(key=lambda r: r[0], reverse=True)
    combined = combined[:top_k]
    combined.reverse()  # plotly draws bottom-up

    fig = go.Figure()
    fig.add_trace(
        go.Bar(
            x=[c for c, _, _ in combined],
            y=[f"{n}  ·  {_source_label(s)}" for _, n, s in combined],
            orientation="h",
            marker={"color": [_source_color(s) for _, _, s in combined]},
            text=[str(c) for c, _, _ in combined],
            textposition="outside",
            cliponaxis=False,
            hovertemplate="<b>%{y}</b><br>%{x} tracks<extra></extra>",
        )
    )
    fig.update_layout(
        title={
            "text": title,
            "x": 0.0,
            "xanchor": "left",
            "font": {"size": 14, "color": "#1F2933"},
        },
        height=height,
        margin={"l": 10, "r": 40, "t": 40, "b": 20},
        plot_bgcolor="white",
        paper_bgcolor="white",
        showlegend=False,
        font={"size": 12},
        xaxis={
            "showgrid": False,
            "zeroline": False,
            "showline": False,
            "ticks": "",
            "tickfont": {"color": "#6B7280", "size": 10},
        },
        yaxis={
            "showgrid": False,
            "zeroline": False,
            "showline": False,
            "ticks": "",
        },
    )
    return fig


def language_distribution_chart(
    distribution: dict[str, int],
    *,
    title: str = "Lyrics language distribution",
    height: int = 220,
) -> go.Figure:
    """Tiny horizontal bar chart over the transcription summary's language counts."""
    if not distribution:
        return _empty_figure("no language data")
    items = sorted(distribution.items(), key=lambda kv: kv[1])
    fig = go.Figure(
        go.Bar(
            x=[c for _, c in items],
            y=[lang for lang, _ in items],
            orientation="h",
            marker={"color": CLAP_COLOR},
            text=[str(c) for _, c in items],
            textposition="outside",
            cliponaxis=False,
        )
    )
    fig.update_layout(
        title={"text": title, "x": 0.0, "xanchor": "left", "font": {"size": 13}},
        height=height,
        margin={"l": 10, "r": 30, "t": 40, "b": 20},
        plot_bgcolor="white",
        paper_bgcolor="white",
        showlegend=False,
        xaxis={
            "showgrid": False,
            "zeroline": False,
            "showline": False,
            "ticks": "",
            "tickfont": {"color": "#6B7280", "size": 10},
        },
        yaxis={"showgrid": False, "zeroline": False, "showline": False, "ticks": ""},
    )
    return fig