visualizer.py

import numpy as np
import plotly.graph_objects as go

from data_loader import ARTIFACT_COLORS, ARTIFACT_LABEL_MAP

CHANNEL_SPACING = 1.0


def build_eeg_figure(
    signal: np.ndarray,
    channel_names: list[str],
    sampling_rate_hz: float,
    start_time_s: float = 0.0,
    annotations: list[dict] | None = None,
    title: str = "EEG Signal",
) -> go.Figure:
    n_channels, n_samples = signal.shape
    time_axis = start_time_s + np.arange(n_samples) / sampling_rate_hz

    normalized = np.zeros_like(signal, dtype=np.float32)
    for i in range(n_channels):
        p2p = np.nanpercentile(signal[i], 97.5) - np.nanpercentile(signal[i], 2.5)
        if p2p > 0:
            normalized[i] = signal[i] / p2p * CHANNEL_SPACING * 0.8
        else:
            normalized[i] = 0.0

    fig = go.Figure()

    if annotations:
        _add_artifact_shapes(fig, annotations, start_time_s, time_axis[-1], n_channels)

    for i in range(n_channels):
        ch_data = normalized[i]
        offset = -i * CHANNEL_SPACING
        name = channel_names[i] if i < len(channel_names) else f"Ch{i}"
        fig.add_trace(go.Scatter(
            x=time_axis,
            y=ch_data + offset,
            mode="lines",
            name=name,
            line=dict(width=0.8, color=f"hsl({(i * 37) % 360}, 70%, 65%)"),
            hovertemplate=(
                f"<b>{name}</b><br>"
                "Time: %{x:.3f}s<br>"
                "Amplitude: %{customdata:.1f} uV<extra></extra>"
            ),
            customdata=signal[i],
        ))

    y_ticks = [-i * CHANNEL_SPACING for i in range(n_channels)]
    y_labels = [channel_names[i] if i < len(channel_names) else f"Ch{i}" for i in range(n_channels)]

    fig.update_layout(
        title=dict(text=title, font=dict(size=14)),
        xaxis=dict(
            title="Time (s)",
            showgrid=True,
            gridcolor="rgba(128,128,128,0.2)",
            zeroline=False,
            dtick=1.0,
        ),
        yaxis=dict(
            tickvals=y_ticks,
            ticktext=y_labels,
            showgrid=False,
            zeroline=False,
        ),
        height=max(500, n_channels * 30 + 120),
        margin=dict(l=110, r=20, t=50, b=50),
        plot_bgcolor="rgb(15, 15, 25)",
        paper_bgcolor="rgb(10, 10, 20)",
        font=dict(color="rgb(200, 200, 210)", size=11),
        showlegend=False,
        hovermode="x unified",
        dragmode="zoom",
    )

    return fig


def _add_artifact_shapes(
    fig: go.Figure,
    annotations: list[dict],
    view_start: float,
    view_end: float,
    n_channels: int,
) -> None:
    y_top = CHANNEL_SPACING
    y_bottom = -(n_channels - 0.5) * CHANNEL_SPACING

    for ann in annotations:
        onset = max(ann["onset_s"], view_start)
        end = min(ann["end_s"], view_end)
        if onset >= end:
            continue

        label = ann.get("label", ann.get("raw_label", ""))
        color = ann.get("color", "rgba(128, 128, 128, 0.3)")
        strong_color = color.replace("0.25)", "0.35)")

        fig.add_shape(
            type="rect",
            x0=onset, x1=end,
            y0=y_bottom, y1=y_top,
            fillcolor=strong_color,
            line=dict(width=1, color=strong_color.replace("0.35)", "0.7)")),
            layer="below",
        )

        mid_x = (onset + end) / 2
        if (end - onset) > 0.2:
            fig.add_annotation(
                x=mid_x,
                y=y_top + CHANNEL_SPACING * 0.3,
                text=f"<b>{label}</b>",
                showarrow=False,
                font=dict(size=11, color="white"),
                bgcolor=strong_color.replace("0.35)", "0.85)"),
                borderpad=3,
            )


def build_artifact_legend() -> str:
    items = []
    for label, color in ARTIFACT_COLORS.items():
        if label in ("Clean", "Background"):
            continue
        rgba_solid = color.replace("0.25)", "0.85)")
        items.append(
            f'<div style="display:flex;align-items:center;margin:4px 0;">'
            f'<span style="display:inline-block;width:16px;height:16px;'
            f'background:{rgba_solid};border-radius:3px;margin-right:8px;'
            f'flex-shrink:0;"></span>'
            f'<span style="font-size:13px;">{label}</span></div>'
        )
    return '<div style="padding:4px;">' + "".join(items) + "</div>"


def build_metadata_html(info: dict) -> str:
    rows = [
        ("Recording ID", info.get("recording_id", "N/A")),
        ("Dataset", info.get("dataset_id", "N/A")),
        ("Subject", info.get("subject", "N/A")),
        ("Session", info.get("session", "N/A")),
        ("Task", info.get("task", "N/A")),
        ("Duration", f"{info.get('duration_s', 0):.1f} s"),
        ("Channels", f"{info.get('n_channels', 0)} ({info.get('n_eeg_channels', 0)} EEG)"),
        ("Sampling Rate", f"{info.get('sampling_rate_hz', 0):.0f} Hz"),
        ("Reference", info.get("reference", "N/A")),
        ("Montage", info.get("montage_name", "N/A")),
        ("Format", info.get("archival_format", "N/A")),
        ("Roundtrip", info.get("roundtrip_class", "N/A")),
    ]

    html = '<div style="font-family:monospace; font-size:13px; line-height:1.8;">'
    for label, value in rows:
        html += (
            f'<div><span style="color:#888;min-width:130px;display:inline-block;">'
            f'{label}:</span> <span style="color:#e0e0e0;">{value}</span></div>'
        )
    html += "</div>"
    return html


def build_annotation_summary(annotations: list[dict]) -> str:
    if not annotations:
        return (
            '<div style="color:#888; font-size:13px; padding:8px;">'
            'No artifact annotations in this window. '
            'Try navigating to a different time region.</div>'
        )

    counts: dict[str, int] = {}
    total_dur: dict[str, float] = {}
    for ann in annotations:
        label = ann.get("label", "Unknown")
        counts[label] = counts.get(label, 0) + 1
        total_dur[label] = total_dur.get(label, 0) + ann.get("duration_s", 0)

    html = '<div style="font-family:monospace; font-size:13px; line-height:1.8; padding:4px;">'
    html += f'<div style="color:#fff; margin-bottom:6px;"><b>Artifacts found: {len(annotations)}</b></div>'
    for label in sorted(counts.keys()):
        color = ARTIFACT_COLORS.get(label, "rgba(128,128,128,0.5)")
        rgba_solid = color.replace("0.25)", "0.85)")
        html += (
            f'<div style="display:flex;align-items:center;margin:2px 0;">'
            f'<span style="display:inline-block;width:12px;height:12px;'
            f'background:{rgba_solid};border-radius:2px;margin-right:8px;'
            f'flex-shrink:0;"></span>'
            f'<span style="color:#e0e0e0;">{label}</span>: '
            f'<span style="color:#aaa;">{counts[label]}x, {total_dur[label]:.1f}s total</span>'
            f'</div>'
        )
    html += "</div>"
    return html