analysis/viz/umap_plot.py

from __future__ import annotations

import copy
import io
import base64
from typing import Any, List, Mapping, Tuple

import numpy as np
from PIL import Image
import plotly.graph_objects as go

from log_config import get_logger

logger = get_logger(__name__)


SELECTED_FEATURE_TRACE_NAME = "selected_feature"
NO_ACTIVE_FEATURES_ANNOTATION = "No active features on this image."

CATEGORY_PALETTE = [
    '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#17becf', '#bcbd22',
    '#8c564b', '#e377c2', '#7f7f7f', '#9467bd', '#aec7e8', '#ffbb78', '#98df8a',
]


def _img_to_base64(img: Image.Image, fmt: str = 'JPEG', quality: int = 80) -> str:
    buf = io.BytesIO()
    img.save(buf, format=fmt, quality=quality)
    return base64.b64encode(buf.getvalue()).decode('utf-8')


def make_thumb_b64(img_idx: int, kadid_ds: Any, size: int = 140) -> str | None:
    """Return base64 thumbnail for image index using dataset object providing `images`.

    `kadid_ds` is expected to expose sequence-like `.images` where each entry is a path.
    """
    try:
        p = str(kadid_ds.images[int(img_idx)])
        img = Image.open(p).convert('RGB').resize((size, size), Image.LANCZOS)
        return _img_to_base64(img)
    except Exception as e:
        logger.error('Error occurred while processing image %s: %s', img_idx, e)
        return None


def _align_umap_series_length(
    emb: np.ndarray,
    labels: np.ndarray,
    customdata: np.ndarray | list | None,
    hovertext: list[str] | None,
    marker_sizes: np.ndarray | None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray | list | None, list[str] | None, np.ndarray | None]:
    n = int(emb.shape[0])
    labels_arr = np.asarray(labels)
    if labels_arr.shape[0] == n:
        return emb, labels_arr, customdata, hovertext, marker_sizes

    n = min(n, int(labels_arr.shape[0]))
    emb = emb[:n]
    labels_arr = labels_arr[:n]
    if customdata is not None:
        customdata = list(customdata)[:n]
    if hovertext is not None:
        hovertext = list(hovertext)[:n]
    if marker_sizes is not None:
        marker_sizes = np.asarray(marker_sizes)[:n]
    return emb, labels_arr, customdata, hovertext, marker_sizes


def build_categorical_traces(
    emb: np.ndarray,
    labels: np.ndarray,
    selected_category: str | None = None,
    palette: List[str] | None = None,
    customdata: np.ndarray | list | None = None,
    hovertext: list[str] | None = None,
    marker_sizes: np.ndarray | None = None,
    legendonly_categories: frozenset[str] | None = None,
) -> Tuple[list[go.Scatter], list[str]]:
    emb, labels, customdata, hovertext, marker_sizes = _align_umap_series_length(
        emb, labels, customdata, hovertext, marker_sizes
    )
    palette = CATEGORY_PALETTE if palette is None else palette
    legendonly_categories = legendonly_categories or frozenset()
    unique_cats = sorted(str(x) for x in np.unique(labels))
    selected_category = None if selected_category in {None, ''} else str(selected_category)
    if selected_category not in unique_cats:
        selected_category = None

    traces: list[go.Scatter] = []
    for i, cat in enumerate(unique_cats):
        if selected_category is not None and cat != selected_category:
            continue

        mask = np.asarray([str(v) == cat for v in labels], dtype=bool)
        if not np.any(mask):
            continue
        img_indices = np.flatnonzero(mask).astype(np.int32)

        # Prepare customdata as list of [img_idx, label]
        if customdata is None:
            trace_customdata = [[int(idx), str(labels[int(idx)])] for idx in img_indices]
        else:
            trace_customdata = [customdata[int(idx)] for idx in img_indices]

        trace_hovertext = None
        if hovertext is not None:
            trace_hovertext = [hovertext[int(idx)] for idx in img_indices]

        trace_sizes = None
        if marker_sizes is not None:
            trace_sizes = [float(marker_sizes[int(idx)]) for idx in img_indices]

        color = palette[i % len(palette)]
        trace_visible: bool | str = 'legendonly' if cat in legendonly_categories else True
        traces.append(
            go.Scatter(
                x=emb[mask, 0],
                y=emb[mask, 1],
                mode='markers',
                name=cat,
                legendgroup=cat,
                marker=dict(size=trace_sizes if trace_sizes is not None else 7, color=color, opacity=0.85, line=dict(width=0)),
                customdata=trace_customdata,
                hoverinfo='none',
                hovertext=trace_hovertext if trace_hovertext is not None else [f"image_idx={int(idx)} | category={cat}" for idx in img_indices],
                showlegend=True,
                visible=trace_visible,
            )
        )

    return traces, unique_cats


