{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, write_jsonl, read_json\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "### Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tag_data = read_jsonl(\"/app2/suno/data/sara/sfx_synthetic_tags/combined_metas_filter_v0_with_key_data.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "og_meta = read_jsonl(\"/app2/suno/data/diffusion/sfx/v1/combined_metas_filter_v0.jsonl\")\n",
    "metas_map = {meta[\"id\"]: meta for meta in og_meta}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for meta_row in tag_data:\n",
    "    id = meta_row['id']\n",
    "    og_row = metas_map[id]\n",
    "    tags = og_row[\"tags\"]\n",
    "\n",
    "    og_row[\"key\"] = None\n",
    "    og_key_center = None\n",
    "    if tags is None or len(tags) == 0:\n",
    "        pass\n",
    "    elif tags[-1].startswith(\"key of\"):\n",
    "        og_key_center = tags[-1][len(\"key of \"):]\n",
    "        og_row[\"tags\"] = tags[:-1]\n",
    "\n",
    "    inferred_key = meta_row.get(\"inferred_key\", None)\n",
    "    confidence = meta_row.get(\"inferred_key_confidence\", None)\n",
    "    if inferred_key is not None and confidence is not None:\n",
    "        key_center = inferred_key.split(\"_\")[0].lower()\n",
    "        tonality = inferred_key.split(\"_\")[1].lower()\n",
    "\n",
    "        final_key = og_key_center\n",
    "        if confidence > 0.6:\n",
    "            if key_center == og_key_center:\n",
    "                final_key = f\"{final_key} {tonality}\"\n",
    "            elif og_key_center is None:\n",
    "                final_key = f\"{key_center} {tonality}\"\n",
    "        og_row[\"key\"] = final_key\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in og_meta[:25]:\n",
    "    if m[\"key\"] is not None:\n",
    "        print(m[\"key\"], m[\"tags\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### BPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bpm_std_to_time_difference(tempo_bpm, std_bpm):\n",
    "    # Beat periods in seconds\n",
    "    mean_period = 60 / tempo_bpm\n",
    "    slower_period = 60 / (tempo_bpm - 2 * std_bpm)  # tempo - 1 std dev\n",
    "    faster_period = 60 / (tempo_bpm + 2 * std_bpm)  # tempo + 1 std dev\n",
    "    \n",
    "    # Time differences from mean\n",
    "    time_diff_slower = slower_period - mean_period\n",
    "    time_diff_faster = mean_period - faster_period\n",
    "    \n",
    "    return time_diff_slower, time_diff_faster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "kept = 0\n",
    "total_found = 0\n",
    "\n",
    "folder_path = \"/app2/suno/data/sara/sfx_synthetic_tags/sfx_data_beats/\"\n",
    "\n",
    "for filename in os.listdir(folder_path):\n",
    "    if filename.endswith('.json'):\n",
    "        file_path = os.path.join(folder_path, filename)\n",
    "        data = read_json(file_path)\n",
    "        for id, row in data.items():\n",
    "            meta_row = metas_map[id]\n",
    "            tags = meta_row[\"tags\"]\n",
    "\n",
    "            meta_row[\"bpm\"] = None\n",
    "            og_bpm = None\n",
    "            if tags is None or len(tags) == 0:\n",
    "                pass\n",
    "            elif tags[-1].endswith(\" bpm\"):\n",
    "                og_bpm = tags[-1][:-len(\" bpm\")]\n",
    "                meta_row[\"tags\"] = tags[:-1]\n",
    "            \n",
    "            inferred_bpm = None\n",
    "            if not row[\"processing_failed\"]:\n",
    "                detected_bpm = row[\"inferred_tempo\"]\n",
    "                beats_detected = row[\"beats_detected\"]\n",
    "                tempo_std = row[\"tempo_std\"]\n",
    "\n",
    "                if detected_bpm is not None and beats_detected is not None and tempo_std is not None:\n",
    "                    std_time = max(bpm_std_to_time_difference(detected_bpm, tempo_std))\n",
    "                    total_found += 1\n",
    "                    if beats_detected >= 4 and std_time < 0.01:\n",
    "                        inferred_bpm = detected_bpm\n",
    "                        kept += 1\n",
    "\n",
    "            meta_row[\"bpm\"] = og_bpm if og_bpm is not None else inferred_bpm\n",
    "\n",
    "print(kept / total_found)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(og_meta[0].keys())\n",
    "for m in og_meta[:25]:\n",
    "    if m[\"bpm\"] is not None or m[\"key\"] is not None:\n",
    "        print(m[\"bpm\"], m[\"key\"], m[\"dataset\"], m[\"tags\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(og_meta, \"metas_filter_v0_add_key_bpm.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Merge In"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_tr_path = \"/app2/suno/data/diffusion/sfx/v0/metas_tr.jsonl\"\n",
    "metas_val_path = \"/app2/suno/data/diffusion/sfx/v0/metas_val.jsonl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_tr = read_jsonl(metas_tr_path)\n",
    "metas_val = read_jsonl(metas_val_path)\n",
    "og_meta = read_jsonl(\"/app2/suno/data/diffusion/sfx/v1/metas_filter_v0_add_key_bpm_to_tags.jsonl\")\n",
    "metas_map = {meta[\"id\"]: meta for meta in og_meta}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_tr[0].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "for row in metas_tr:\n",
    "    updated_tags = metas_map[row[\"id\"]]\n",
    "    row[\"tags\"] = updated_tags\n",
    "\n",
    "for row in metas_val:\n",
    "    updated_tags = metas_map[row[\"id\"]]\n",
    "    row[\"tags\"] = updated_tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(metas_tr, \"metas_tr_v1.jsonl\")\n",
    "write_jsonl(metas_val, \"metas_val_v1.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "for row in metas_tr[:10]:\n",
    "    print(row[\"tags\"])\n",
    "\n",
    "for row in metas_val[:10]:\n",
    "    print(row[\"tags\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(metas_tr))\n",
    "print(len(metas_val))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}