Blame for patch_view/visualization.py

"""
Visualization generation for BLT patching explorer.
"""

import numpy as np
import librosa
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from typing import List

from .audio import load_audio_file, TOKENS_PER_SECOND


def compute_spectrogram(audio_path: str, sr: int | None = None):
    """
    Compute mel-spectrogram from audio file.

    Args:
        audio_path: Path to audio file
        sr: Optional sample rate override. When None, uses the file's native rate.

    Returns:
        times: Time points (seconds)
        frequencies: Frequency bins (Hz)
        magnitudes: 2D array (freq x time) in dB
    """
    file_sr, audio = load_audio_file(audio_path)
    sr = sr or file_sr

    n_fft = 1024
    hop_length = 64
    n_mels = 256

    S = librosa.feature.melspectrogram(
        y=audio,
        sr=sr,
        n_fft=n_fft,
        hop_length=hop_length,
        n_mels=n_mels,
        fmax=8000
    )
    S_db = librosa.power_to_db(S, ref=np.max)

    times = librosa.frames_to_time(
        np.arange(S.shape[1]), sr=sr, hop_length=hop_length
    )
    frequencies = librosa.mel_frequencies(n_mels=n_mels, fmax=8000)

    return times, frequencies, S_db


def create_visualization(
    data: dict,
    patches: List[List[int]],
    metric_name: str,
    threshold: float
):
    """
    Create 4-panel synchronized visualization.

    Args:
        data: Precomputed data dict
        patches: Patch boundaries
        metric_name: "entropy" or "perplexity"
        threshold: Threshold value

    Returns:
        Plotly figure
    """
    import time
    start_time = time.time()
    print(f"\n[Visualization] Starting for {data['file_id']}")
    print(f"  Patches: {len(patches)}, Metric: {metric_name}, Threshold: {threshold}")

    fig = make_subplots(
        rows=3, cols=1,
        subplot_titles=(
            'Spectrogram',
            f'{metric_name.capitalize()} + Patch Boundaries',
            'Patch Size Distribution'
        ),
        specs=[
            [{"type": "heatmap"}],
            [{"type": "scatter"}],
            [{"type": "bar"}]
        ],
        vertical_spacing=0.1,
        row_heights=[0.4, 0.4, 0.2]
    )

    # Get metric data
    metric = np.array(
        data['entropies'] if metric_name == 'entropy' else data['perplexities']
    )

    # Calculate patch boundary times
    boundary_times = []
    for patch in patches:
        if patch:
            boundary_times.append(patch[0] / TOKENS_PER_SECOND)

    duration = data['audio_duration_sec']

    # Limit vlines to avoid slowdown
    MAX_VLINES = 200
    draw_boundaries = len(boundary_times) <= MAX_VLINES

    if draw_boundaries:
        print(f"  Will draw {len(boundary_times)} patch boundaries")
    else:
        print(f"  Skipping {len(boundary_times)} boundaries (too many, would be slow)")

    # 1. Spectrogram
    spec_start = time.time()
    times, freqs, mags = compute_spectrogram(data['audio_path'])
    print(f"  Spectrogram computed in {time.time() - spec_start:.2f}s")

    fig.add_trace(
        go.Heatmap(
            x=times,
            y=freqs,
            z=mags,
            colorscale='Viridis',
            showscale=True,
            colorbar=dict(title="dB", x=1.05)
        ),
        row=1, col=1
    )

    # Add patch boundaries to spectrogram
    if draw_boundaries:
        for t in boundary_times:
            fig.add_vline(
                x=t, line_width=2, line_dash="solid",
                line_color="red", opacity=0.7,
                row=1, col=1
            )

    fig.update_xaxes(title_text="Time (s)", row=1, col=1, range=[0, duration])
    fig.update_yaxes(title_text="Frequency (Hz)", row=1, col=1)

    # 2. Entropy/Perplexity plot
    token_positions = np.arange(len(metric))

    fig.add_trace(
        go.Scatter(
            x=token_positions,
            y=metric,
            mode='lines+markers',
            name=metric_name.capitalize(),
            line=dict(color='blue', width=2),
            marker=dict(size=3)
        ),
        row=2, col=1
    )

    # Add threshold line
    fig.add_hline(
        y=threshold, line_width=2, line_dash="dash",
        line_color="red", opacity=0.7,
        annotation_text=f"Threshold: {threshold:.2f}",
        annotation_position="right",
        row=2, col=1
    )

    # Add patch boundaries
    if draw_boundaries:
        for patch in patches:
            if patch:
                fig.add_vline(
                    x=patch[0], line_width=1, line_dash="solid",
                    line_color="gray", opacity=0.4,
                    row=2, col=1
                )

    fig.update_xaxes(title_text="Token Position", row=2, col=1, range=[0, len(metric)])
    fig.update_yaxes(title_text=metric_name.capitalize(), row=2, col=1)

    # 3. Patch size distribution
    sizes = [len(p) for p in patches]
    size_counts = {}
    for size in sizes:
        size_counts[size] = size_counts.get(size, 0) + 1

    sorted_sizes = sorted(size_counts.keys())
    sorted_counts = [size_counts[s] for s in sorted_sizes]

    fig.add_trace(
        go.Bar(
            x=sorted_sizes,
            y=sorted_counts,
            name='Patch Sizes',
            marker_color='steelblue',
            hovertemplate='Size: %{x} tokens<br>Count: %{y}<extra></extra>'
        ),
        row=3, col=1
    )

    fig.update_xaxes(title_text="Patch Size (tokens)", row=3, col=1)
    fig.update_yaxes(title_text="Count", row=3, col=1)

    # Update layout
    fig.update_layout(
        height=1200,
        showlegend=False,
        hovermode='closest',
        title_text=f"BLT Patching Visualization - {data['file_id']}"
    )

    total_time = time.time() - start_time
    print(f"  Total visualization time: {total_time:.2f}s")
    if not draw_boundaries:
        print(f"  ⚠️  Note: Patch boundaries not drawn (>{MAX_VLINES} patches)")

    return fig