{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import torch\n",
    "import torchaudio\n",
    "from suno_utils.audio import Audio\n",
    "from itertools import combinations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_most_dissimilar_pair(mfccs_list):\n",
    "    \"\"\"\n",
    "    Find the most dissimilar pair of audio files using full MFCC sequences.\n",
    "    Assumes all MFCCs have the same time dimension length.\n",
    "    Args:\n",
    "        mfccs_list: List of MFCC tensors, each of shape (n_mfcc, time_steps)\n",
    "    Returns:\n",
    "        tuple: (index1, index2, distance) of the most dissimilar pair\n",
    "    \"\"\"\n",
    "    max_distance = -1\n",
    "    most_dissimilar_pair = (0, 0)\n",
    "    \n",
    "    # Compare all pairs\n",
    "    for (i, mfcc1), (j, mfcc2) in combinations(enumerate(mfccs_list), 2):\n",
    "        # Calculate Frobenius norm between full MFCC sequences\n",
    "        # This is equivalent to sqrt(sum of squared differences of all elements)\n",
    "        distance = torch.norm(mfcc1 - mfcc2, p='fro')\n",
    "        \n",
    "        if distance > max_distance:\n",
    "            max_distance = distance\n",
    "            most_dissimilar_pair = (i, j)\n",
    "    \n",
    "    return most_dissimilar_pair[0], most_dissimilar_pair[1], max_distance.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_filepaths = glob.glob(\"/home/christian/audio/reference-audio-wav/*.wav\")\n",
    "\n",
    "upsampled_audios = []\n",
    "for audio_filepath in audio_filepaths[:4]:\n",
    "    audio = Audio.from_file(audio_filepath)\n",
    "    upsampled_audios.append(audio)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "mfccs_list = []\n",
    "for audio in upsampled_audios:\n",
    "    audio_tensor = torch.from_numpy(audio.array_float)[:48000]\n",
    "    mfccs = torchaudio.transforms.MFCC(sample_rate=48000)(audio_tensor)\n",
    "    print(mfccs.shape)\n",
    "    mfccs_list.append(mfccs)\n",
    "\n",
    "i, j, dist = find_most_dissimilar_pair(mfccs_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(i, j, dist\n",
    "      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "meta_id = \"ee697b50-292e-4dbd-b019-2a20e5ca60dd\"\n",
    "basepath = f\"s3://suno-data/christian/data/upsample_v4_t_5_20241018/25hz_20241031_v1/{meta_id}\"\n",
    "pair_info = read_from_s3(f\"{basepath}/pair_info.json\")\n",
    "pair_info = json.loads(pair_info)\n",
    "print(pair_info)\n",
    "\n",
    "# listen to the two mp3 audio file\n",
    "\n",
    "\n",
    "a_audio = Audio.from_s3(f\"{basepath}/{pair_info['a_filename']}.mp3\", n_channels=2)  \n",
    "b_audio = Audio.from_s3(f\"{basepath}/{pair_info['b_filename']}.mp3\", n_channels=2)\n",
    "\n",
    "a_audio.play()\n",
    "b_audio.play()\n",
    "\n"
   ]
  },
  {
   "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
}