data_visualization.py

"""
Data Visualization Module

This module provides functions to visualize song embeddings using PCA.
PCA coordinates are saved to songs.json for persistence.
"""

import numpy as np
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

try:
    from songs_database.curation import load_database, save_database
    from songs_database.embed import get_embeddings_for_song
except ModuleNotFoundError:
    from curation import load_database, save_database
    from embed import get_embeddings_for_song


def get_averaged_embeddings() -> tuple[np.ndarray, list[int]]:
    """
    Get averaged embeddings for all embedded songs.
    
    Returns:
        Tuple of:
        - numpy array of shape (num_songs, 768) with averaged embeddings
        - list of internal_ids (for mapping back to songs)
    """
    db = load_database()
    
    embeddings = []
    internal_ids = []
    
    for song in db["songs"]:
        if not song.get("embedded", False):
            continue
        
        internal_id = song["internal_id"]
        song_embeddings = get_embeddings_for_song(internal_id)
        
        if song_embeddings is None or len(song_embeddings) == 0:
            continue
        
        # Average all segment embeddings
        avg_embedding = np.mean(song_embeddings, axis=0)
        embeddings.append(avg_embedding)
        internal_ids.append(internal_id)
    
    if not embeddings:
        return np.array([]), []
    
    return np.array(embeddings), internal_ids


def get_pca() -> dict:
    """
    Compute PCA for all embedded songs and save coordinates to songs.json.
    
    Adds/updates the 'PCA_2d' attribute for each embedded song:
        PCA_2d: {"x": float, "y": float}
    
    Returns:
        Dictionary with summary:
        {
            "processed": number of songs with PCA coordinates saved,
            "skipped": number of songs skipped (not embedded)
        }
    """
    print("Loading embeddings...")
    embeddings, internal_ids = get_averaged_embeddings()
    
    if len(embeddings) < 2:
        print("Need at least 2 embedded songs for PCA")
        return {"processed": 0, "skipped": 0}
    
    print(f"Found {len(embeddings)} embedded songs")
    print("Performing PCA...")
    
    # Compute PCA
    pca = PCA(n_components=2)
    coords_2d = pca.fit_transform(embeddings)
    
    # Save to songs.json
    print("Saving PCA coordinates to songs.json...")
    db = load_database()
    
    # Create mapping from internal_id to PCA coordinates
    id_to_coords = {
        internal_ids[i]: {"x": float(coords_2d[i, 0]), "y": float(coords_2d[i, 1])}
        for i in range(len(internal_ids))
    }
    
    processed = 0
    for song in db["songs"]:
        internal_id = song["internal_id"]
        if internal_id in id_to_coords:
            song["PCA_2d"] = id_to_coords[internal_id]
            processed += 1
    
    save_database(db)
    
    skipped = len(db["songs"]) - processed
    print(f"\n--- PCA Summary ---")
    print(f"Processed: {processed}")
    print(f"Skipped: {skipped}")
    
    return {"processed": processed, "skipped": skipped}


def visualize_embeddings(
    save_path: str = None,
    show_labels: bool = True,
    figsize: tuple = (12, 10)
) -> None:
    """
    Visualize songs using PCA coordinates from songs.json.
    
    Only displays songs that have the PCA_2d attribute.
    Run get_pca() first to compute and save coordinates.
    
    Args:
        save_path: Optional path to save the figure
        show_labels: Whether to show song names on the plot
        figsize: Figure size (width, height)
    """
    db = load_database()
    
    # Get songs with PCA_2d coordinates
    songs_with_pca = [s for s in db["songs"] if s.get("PCA_2d") is not None]
    
    if not songs_with_pca:
        print("No songs with PCA coordinates found!")
        print("Run get_pca() first to compute coordinates.")
        return
    
    print(f"Found {len(songs_with_pca)} songs with PCA coordinates")
    
    # Extract coordinates
    x_coords = [s["PCA_2d"]["x"] for s in songs_with_pca]
    y_coords = [s["PCA_2d"]["y"] for s in songs_with_pca]
    
    # Create plot
    plt.figure(figsize=figsize)
    plt.scatter(x_coords, y_coords, alpha=0.7, s=50)
    
    if show_labels:
        for i, song in enumerate(songs_with_pca):
            name = song.get("song_name", f"ID {song['internal_id']}")
            label = f"{name[:20]}"  # Truncate long names
            plt.annotate(
                label,
                (x_coords[i], y_coords[i]),
                fontsize=8,
                alpha=0.7
            )
    
    plt.xlabel("PC1")
    plt.ylabel("PC2")
    plt.title(f"Song Embeddings PCA (n={len(songs_with_pca)})")
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150)
        print(f"Saved to {save_path}")
    
    plt.show()


def get_2d_coordinates() -> list[dict]:
    """
    Get 2D PCA coordinates for all songs from songs.json.
    
    Returns:
        List of dicts with song info and coordinates:
        [{"internal_id": 0, "song_name": "...", "artist": "...", "x": 0.5, "y": -0.3}, ...]
    """
    db = load_database()
    
    results = []
    for song in db["songs"]:
        if song.get("PCA_2d") is None:
            continue
        
        results.append({
            "internal_id": song["internal_id"],
            "song_name": song.get("song_name", "Unknown"),
            "artist": song.get("artists", ["Unknown"])[0],
            "x": song["PCA_2d"]["x"],
            "y": song["PCA_2d"]["y"]
        })
    
    return results


if __name__ == "__main__":
    # First compute PCA and save to songs.json
    get_pca()
    
    # Then visualize from songs.json
    visualize_embeddings(save_path="pca_visualization.png", show_labels=True)