src/plot_config.py

"""Shared Plotly chart configuration and color constants for the Pangenome Atlas.

Every chart-building callback should import from this module to ensure
visual consistency across all plots.

Exports
-------
COLORS           – dict of semantic color tokens.
COUNTRY_COLORS   – per-country color mapping (17 countries + Unknown).
CLUSTER_COLORS   – ordered list for the 3 optimal clusters.
PLOTLY_TEMPLATE  – light-background Plotly layout dict.
PLOTLY_TEMPLATE_DARK – dark-background variant for globe / 3-D views.
apply_template() – apply the atlas template to any ``plotly.graph_objects.Figure``.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import plotly.graph_objects as go  # pragma: no cover

# =====================================================================
# Color Palette
# =====================================================================

COLORS: dict[str, str] = {
    # Gene classification
    "core": "#2E7D32",           # forest green
    "shell": "#F9A825",          # warm amber
    "cloud": "#C62828",          # warm red
    # UI accents
    "selected": "#D4A017",       # gold — user's selected line
    "accent": "#1565C0",         # deep blue — links / actions
    # Backgrounds
    "bg_dark": "#1a2332",        # dark navy — globe / hero header
    "bg_light": "#FAFAF5",       # warm white — page background
    "bg_card": "#FFFFFF",        # card surface
    # Text
    "text_primary": "#1A1A1A",
    "text_secondary": "#757575",
    "text_muted": "#94a3b8",
    # Structural
    "border": "#E0E0E0",
    "grid_faint": "#F0F0E8",
}

# Per-country colors used by the globe, UMAP, and any country legend.
# 17 countries + Unknown.
COUNTRY_COLORS: dict[str, str] = {
    "India": "#2E7D32",
    "Philippines": "#1565C0",
    "Kenya": "#C62828",
    "Nepal": "#D4A017",
    "Myanmar": "#6A1B9A",
    "Uganda": "#00838F",
    "Zaire": "#E65100",
    "Indonesia": "#455A64",
    "Jamaica": "#AD1457",
    "South_Africa": "#1B5E20",
    "Puerto_Rico": "#0277BD",
    "Sierra_Leone": "#BF360C",
    "Nigeria": "#827717",
    "Malawi": "#4E342E",
    "Italy": "#283593",
    "Sri_Lanka": "#00695C",
    "Thailand": "#FF6F00",
    "Unknown": "#9E9E9E",
}

# Cluster colors for the 3 optimal clusters (silhouette = 0.649).
CLUSTER_COLORS: list[str] = ["#2E7D32", "#1565C0", "#C62828"]

# =====================================================================
# Shared font stack (matches the Gradio theme)
# =====================================================================

_FONT_STACK = "-apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif"

# =====================================================================
# Plotly Layout Templates
# =====================================================================

PLOTLY_TEMPLATE: dict = dict(
    layout=dict(
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(
            family=_FONT_STACK,
            color=COLORS["text_primary"],
            size=13,
        ),
        title=dict(
            font=dict(size=18, color=COLORS["text_primary"]),
            x=0.02,
            xanchor="left",
        ),
        xaxis=dict(
            showgrid=False,
            zeroline=False,
            showline=False,
        ),
        yaxis=dict(
            showgrid=False,
            zeroline=False,
            showline=False,
        ),
        margin=dict(l=30, r=20, t=50, b=30),
        hoverlabel=dict(
            bgcolor="white",
            font_size=13,
            font_family=_FONT_STACK,
            bordercolor=COLORS["border"],
        ),
        legend=dict(
            bgcolor="rgba(255,255,255,0.85)",
            bordercolor="rgba(0,0,0,0)",
            font=dict(size=12),
        ),
    )
)

# Dark variant for 3-D plots, globe visualisations, and hero panels.
PLOTLY_TEMPLATE_DARK: dict = dict(
    layout=dict(
        paper_bgcolor=COLORS["bg_dark"],
        plot_bgcolor=COLORS["bg_dark"],
        font=dict(
            family=_FONT_STACK,
            color="#E0E0E0",
            size=13,
        ),
        title=dict(
            font=dict(size=18, color="#E0E0E0"),
            x=0.02,
            xanchor="left",
        ),
        xaxis=dict(
            showgrid=False,
            zeroline=False,
            showline=False,
            color="#E0E0E0",
        ),
        yaxis=dict(
            showgrid=False,
            zeroline=False,
            showline=False,
            color="#E0E0E0",
        ),
        margin=dict(l=30, r=20, t=50, b=30),
        hoverlabel=dict(
            bgcolor="#263245",
            font_size=13,
            font_family=_FONT_STACK,
            bordercolor="#3a4a60",
        ),
        legend=dict(
            bgcolor="rgba(26,35,50,0.85)",
            bordercolor="rgba(0,0,0,0)",
            font=dict(size=12, color="#E0E0E0"),
        ),
    )
)

# Modebar buttons to hide for a cleaner look.
_MODEBAR_REMOVE = [
    "toImage",
    "zoom2d",
    "pan2d",
    "select2d",
    "lasso2d",
    "autoScale2d",
    "resetScale2d",
    "zoomIn2d",
    "zoomOut2d",
]


def apply_template(fig: "go.Figure", dark: bool = False) -> "go.Figure":
    """Apply the Atlas Plotly template to *fig* and hide the modebar.

    Parameters
    ----------
    fig : plotly.graph_objects.Figure
        The figure to style in-place.
    dark : bool, optional
        If ``True``, use the dark-background variant suitable for globe
        and 3-D scenes.  Defaults to ``False``.

    Returns
    -------
    plotly.graph_objects.Figure
        The same figure, for chaining.
    """
    template = PLOTLY_TEMPLATE_DARK if dark else PLOTLY_TEMPLATE
    fig.update_layout(**template["layout"])
    fig.update_layout(modebar=dict(remove=_MODEBAR_REMOVE))
    return fig