def build_umap_figure(
    embedding_2d: np.ndarray,
    labels: np.ndarray,
    *,
    title: str,
    selected_category: str | None = None,
    palette: List[str] | None = None,
    customdata: np.ndarray | list | None = None,
    hovertext: list[str] | None = None,
    marker_sizes: np.ndarray | None = None,
    legendonly_categories: frozenset[str] | None = None,
) -> Tuple[go.Figure, list[str]]:
    traces, unique_cats = build_categorical_traces(
        embedding_2d,
        labels,
        selected_category=selected_category,
        palette=palette,
        customdata=customdata,
        hovertext=hovertext,
        marker_sizes=marker_sizes,
        legendonly_categories=legendonly_categories,
    )
    fig = go.Figure(data=traces)
    fig.update_layout(
        title=dict(text=title, x=0.01, xanchor='left', y=0.97, yanchor='top', font=dict(size=16)),
        xaxis_title='UMAP 1',
        yaxis_title='UMAP 2',
        legend=dict(
            orientation='h',
            yanchor='top',
            y=-0.16,
            xanchor='left',
            x=0.0,
            title_text='',
            font=dict(size=11),
            tracegroupgap=8,
        ),
        margin=dict(l=40, r=24, t=96, b=92),
        hovermode='closest',
        autosize=True,
    )
    return fig, unique_cats


def _decode_plotly_array(values: Any) -> list[float]:
    if values is None:
        return []
    if isinstance(values, (list, tuple, np.ndarray)):
        return [float(v) for v in values]
    if isinstance(values, Mapping) and "bdata" in values:
        dtype_map = {"f8": np.float64, "f4": np.float32, "i4": np.int32, "i2": np.int16}
        dtype = dtype_map.get(str(values.get("dtype")), np.float64)
        raw = base64.b64decode(values["bdata"])
        return np.frombuffer(raw, dtype=dtype).astype(float).tolist()
    return [float(values)]


def _find_feature_point_in_figure(
    traces: list[go.Scatter],
    highlight_feature_id: int,
) -> tuple[float, float, list[Any]] | None:
    for trace in traces:
        if trace.name == SELECTED_FEATURE_TRACE_NAME:
            continue
        customdata = trace.customdata
        if customdata is None:
            continue
        x_vals = _decode_plotly_array(trace.x)
        y_vals = _decode_plotly_array(trace.y)
        if not x_vals or not y_vals:
            continue
        for i, row in enumerate(customdata):
            if not isinstance(row, (list, tuple, np.ndarray)) or len(row) < 1:
                continue
            try:
                if int(row[0]) != highlight_feature_id:
                    continue
            except (TypeError, ValueError):
                continue
            if i >= len(x_vals) or i >= len(y_vals):
                continue
            return x_vals[i], y_vals[i], list(row)
    return None


