{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the audio production features for genius, imslp, and youtube_music\n",
    "genius_ap = pd.read_csv(\"/home/christian/code/christian/metadata/v4/genius_audio_production_features_v2.csv\")\n",
    "imslp_ap = pd.read_csv(\"/home/christian/code/christian/metadata/v4/imslp_audio_production_features_v2.csv\")\n",
    "youtube_music_ap = pd.read_csv(\"/home/christian/code/christian/metadata/v4/youtube_music_audio_production_features_v2.csv\")\n",
    "\n",
    "# load the base jsonl file \n",
    "# combine into one dataframe\n",
    "combined_ap = pd.concat([genius_ap, imslp_ap, youtube_music_ap])\n",
    "combined_ap.reset_index(drop=True, inplace=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_ap.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_ap[combined_ap[\"id\"] == \"OzrkFekIbxE\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# before we merge, lets create psuedo tags for all rows in the combined_ap dataframe and save it to a new csv\n",
    "\n",
    "# define the mean and std for each feature (hard code for reproducibility)\n",
    "# these numbers are based on genius_hq, youtube_music, imslp, and discogs\n",
    "feature_stats = {\n",
    "    \"spectral_centroid\": {\"mean\": 3100.0, \"std\": 900.0},\n",
    "    \"bass\": {\"mean\": 0.3, \"std\": 0.1},\n",
    "    \"mid\": {\"mean\": 0.6, \"std\": 0.1}, \n",
    "    \"high\": {\"mean\": 0.6, \"std\": 0.3},\n",
    "    \"stereo_width\": {\"mean\": 0.2, \"std\": 0.1},\n",
    "    \"spectral_flatness\": {\"mean\": 0.08, \"std\": 0.05},\n",
    "    \"crest_factor\": {\"mean\": 1.8, \"std\": 0.5},\n",
    "    \"silence_percentage\": {\"mean\": 1.5, \"std\": 3.0},\n",
    "    \"loudness\": {\"mean\": -12.0, \"std\": 5.0}\n",
    "}\n",
    "\n",
    "def describe_spectrum(bass_energy, mid_energy, high_energy):\n",
    "    # Calculate ratios\n",
    "    if mid_energy == 0:  # Avoid division by zero\n",
    "        mid_energy = 1e-10\n",
    "    if high_energy == 0:  # Avoid division by zero\n",
    "        high_energy = 1e-10\n",
    "\n",
    "    bass_mid_ratio = bass_energy / mid_energy\n",
    "    bass_high_ratio = bass_energy / high_energy\n",
    "    mid_high_ratio = mid_energy / high_energy\n",
    "\n",
    "    # Descriptors based on the ratios\n",
    "    descriptors = []\n",
    "\n",
    "    # Bass dominant\n",
    "    if bass_mid_ratio > 1.5 and bass_high_ratio > 1.5:\n",
    "        descriptors.extend([\"bassy\", \"boomy\", \"thumpy\"])\n",
    "\n",
    "    # Mid dominant\n",
    "    if bass_mid_ratio < 0.67 and mid_high_ratio > 1.5:\n",
    "        descriptors.extend([\"warm\", \"full-bodied\", \"nasal\", \"boxy\"])\n",
    "\n",
    "    # High dominant\n",
    "    if bass_high_ratio < 0.67 and mid_high_ratio < 0.67:\n",
    "        descriptors.extend([\"bright\", \"sibilant\", \"tinny\"])\n",
    "\n",
    "    # Balanced audio\n",
    "    if 0.9 <= bass_mid_ratio <= 1.1 and 0.9 <= bass_high_ratio <= 1.1 and 0.9 <= mid_high_ratio <= 1.1:\n",
    "        descriptors.extend([\"balanced\", \"even\", \"neutral\"])\n",
    "\n",
    "    # Muddy (overlapping with mid dominant but more specific to lack of clarity)\n",
    "    if bass_mid_ratio > 0.8 and mid_high_ratio < 1.2:\n",
    "        descriptors.append(\"muddy\")\n",
    "\n",
    "    # Detailed or clear, considering mid-high clarity\n",
    "    if mid_high_ratio > 1.2:\n",
    "        descriptors.append(\"detailed\")\n",
    "\n",
    "    # Forward (applicable if mids are clearly dominant over highs and slightly over bass)\n",
    "    if mid_high_ratio > 1.2 and bass_mid_ratio < 1.2:\n",
    "        descriptors.append(\"forward\")\n",
    "\n",
    "    return descriptors\n",
    "\n",
    "\n",
    "# given a row of audio features, return a list of descriptor features\n",
    "def get_descriptor_features(row):\n",
    "    tags = []\n",
    "\n",
    "    stereo_width = row[\"stereo_width\"]\n",
    "    stereo_width_normalized = row[\"stereo_width_normalized\"] \n",
    "    crest_factor = row[\"crest_factor\"]\n",
    "    crest_factor_normalized = row[\"crest_factor_normalized\"]\n",
    "    spectral_flatness = row[\"spectral_flatness\"]\n",
    "    spectral_flatness_normalized = row[\"spectral_flatness_normalized\"]\n",
    "    spectral_centroid = row[\"spectral_centroid\"]\n",
    "    spectral_centroid_normalized = row[\"spectral_centroid_normalized\"]\n",
    "    bass = row[\"bass\"]\n",
    "    mid = row[\"mid\"]\n",
    "    high = row[\"high\"]\n",
    "    loudness = row[\"loudness\"]\n",
    "\n",
    "    # -------- spectral flatness --------\n",
    "    if not np.isnan(spectral_flatness) and not np.isnan(spectral_flatness_normalized):\n",
    "        if spectral_flatness < 0.05:\n",
    "            tags.append(\"rolled-off\")\n",
    "        elif spectral_flatness > 0.2:\n",
    "            tags.append(\"noisy\")\n",
    "        tags.append(f\"sf:{round(spectral_flatness_normalized)}\")\n",
    "\n",
    "    # -------- spectrum analysis --------\n",
    "    if not np.isnan(bass) and not np.isnan(mid) and not np.isnan(high):\n",
    "        tags.extend(describe_spectrum(bass, mid, high))\n",
    "\n",
    "    # -------- spectral centroid --------\n",
    "    if not np.isnan(spectral_centroid) and not np.isnan(spectral_centroid_normalized):\n",
    "        if spectral_centroid < 1500:\n",
    "            tags.append(\"very warm\")\n",
    "            tags.append(\"very dark\")\n",
    "        elif spectral_centroid >= 1500 and spectral_centroid < 2000:\n",
    "            tags.append(\"warm\")\n",
    "            tags.append(\"dark\")\n",
    "        elif spectral_centroid >= 2000 and spectral_centroid < 3000:\n",
    "            pass\n",
    "        elif spectral_centroid >= 3750 and spectral_centroid < 4500:\n",
    "            tags.append(\"bright\")\n",
    "        else:\n",
    "            tags.append(\"bright\")\n",
    "            tags.append(\"very bright\")\n",
    "            tags.append(\"sharp\")\n",
    "        tags.append(f\"sc:{round(spectral_centroid_normalized)}\")\n",
    "\n",
    "    # -------- stereo width --------\n",
    "    if not np.isnan(stereo_width) and not np.isnan(stereo_width_normalized):\n",
    "        if stereo_width < 0.05:\n",
    "            tags.append(\"mono\")\n",
    "        elif stereo_width >= 0.05 and stereo_width < 0.1:\n",
    "            tags.append(\"narrow\")\n",
    "        elif stereo_width >= 0.1 and stereo_width < 0.3:\n",
    "            tags.append(\"stereo\")\n",
    "        else:\n",
    "            tags.append(\"stereo\")\n",
    "            tags.append(\"wide stereo\")\n",
    "            tags.append(\"wide\")\n",
    "        tags.append(f\"sw:{round(stereo_width_normalized)}\")\n",
    "\n",
    "    # -------- crest factor --------\n",
    "    if not np.isnan(crest_factor) and not np.isnan(crest_factor_normalized):\n",
    "        if crest_factor < 1.0:\n",
    "            tags.append(\"slammed\")\n",
    "            tags.append(\"maximized loudness\")\n",
    "            tags.append(\"very compressed\")\n",
    "        elif crest_factor >= 1.0 and crest_factor < 1.5:\n",
    "            tags.append(\"compressed\")\n",
    "        elif crest_factor >= 1.5 and crest_factor < 2.5:\n",
    "            tags.append(\"dynamic\")\n",
    "        else:\n",
    "            tags.append(\"very dynamic\")\n",
    "            tags.append(\"dynamic\")\n",
    "        tags.append(f\"cf:{round(crest_factor_normalized)}\")\n",
    "\n",
    "    return tags\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalize each feature using the predefined means and stds\n",
    "for feature, stats in feature_stats.items():\n",
    "    combined_ap[f\"{feature}_normalized\"] = (combined_ap[feature] - stats[\"mean\"]) / stats[\"std\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_ap[\"tags\"] = None\n",
    "\n",
    "# iterate over each row in the combined_ap dataframe and get the descriptor features\n",
    "for index, row in tqdm(combined_ap.iterrows()):\n",
    "    tags = get_descriptor_features(row)\n",
    "    combined_ap.at[index, \"tags\"] = tags\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the combined_ap dataframe to a new csv\n",
    "combined_ap.to_csv(\"/home/christian/code/christian/metadata/v4/combined_audio_production_features_v2.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read combined_ap from csv\n",
    "combined_ap = pd.read_csv(\"/home/christian/code/christian/metadata/v4/combined_audio_production_features_v2.csv\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ap_dict = combined_ap.set_index('id').to_dict(orient='index')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ast import literal_eval\n",
    "\n",
    "base_meta = base_metas[2]\n",
    "tags_dict = ap_dict[base_meta[\"id\"]]\n",
    "tags = list(literal_eval(tags_dict[\"tags\"]))\n",
    "# convert string to list of strings\n",
    "print(tags)\n",
    "# convert string to list \n",
    "print(type(tags))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ast import literal_eval\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the base jsonl file \n",
    "base_dir = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz/\"\n",
    "#base_dir = \"/app/suno/data/chirp_v5_ft/v0\"\n",
    "\n",
    "for subset in [\"val\", \"tr\"]:\n",
    "    base_jsonl = os.path.join(base_dir, f\"metas_{subset}.jsonl\")\n",
    "    base_metas = read_jsonl(base_jsonl)\n",
    "    print(len(base_metas))\n",
    "\n",
    "    new_subset_metas = []\n",
    "\n",
    "    found_ids = 0\n",
    "    for base_meta in tqdm(base_metas):\n",
    "        new_meta = base_meta.copy()\n",
    "\n",
    "        if \"tags\" not in new_meta:\n",
    "            new_meta[\"tags\"] = []\n",
    "\n",
    "        # check if the id is in the combined_ap dataframe\n",
    "        if base_meta[\"id\"] in ap_dict:\n",
    "            # add the audio production features to the base meta\n",
    "            tags_dict = ap_dict[base_meta[\"id\"]]\n",
    "            tags = list(literal_eval(tags_dict[\"tags\"]))\n",
    "            new_meta[\"tags\"].extend(tags)\n",
    "            found_ids += 1\n",
    "        \n",
    "        new_subset_metas.append(new_meta)\n",
    "\n",
    "    print(f\"Found {found_ids}/{len(base_metas)} ids in the combined_ap dataframe\")\n",
    "    # write the updated base metas to a new jsonl file\n",
    "    write_jsonl(new_subset_metas, os.path.join(base_dir, f\"metas_{subset}_v1.jsonl\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_metas[116]"
   ]
  },
  {
   "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
}