{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import json\n",
    "import torch\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "from suno_utils.utils.text import (\n",
    "    write_jsonl,\n",
    "    read_jsonl,\n",
    "    write_json,\n",
    "    read_json,\n",
    "    normalize_whitespace,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with open(\"/home/christian/code/christian/metadata/genius_hq_metas_quality.json\", \"r\") as f:\n",
    "    metas_map = json.load(f)\n",
    "\n",
    "print(len(metas_map))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/home/christian/code/christian/metadata/genius_hq_metas_tempo.json\", \"r\") as f:\n",
    "    tempo_metas_map = json.load(f)\n",
    "\n",
    "print(len(tempo_metas_map))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot hitsogram of tempo scores\n",
    "fig, axs = plt.subplots(1, 1, figsize=(8, 4), sharex=True, sharey=True)\n",
    "axs = np.reshape(axs, -1)\n",
    "bins = np.linspace(50, 200, 50)\n",
    "scores = []\n",
    "for idx, (meta_id, meta_tempo) in enumerate(tempo_metas_map.items()):\n",
    "    score = float(meta_tempo[\"tempo\"][\"bpm\"])\n",
    "    scores.append(score)\n",
    "\n",
    "axs[0].hist(scores, bins=bins, density=False)\n",
    "axs[0].set_yscale(\"log\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the main meta\n",
    "main_metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas.jsonl\", progress=True)\n",
    "\n",
    "main_metas_map = {}\n",
    "for main_meta in main_metas:\n",
    "    main_metas_map[main_meta[\"id\"]] = main_meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot hitsogram of quality scores\n",
    "fig, axs = plt.subplots(1, 1, figsize=(8, 4), sharex=True, sharey=True)\n",
    "axs = np.reshape(axs, -1)\n",
    "bins = np.linspace(-32, 5, 200)\n",
    "scores = []\n",
    "\n",
    "cutoff = -0.5\n",
    "\n",
    "low_quality_count = 0\n",
    "std_quality_count = 0\n",
    "high_quality_count = 0\n",
    "for idx, (meta_id, meta_quality) in enumerate(metas_map.items()):\n",
    "    score = float(meta_quality[\"audio_quality\"][\"score\"])\n",
    "    scores.append(score)\n",
    "\n",
    "    if score < -1.0:\n",
    "        low_quality_count += 1\n",
    "    elif score < -0.0:\n",
    "        std_quality_count += 1\n",
    "    else:\n",
    "        high_quality_count += 1\n",
    "\n",
    "axs[0].hist(scores, bins=bins, density=False)\n",
    "axs[0].set_yscale(\"log\")\n",
    "axs[0].vlines(-1.5, 0, 1000000, color=\"red\")\n",
    "axs[0].vlines(-0.5, 0, 1000000, color=\"red\")\n",
    "\n",
    "\n",
    "total_count = low_quality_count + std_quality_count + high_quality_count\n",
    "low_quality_percentage = (low_quality_count / (total_count)) * 100\n",
    "std_quality_percentage = (std_quality_count / (total_count)) * 100\n",
    "high_quality_percentage = (high_quality_count / (total_count)) * 100\n",
    "print(f\"low quality count: {low_quality_count} ({low_quality_percentage:0.1f}%)\")\n",
    "print(f\"std quality count: {std_quality_count} ({std_quality_percentage:0.1f}%)\")\n",
    "print(f\"high quality count: {high_quality_count} ({high_quality_percentage:0.1f}%)\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort metas by the audio quality score\n",
    "sorted_metas = sorted(metas_map.items(), key=lambda x: float(x[1][\"audio_quality\"][\"score\"]))\n",
    "\n",
    "\n",
    "# get worst quality\n",
    "worst_metas = sorted_metas[:10]\n",
    "for worst_meta_id, worst_meta in worst_metas:\n",
    "    main_meta = main_metas_map[worst_meta_id]\n",
    "    for key, val in main_meta.items():\n",
    "        print(key, val)\n",
    "    print(worst_meta_id, worst_meta[\"audio_quality\"][\"score\"])\n",
    "\n",
    "    # download from s3 \n",
    "    filename = os.path.basename(worst_meta[\"audio_filepath\"])\n",
    "    print(filename)\n",
    "    os.system(f\"\"\"aws s3 cp {worst_meta[\"audio_filepath\"]} /home/christian/code/christian/notebooks/outputs/s3\"\"\")\n",
    "    local_path = f\"/home/christian/code/christian/notebooks/outputs/s3/{filename}\"\n",
    "    print(local_path)\n",
    "    audio, sr = torchaudio.load(local_path)\n",
    "    if audio.abs().max() > 1.0:\n",
    "        audio = audio / audio.abs().max()\n",
    "\n",
    "    IPython.display.display(IPython.display.Audio(data=audio, rate=sr, normalize=False))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load full metas\n",
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas.jsonl\", progress=True)\n",
    "\n",
    "tags = {}\n",
    "# find the most popular tags\n",
    "for meta in tqdm(metas):\n",
    "    text_tags = meta[\"tags_text\"]\n",
    "    for text_tag in text_tags:\n",
    "        if text_tag not in tags:\n",
    "            tags[text_tag] = 1\n",
    "        else:\n",
    "            tags[text_tags] += 1\n",
    "\n",
    "# print 100 most popular tags\n",
    "sorted_tags = \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find the average quality score for each tag\n",
    "\n",
    "for tag in sorted_tags:\n",
    "    for meta in metas:\n",
    "        text_tags = "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for quality_threshold in [0.0, -0.5, -1, -2, -4, -6, -8]:\n",
    "    filtered_subset_metas = []\n",
    "    for metas in subset_metas:\n",
    "        if float(metas[\"audio_quality\"][\"score\"]) > quality_threshold:\n",
    "            filtered_subset_metas.append(metas)\n",
    "\n",
    "    keep_percent = len(filtered_subset_metas) / len(subset_metas) * 100\n",
    "    num_removed = len(subset_metas) - len(filtered_subset_metas)\n",
    "\n",
    "    print(f\"Quality threshold: {quality_threshold} keep: {len(filtered_subset_metas)} ({keep_percent:0.2f}%) remove: {num_removed} \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quality_threshold = 0.0\n",
    "subset_name = \"genius_hq\"\n",
    "root_dir = \"/home/christian/code/christian/outputs/quality_metas/audio_2ch_48khz_lg/\"\n",
    "\n",
    "for sub_dir in [\"train\", \"val\"]:\n",
    "    subset_metas = []\n",
    "\n",
    "    meta_filepath = os.path.join(root_dir, sub_dir, f\"{subset_name}_quality_metas.jsonl\")\n",
    "    subset_metas = read_jsonl(meta_filepath)\n",
    "\n",
    "    filtered_subset_metas = []\n",
    "    for metas in subset_metas:\n",
    "        if float(metas[\"audio_quality\"][\"score\"]) > quality_threshold:\n",
    "            filtered_subset_metas.append(metas)\n",
    "\n",
    "    keep_percent = len(filtered_subset_metas) / len(subset_metas) * 100\n",
    "    num_removed = len(subset_metas) - len(filtered_subset_metas)\n",
    "    print(f\"({sub_dir}) Quality threshold: {quality_threshold} keep: {len(filtered_subset_metas)} ({keep_percent:0.2f}%) remove: {num_removed} \")\n",
    "\n",
    "    write_jsonl(filtered_subset_metas, f\"/home/christian/code/christian/outputs/quality_metas/audio_2ch_48khz_lg/{sub_dir}/{subset_name}_filtered_0.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# listen to some tiktok covers data\n",
    "root_dir = \"/app/suno/christian/data/tiktok_covers_48khz/train\"\n",
    "filepaths = glob.glob(os.path.join(root_dir, \"*.wav\"))\n",
    "print(len(filepaths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# random filepath \n",
    "filepath = np.random.choice(filepaths)\n",
    "IPython.display.display(IPython.display.Audio(filepath, normalize=False))\n",
    "\n",
    "# s3://suno-data/datasets/metadata/quality/genius_hq_metas_quality.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load codec model\n",
    "import funcy\n",
    "import torchaudio\n",
    "from dac.model.dac4 import DAC\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "\n",
    "device = \"cuda:0\"\n",
    "# checkpoint_filepath = \"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth\"\n",
    "checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n",
    "load_f = funcy.partial(torch.load, map_location=\"cpu\")\n",
    "\n",
    "if checkpoint_filepath.startswith(\"s3://\"):\n",
    "    sd = read_from_s3(checkpoint_filepath, read_f=load_f)\n",
    "else:\n",
    "    sd = load_f(checkpoint_filepath)\n",
    "\n",
    "sd[\"metadata\"][\"kwargs\"] = {\n",
    "    k: v\n",
    "    for k, v in sd[\"metadata\"][\"kwargs\"].items()\n",
    "    if k in DAC.__init__.__code__.co_varnames\n",
    "}\n",
    "model_100hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n",
    "model_100hz.load_state_dict(sd[\"state_dict\"])\n",
    "model_100hz.eval()\n",
    "model_100hz.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# codec cycle random audio\n",
    "filepath = np.random.choice(filepaths)\n",
    "audio, sr = torchaudio.load(filepath)\n",
    "print(filepath)\n",
    "with torch.no_grad():\n",
    "    cycled_audio = model_100hz(audio.unsqueeze(0).to(device))[\"audio\"].squeeze(0).cpu()\n",
    "    peak = cycled_audio.abs().max()\n",
    "    if peak > 1:\n",
    "        cycled_audio /= peak\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(data=audio, rate=48_000, normalize=False))\n",
    "IPython.display.display(IPython.display.Audio(data=cycled_audio, rate=48_000, normalize=False))"
   ]
  },
  {
   "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
}