{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import json\n",
    "import pandas as pd\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find all json files in cumulative_playlists\n",
    "cumulative_playlists_dir = \"/home/christian/code/christian/metadata/playlists/spotify-playlist-archive-main/playlists/cumulative\"\n",
    "json_files = glob.glob(os.path.join(cumulative_playlists_dir, \"*.json\"))\n",
    "\n",
    "print(len(json_files))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "# merge all json files into a single jsonl file \n",
    "metas = []\n",
    "\n",
    "for json_file in tqdm(json_files):\n",
    "    with open(json_file, \"r\") as f:\n",
    "        data = json.load(f)\n",
    "\n",
    "    playlist_name = data[\"name\"]\n",
    "    description = data[\"description\"]\n",
    "    tracks = data[\"tracks\"]\n",
    "    track_metas = []\n",
    "    for track in tracks:\n",
    "        artists = [artist[\"name\"] for artist in track[\"artists\"]]\n",
    "        track_meta = {\n",
    "            \"name\": track[\"name\"],\n",
    "            \"artists\": artists,\n",
    "        }\n",
    "        track_metas.append(track_meta)\n",
    "\n",
    "    meta = {\n",
    "        \"playlist_name\": playlist_name,\n",
    "        \"description\": description,\n",
    "        \"tracks\": track_metas,\n",
    "    }\n",
    "    metas.append(meta)\n",
    "\n",
    "print(len(metas))\n",
    "\n",
    "write_jsonl(metas, \"/home/christian/code/christian/metadata/playlists/spotify_playlist_archive.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/playlists/spotify_playlist_archive.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect a list of all songs + artists and count unique ones\n",
    "songs = []\n",
    "for meta in metas:\n",
    "    for track in meta[\"tracks\"]:\n",
    "        full_name = f\"{track['name']} - {', '.join(track['artists'])}\"\n",
    "        songs.append(full_name)\n",
    "\n",
    "songs = list(set(songs))\n",
    "\n",
    "print(len(songs))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a list of all playlist names and descriptions\n",
    "playlist_names = [meta[\"playlist_name\"] for meta in metas]\n",
    "playlist_descriptions = [meta[\"description\"] for meta in metas]\n",
    "\n",
    "for playlist_name, playlist_description in zip(playlist_names, playlist_descriptions):\n",
    "    print(playlist_name)\n",
    "    print(playlist_description)\n",
    "    print()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# can we intersect this with the discogs subset metas\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "\n",
    "discogs_subset_metas_filepath = \"/home/christian/code/christian/metadata/v4/discogs_subset_metas_raw.jsonl\"\n",
    "discogs_subset_metas = read_jsonl(discogs_subset_metas_filepath)\n",
    "print(len(discogs_subset_metas))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rapidfuzz import process, fuzz\n",
    "from rapidfuzz.process import cdist\n",
    "from multiprocessing import Pool\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "def create_song_lookups(metas):\n",
    "    return {'combined': {\n",
    "        f\"{meta['title']} {meta['artists'][0]['name']}\".lower().strip(): meta \n",
    "        for meta in metas \n",
    "        if meta.get('title') and meta.get('artists') and meta['artists']\n",
    "    }}\n",
    "\n",
    "def process_chunk(args):\n",
    "    chunk, library_keys, threshold = args\n",
    "    return process.extract(chunk, library_keys, scorer=fuzz.ratio, \n",
    "                         score_cutoff=threshold, limit=3)\n",
    "\n",
    "def find_matches_parallel(playlist_songs, library_songs, threshold=85):\n",
    "    library_lookups = create_song_lookups(library_songs)\n",
    "    library_keys = list(library_lookups['combined'].keys())\n",
    "    \n",
    "    search_strings = [f\"{song['name']} {', '.join(song['artists'])}\".lower().strip() \n",
    "                     for song in playlist_songs]\n",
    "    \n",
    "    chunks = [(s, library_keys, threshold) for s in search_strings]\n",
    "    \n",
    "    matches = []\n",
    "    with Pool() as pool:\n",
    "        for i, combined_matches in enumerate(tqdm(\n",
    "            pool.imap(process_chunk, chunks), \n",
    "            total=len(chunks)\n",
    "        )):\n",
    "            for match_str, score, _ in combined_matches:\n",
    "                matches.append({\n",
    "                    'chart_song': chart_songs[i],\n",
    "                    'library_song': library_lookups['combined'][match_str],\n",
    "                    'match_score': score\n",
    "                })\n",
    "    return matches\n",
    "\n",
    "def find_matches_vectorized(playlist_songs, library_songs, threshold=85):\n",
    "    library_lookups = create_song_lookups(library_songs)\n",
    "    library_keys = list(library_lookups['combined'].keys())\n",
    "    \n",
    "    search_strings = [f\"{song['name']} {', '.join(song['artists'])}\".lower().strip() \n",
    "                     for song in playlist_songs]\n",
    "    \n",
    "    distances = cdist(search_strings, library_keys, scorer=fuzz.ratio, workers=128)\n",
    "    \n",
    "    matches = []\n",
    "    for i, row in enumerate(distances):\n",
    "        match_indices = np.where(row >= threshold)[0][:3]\n",
    "        \n",
    "        for idx in match_indices:\n",
    "            match_str = library_keys[idx]\n",
    "            matches.append({\n",
    "                'playlist_song': playlist_songs[i],\n",
    "                'library_song': library_lookups['combined'][match_str],\n",
    "                'match_score': row[idx]\n",
    "            })\n",
    "    \n",
    "    return matches\n",
    "\n",
    "# Usage example:\n",
    "def process_playlist_matches(playlist_songs, discogs_subset_metas, method='vectorized'):    \n",
    "    \n",
    "    if method == 'parallel':\n",
    "        matches = find_matches_parallel(playlist_songs, discogs_subset_metas)\n",
    "    else:\n",
    "        matches = find_matches_vectorized(playlist_songs, discogs_subset_metas)\n",
    "    \n",
    "    playlist_metas = []\n",
    "    for match in matches:\n",
    "        billboard_meta = {\n",
    "            \"id\": match[\"library_song\"][\"id\"],\n",
    "            \"title\": match[\"library_song\"][\"title\"],\n",
    "            \"artist\": match[\"library_song\"][\"artists\"][0][\"name\"],\n",
    "            \"match_score\": match[\"match_score\"],\n",
    "        }\n",
    "        # add tags\n",
    "        playlist_metas.append(billboard_meta)\n",
    "    \n",
    "    return playlist_metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas[1][\"tracks\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_metas = []\n",
    "\n",
    "for meta in tqdm(metas):\n",
    "    playlist_songs = meta[\"tracks\"]\n",
    "    playlist_name = meta[\"playlist_name\"]\n",
    "    playlist_description = meta[\"description\"]\n",
    "    results = process_playlist_matches(playlist_songs, discogs_subset_metas)\n",
    "\n",
    "    new_meta = {\n",
    "        \"playlist_name\": playlist_name,\n",
    "        \"playlist_description\": playlist_description,\n",
    "        \"tracks\": [],\n",
    "    }\n",
    "\n",
    "    for result in results:\n",
    "        new_meta[\"tracks\"].append({\n",
    "            \"id\": result[\"id\"],\n",
    "            \"song_name\": result[\"title\"],\n",
    "            \"artist_name\": result[\"artist\"],\n",
    "            \"match_score\": float(result[\"match_score\"]),\n",
    "        })\n",
    "    output_metas.append(new_meta)\n",
    "    write_jsonl(output_metas, \"/home/christian/code/christian/metadata/playlists/spotify_playlist_archive_discogs_subset_matches.jsonl\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}