def filter_umap_figure_to_feature_ids(
    figure: go.Figure | dict[str, Any],
    active_feature_ids: frozenset[int],
) -> dict[str, Any]:
    """Keep only UMAP points whose feature id (``customdata[0]``) is in ``active_feature_ids``."""
    if isinstance(figure, go.Figure):
        fig = copy.deepcopy(figure)
    else:
        fig = go.Figure(copy.deepcopy(figure))

    active_ids = {int(fid) for fid in active_feature_ids}
    filtered_traces: list[go.Scatter] = []

    for trace in fig.data:
        if trace.name == SELECTED_FEATURE_TRACE_NAME:
            continue
        customdata = trace.customdata
        if customdata is None:
            continue

        x_vals = _decode_plotly_array(trace.x)
        y_vals = _decode_plotly_array(trace.y)
        if not x_vals or not y_vals:
            continue

        sizes = trace.marker.size if trace.marker else None
        if isinstance(sizes, (list, tuple, np.ndarray)):
            size_list = [float(s) for s in sizes]
        else:
            size_list = None

        new_x: list[float] = []
        new_y: list[float] = []
        new_cd: list[Any] = []
        new_sizes: list[float] = []

        for i, row in enumerate(customdata):
            if not isinstance(row, (list, tuple, np.ndarray)) or len(row) < 1:
                continue
            try:
                feature_id = int(row[0])
            except (TypeError, ValueError):
                continue
            if feature_id not in active_ids:
                continue
            if i >= len(x_vals) or i >= len(y_vals):
                continue
            new_x.append(x_vals[i])
            new_y.append(y_vals[i])
            new_cd.append(list(row))
            if size_list is not None and i < len(size_list):
                new_sizes.append(size_list[i])

        if not new_x:
            continue

        if trace.marker is not None:
            marker = trace.marker.to_plotly_json()
        else:
            marker = {"size": 7}
        if new_sizes:
            marker["size"] = new_sizes

        filtered_traces.append(
            go.Scatter(
                x=new_x,
                y=new_y,
                mode=trace.mode,
                name=trace.name,
                legendgroup=trace.legendgroup,
                marker=marker,
                customdata=new_cd,
                hoverinfo=trace.hoverinfo,
                hovertext=trace.hovertext,
                showlegend=trace.showlegend,
                visible=trace.visible,
            )
        )

    if isinstance(figure, go.Figure):
        base_layout = copy.deepcopy(figure.layout)
    else:
        base_layout = copy.deepcopy(figure.get("layout", {}))

    filtered_fig = go.Figure(data=filtered_traces, layout=base_layout)
    if not filtered_traces:
        filtered_fig.update_layout(
            annotations=[
                dict(
                    text=NO_ACTIVE_FEATURES_ANNOTATION,
                    xref="paper",
                    yref="paper",
                    x=0.5,
                    y=0.5,
                    showarrow=False,
                    font=dict(size=14),
                )
            ]
        )
    else:
        filtered_fig.update_layout(annotations=[])

    return filtered_fig.to_plotly_json()


def apply_feature_umap_highlight(
    figure: go.Figure | dict[str, Any],
    highlight_feature_id: int | None,
) -> dict[str, Any]:
    """Add or remove a star marker for the selected feature on a feature UMAP figure."""
    if isinstance(figure, go.Figure):
        fig = copy.deepcopy(figure)
    else:
        fig = go.Figure(copy.deepcopy(figure))

    fig.data = tuple(trace for trace in fig.data if trace.name != SELECTED_FEATURE_TRACE_NAME)

    if highlight_feature_id is not None:
        try:
            feature_id = int(highlight_feature_id)
        except (TypeError, ValueError):
            feature_id = None

        if feature_id is not None:
            point = _find_feature_point_in_figure(list(fig.data), feature_id)
            if point is not None:
                x_val, y_val, row = point
                fig.add_trace(
                    go.Scatter(
                        x=[x_val],
                        y=[y_val],
                        mode="markers",
                        name=SELECTED_FEATURE_TRACE_NAME,
                        legendgroup=SELECTED_FEATURE_TRACE_NAME,
                        showlegend=False,
                        marker=dict(
                            symbol="star",
                            size=18,
                            color="#f59e0b",
                            line=dict(width=2, color="#111827"),
                        ),
                        customdata=[row],
                        hoverinfo="none",
                    )
                )

    return fig.to_plotly_json()


def plot_categorical_scatter_matplotlib(
    emb: np.ndarray,
    labels: np.ndarray,
    title: str,
    legend_title: str = 'category',
    exclude_categories: frozenset[str] | None = None,
) -> None:
    import matplotlib.pyplot as plt

    def _label_str(value: object) -> str:
        if value is None:
            return 'unknown'
        if isinstance(value, float) and np.isnan(value):
            return 'unknown'
        text = str(value).strip()
        return 'unknown' if text.lower() in {'', 'nan', 'none'} else text

    exclude = exclude_categories or frozenset()
    if exclude:
        keep_mask = np.array([_label_str(v) not in exclude for v in labels], dtype=bool)
        emb = np.asarray(emb)[keep_mask]
        labels = np.asarray(labels)[keep_mask]

    unique_cats = sorted(dict.fromkeys(_label_str(x) for x in labels))
    cmap = plt.get_cmap('tab10')
    n = max(1, len(unique_cats))
    colors = [cmap(i % cmap.N) for i in range(n)]
    plt.figure(figsize=(10, 7))
    for i, cat in enumerate(unique_cats):
        mask = np.array([_label_str(v) == cat for v in labels])
        if not np.any(mask):
            continue
        plt.scatter(
            emb[mask, 0],
            emb[mask, 1],
            c=[colors[i]],
            label=cat,
            s=28,
            alpha=0.85,
            edgecolors='none',
        )
    plt.title(title)
    plt.xlabel('UMAP 1')
    plt.ylabel('UMAP 2')
    plt.grid(True, alpha=0.3)
    plt.legend(title=legend_title, bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.tight_layout()
    plt.show()