{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, load_json\n",
    "from suno_utils.audio import Audio\n",
    "\n",
    "from IPython.display import clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_covers = read_jsonl(\"filtered_cover_2025_03_11-15_39_26.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_centroids = []\n",
    "cover_centroids = []\n",
    "for source in filtered_covers:\n",
    "    weighted_centroid = source[\"centroid\"][1]\n",
    "    source_centroids.append(float(weighted_centroid))\n",
    "    for cover in source[\"covers\"]:\n",
    "        weighted_centroid_c = cover[\"centroid\"][1]\n",
    "        cover_centroids.append(float(weighted_centroid_c))\n",
    "        # if float(cover[\"similarity\"]) > 0.9:\n",
    "        #    print(cover[\"similarity\"], cover[\"cover_sim\"])\n",
    "        #    print(weighted_centroid, weighted_centroid_c)\n",
    "        #    Audio.from_s3(cover[\"s3_filepath\"]).play()\n",
    "        #    Audio.from_s3(source[\"s3_filepath\"]).play()\n",
    "        #    clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(source_centroids)\n",
    "print(cover_centroids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Create the plot\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.hist(source_centroids, bins=50, alpha=0.5, label=\"Source Audio\")\n",
    "plt.hist(cover_centroids, bins=50, alpha=0.5, label=\"Cover Audio\")\n",
    "\n",
    "# Add labels and legend\n",
    "plt.xlabel(\"Spectral Centroid (Hz)\")\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.title(\"Spectral Centroid Comparison\")\n",
    "plt.legend()\n",
    "\n",
    "# Display the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(cover_centroids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = read_jsonl(\"/app/suno/data/chirp_v4/base/metas_val.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for d in data:\n",
    "    if \"parent_id\" in d:\n",
    "        print(d.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for val in data:\n",
    "    print(val.keys())\n",
    "    # task = val['task']\n",
    "    # if task == \"covers\":\n",
    "    #    print(val.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "info = load_json(\"/app/suno/data/chirp_v5_ft/v2/info_val.json\")\n",
    "print(info.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covers = info[\"covers\"]\n",
    "print(covers[\"idx_map\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "info = read_jsonl(\"/app/suno/sara/metas_v0_cover.jsonl\")\n",
    "print(len(info))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for data in info[:30]:\n",
    "    print(data.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "by_parent = {}\n",
    "for data in info:\n",
    "    parent_id = data[\"parent_id\"]\n",
    "    child_id = data[\"child_id\"]\n",
    "    if parent_id not in by_parent:\n",
    "        by_parent[parent_id] = {}\n",
    "    by_parent[parent_id][child_id] = data[\"score_self_sim\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/app/suno/sara/self_sim_by_parent.json\", \"w\") as json_file:\n",
    "    json.dump(by_parent, json_file, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "import torch\n",
    "from suno_utils.worker.feature_eval import CentroidEval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test, sr = torchaudio.load(\"../samples/welcome_to_nowhere.wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval = CentroidEval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval.calculate_loudness(torch.zeros((1, 5000)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval.calculate_channel_diff(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval.calculate_centroid(test, sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval.calculate_channel_diff(\n",
    "    torch.stack(\n",
    "        [torch.mean(test, dim=0).unsqueeze(0), torch.mean(test, dim=0).unsqueeze(0)]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.data_loader import load_audio_mp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_path = \"../samples/test_train.webm\"\n",
    "result = load_audio_mp([sample_path], n_channels=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "import numpy as np\n",
    "\n",
    "SAMPLE_RATE = 16000\n",
    "\n",
    "import torch\n",
    "\n",
    "import torchaudio\n",
    "\n",
    "centroid_transform = torchaudio.transforms.SpectralCentroid(\n",
    "    sample_rate=SAMPLE_RATE,\n",
    "    hop_length=312,\n",
    ")\n",
    "\n",
    "\n",
    "def encode(arr: list):\n",
    "    arrays = []\n",
    "    for arr in arr:\n",
    "        if arr is None:\n",
    "            arrays.append(None)\n",
    "            continue\n",
    "\n",
    "        print(arr.shape[1] // 320)\n",
    "        # make mono\n",
    "        if arr.ndim == 2:\n",
    "            arr = arr.mean(axis=0)\n",
    "        arr = torch.nn.functional.unfold(\n",
    "            arr, kernel_size=(1, arr.shape[0] // 32), stride=(1, arr.shape[0] // 32)\n",
    "        )\n",
    "        print(arr)\n",
    "        rms_amp = torch.sqrt(torch.mean(arr**2))\n",
    "        epsilon = 1e-10\n",
    "        rms_db = 20 * torch.log10(rms_amp + epsilon)\n",
    "        rms_db = np.array([rms_db], dtype=np.float16)\n",
    "        arrays.append(rms_db)\n",
    "    return arrays\n",
    "\n",
    "\n",
    "results = encode(result)\n",
    "print(len(results))\n",
    "print(results[0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = np.load(\"centroid_raw/part_4.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 2
}