visualization/timeline.py

# Created: 2026-02-18
# Purpose: P(AI) per-segment timeline bar chart with waveform (plotly)
# Dependencies: plotly, numpy

"""Per-segment (chunk) AI probability timeline visualization with waveform."""

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

from config import CHUNK_SEC, SR, CHUNK_SAMPLES


def plot_timeline(
    chunk_probs: list[float],
    waveform: np.ndarray = None,
    chunk_metadata: list[dict] = None,
    weighted_median: float = None
) -> go.Figure:
    """Per-chunk P(AI) timeline bar chart with optional waveform.

    Args:
        chunk_probs: P(AI) list for each 4-second chunk
        waveform: Optional mono waveform array for envelope visualization
        chunk_metadata: Optional metadata with start_sample info
        weighted_median: Energy-weighted median P(AI) for reference line

    Returns:
        plotly Figure with waveform (top) + P(AI) bars (bottom)
    """
    n = len(chunk_probs)
    times = [f"{i * CHUNK_SEC:.0f}-{(i + 1) * CHUNK_SEC:.0f}s" for i in range(n)]
    colors = ['#ff4757' if p >= 0.5 else '#2ed573' for p in chunk_probs]

    # 파형이 있으면 subplot, 없으면 단순 bar chart
    if waveform is not None and len(waveform) > 0:
        fig = make_subplots(
            rows=2, cols=1,
            row_heights=[0.3, 0.7],
            vertical_spacing=0.08,
            subplot_titles=("Waveform Envelope", "Segment-level AI Probability"),
        )

        # Waveform envelope (unipolar - 절댓값의 상단만)
        time_axis = np.arange(len(waveform)) / SR
        envelope = np.abs(waveform)

        # Downsample for plotting (매 100 샘플마다)
        downsample_factor = 100
        time_ds = time_axis[::downsample_factor]
        envelope_ds = envelope[::downsample_factor]

        fig.add_trace(
            go.Scatter(
                x=time_ds,
                y=envelope_ds,
                mode='lines',
                line=dict(color='#5f9ea0', width=0.5),
                fill='tozeroy',
                fillcolor='rgba(95, 158, 160, 0.3)',
                name='Envelope',
                hovertemplate="Time: %{x:.2f}s<br>Amplitude: %{y:.3f}<extra></extra>",
            ),
            row=1, col=1
        )

        # 세그먼트 경계선 표시 (chunk metadata 사용)
        if chunk_metadata:
            for meta in chunk_metadata:
                start_sec = meta['start_sample'] / SR
                fig.add_vline(
                    x=start_sec,
                    line=dict(color='#ffa502', width=1, dash='dot'),
                    opacity=0.5,
                    row=1, col=1
                )

        # P(AI) bar chart
        fig.add_trace(
            go.Bar(
                x=list(range(n)),
                y=chunk_probs,
                marker_color=colors,
                text=[f"{p:.2f}" for p in chunk_probs],
                textposition='outside',
                textfont=dict(size=10, color='white'),
                hovertemplate="<b>%{customdata}</b><br>P(AI): %{y:.3f}<extra></extra>",
                customdata=times,
                name='P(AI)',
            ),
            row=2, col=1
        )

        # Energy-weighted median reference line
        if weighted_median is not None:
            fig.add_hline(
                y=weighted_median, line_dash="dash", line_color="#00d2ff",
                annotation_text=f"Weighted Median ({weighted_median:.2f})",
                annotation_position="top right",
                annotation_font_color="#00d2ff",
                annotation_font_size=10,
                row=2, col=1
            )

        # Layout
        fig.update_xaxes(title_text="Time (s)", row=1, col=1)
        fig.update_yaxes(title_text="Amplitude", row=1, col=1)
        fig.update_xaxes(
            title_text="Segment",
            tickvals=list(range(n)),
            ticktext=times,
            tickangle=-45,
            tickfont=dict(size=9),
            row=2, col=1
        )
        fig.update_yaxes(title_text="P(AI)", range=[0, 1.05], row=2, col=1)

        fig.update_layout(
            plot_bgcolor='#1a1a2e',
            paper_bgcolor='#1a1a2e',
            font=dict(color='white'),
            margin=dict(l=50, r=20, t=60, b=60),
            height=500,
            showlegend=False,
        )

    else:
        # Fallback: 기존 단순 bar chart
        fig = go.Figure()

        fig.add_trace(go.Bar(
            x=list(range(n)),
            y=chunk_probs,
            marker_color=colors,
            text=[f"{p:.2f}" for p in chunk_probs],
            textposition='outside',
            textfont=dict(size=10, color='white'),
            hovertemplate="<b>%{customdata}</b><br>P(AI): %{y:.3f}<extra></extra>",
            customdata=times,
        ))

        if weighted_median is not None:
            fig.add_hline(y=weighted_median, line_dash="dash", line_color="#00d2ff",
                          annotation_text=f"Weighted Median ({weighted_median:.2f})",
                          annotation_position="top right",
                          annotation_font_color="#00d2ff")

        fig.update_layout(
            title=dict(text="Segment-level AI Probability", font=dict(size=14)),
            xaxis=dict(
                title="Segment",
                tickvals=list(range(n)),
                ticktext=times,
                tickangle=-45,
                tickfont=dict(size=9),
            ),
            yaxis=dict(title="P(AI)", range=[0, 1.05]),
            plot_bgcolor='#1a1a2e',
            paper_bgcolor='#1a1a2e',
            font=dict(color='white'),
            margin=dict(l=50, r=20, t=40, b=60),
            height=300,
            showlegend=False,
        )

    return fig