{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e6ad8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a200465",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = \"/app2/suno/data/diffusion/v1\"\n",
    "#metas_filename = \"metas_v9_tr.jsonl\"\n",
    "metas_filename = \"metas_v9_val_filtered.jsonl\"\n",
    "\n",
    "auk_metas = read_jsonl(os.path.join(base_dir, metas_filename))\n",
    "print(len(auk_metas))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "385cf3f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "auk_metas[100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43b9484",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_metas = []\n",
    "from tqdm import tqdm\n",
    "from collections import Counter\n",
    "\n",
    "# count all the unique stem keys and their occurrences\n",
    "stem_keys = set()\n",
    "stem_key_counter = Counter()\n",
    "\n",
    "# only take metas that have \"tags\" or \"text\" or \"text_aligned\"\n",
    "for meta in tqdm(auk_metas):\n",
    "    if \"tags\" in meta or \"text\" in meta or \"text_aligned\" in meta:\n",
    "        filtered_metas.append(meta)\n",
    "    if \"stems\" in meta:\n",
    "        stem_keys.update(meta[\"stems\"].keys())\n",
    "        stem_key_counter.update(meta[\"stems\"].keys())\n",
    "\n",
    "print(len(filtered_metas))\n",
    "print(len(stem_keys))\n",
    "for stem_key, count in stem_key_counter.items():\n",
    "    if count > 1:\n",
    "        print(stem_key, count)\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7966cd4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = \"/app2/suno/data/diffusion/v1\"\n",
    "\n",
    "write_jsonl(filtered_metas, os.path.join(output_dir, \"metas_v9_val_filtered.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f581da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. load the diffusion memmap metas\n",
    "# 2. load the GPT metas (with the artist names\n",
    "# 3. iterate over the diffusion metas and for each meta:\n",
    "# 4. find the associated GPT meta (which has the artist name)\n",
    "# 5. apply the same hash function as Victor to get the artist hash\n",
    "# 6. add the artist hash as a tag to the diffusion meta (also add as a special tag called \"artist_hash\")\n",
    "# 7. write the updated diffusion metas to a new file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e89e30a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read ids from the file\n",
    "with open(\"/home/christian/code/christian/metadata/auk_id_to_artist_ids.json\", \"r\") as f:\n",
    "    auk_id_to_artist_ids = json.load(f)\n",
    "\n",
    "print(len(auk_id_to_artist_ids))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74dab242",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = \"/app2/suno/data/diffusion/v0\"\n",
    "subset = \"tr\"\n",
    "metas_filename = f\"metas_{subset}_aligned_v3.jsonl\"\n",
    "\n",
    "metas = read_jsonl(os.path.join(base_dir, metas_filename))\n",
    "print(len(metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "098ef5dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "new_metas = []\n",
    "artist_count = 0\n",
    "for meta in tqdm(metas):\n",
    "    # check if this meta id is in the auk_id_to_artist_ids map\n",
    "    if meta[\"id\"] in auk_id_to_artist_ids:\n",
    "        meta[\"artist_ids\"] = auk_id_to_artist_ids[meta[\"id\"]]\n",
    "        artist_count += 1\n",
    "    new_metas.append(meta)\n",
    "\n",
    "print(len(new_metas), artist_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d7b83c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(new_metas, os.path.join(base_dir, f\"metas_{subset}_aligned_v3_artist_ids.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31bfbb8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ae6f3d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = \"/app2/suno/data/auk_v0\"\n",
    "metas_filename = \"metas_v9_tr.jsonl\"\n",
    "#metas_filename = \"metas_v8_val.jsonl\"\n",
    "\n",
    "auk_metas = read_jsonl(os.path.join(base_dir, metas_filename))\n",
    "print(len(auk_metas))\n",
    "\n",
    "artist_id_count = 0\n",
    "for meta in auk_metas:\n",
    "    if \"artist_ids\" in meta:\n",
    "        artist_id_count += 1\n",
    "\n",
    "print(f\"Found {artist_id_count} artists of {len(auk_metas)} metas ({(artist_id_count/len(auk_metas))*100:.2f}%)\")\n",
    "auk_unique_ids = set([meta[\"id\"] for meta in auk_metas])\n",
    "print(len(auk_unique_ids))\n",
    "\n",
    "# get the intersection of the two sets\n",
    "intersection = unique_ids.intersection(auk_unique_ids)\n",
    "print(len(intersection))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b72d59",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from tqdm import tqdm\n",
    "# now lets save a map dict of the diffusion ids to the artist_ids from the auk metas\n",
    "auk_id_to_artist_ids = {}\n",
    "for meta in tqdm(auk_metas):\n",
    "    if \"artist_ids\" in meta:\n",
    "        auk_id_to_artist_ids[meta[\"id\"]] = meta[\"artist_ids\"]\n",
    "        \n",
    "\n",
    "# now lets save the map to a file\n",
    "with open(\"/home/christian/code/christian/metadata/auk_id_to_artist_ids.json\", \"w\") as f:\n",
    "    json.dump(auk_id_to_artist_ids, f)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c37164a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "METAS_DIR = \"/app/suno/tmp\"\n",
    "discogs_subset_metas = {m[\"id\"]: m for m in read_jsonl(os.path.join(METAS_DIR, \"raw_discogs_subset_metas.jsonl\"))}\n",
    "print(len(discogs_subset_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ecc293",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_id = list(discogs_subset_metas.keys())[0]\n",
    "print(discogs_subset_metas[test_id])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4778355",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hashlib\n",
    "from tqdm import tqdm\n",
    "\n",
    "artist_hash_map = {}\n",
    "for meta_id, meta in tqdm(discogs_subset_metas.items()):\n",
    "    artists = meta.get(\"artists\", [])    \n",
    "    artist_names = [a[\"name\"] for a in artists]\n",
    "    if meta_id in auk_id_to_artist_ids:\n",
    "        artist_ids = auk_id_to_artist_ids[meta_id]\n",
    "        # Create a dict mapping artist_name to artist_id\n",
    "        artist_hash = [\n",
    "            hashlib.md5((artist_id + \"fj48x39\").encode()).hexdigest()[:16] for artist_id in artist_ids\n",
    "        ]\n",
    "        for artist_name, artist_id, artist_hash in zip(artist_names, artist_ids, artist_hash):\n",
    "            artist_hash_map[artist_name] = {\n",
    "                \"artist_hash\": artist_hash,\n",
    "                \"artist_id\": artist_id,\n",
    "            }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94f63492",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Found {len(artist_hash_map)} artists\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01fb6074",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort the artist_hash_map by the key\n",
    "artist_hash_map = dict(sorted(artist_hash_map.items(), key=lambda x: x[0]))\n",
    "\n",
    "# save the artist_hash_map to a file\n",
    "with open(\"/home/christian/code/christian/metadata/artist_hash_map.json\", \"w\") as f:\n",
    "    json.dump(artist_hash_map, f)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb0c25b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "artist_hash_map[\"Taylor Swift\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83520d56",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b239937",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "name_to_hash_filepath = \"/home/christian/code/christian/metadata/artist_hash_map.json\"\n",
    "with open(name_to_hash_filepath, \"r\") as f:\n",
    "    artist_name_to_hash = json.load(f)\n",
    "artist_name_to_hash[\"Taylor Swift\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb376dcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now make a reverse map\n",
    "id_to_name_filepath = \"/home/christian/code/christian/metadata/artist_hash_map_reverse.json\"\n",
    "\n",
    "# build the reverse map: map artist_hash to artist_name\n",
    "artist_id_to_name = {\n",
    "    v[\"artist_id\"]: k for k, v in artist_name_to_hash.items()\n",
    "}\n",
    "\n",
    "with open(id_to_name_filepath, \"w\") as f:\n",
    "    json.dump(artist_id_to_name, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1a38799",
   "metadata": {},
   "outputs": [],
   "source": [
    "artist_id_to_name[\"3utyef8u\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16dfaa3b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}