{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import tempfile\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# copy all the original mp3 files from s3 to the local disk\n",
    "LOCAL_DIR = \"/mnt/localdisk/tmp_cjs/v2-infill-data-v1\"\n",
    "\n",
    "# source metas are here\n",
    "work_items = read_jsonl(\n",
    "    \"/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl\"\n",
    ")\n",
    "\n",
    "def download_mp3(work_item):\n",
    "    # get the s3 mp3 filepath \n",
    "    s3_mp3_filepath = work_item[\"s3_filepath\"]\n",
    "    meta_id = work_item[\"id\"]\n",
    "\n",
    "    # make the output directory\n",
    "    os.makedirs(os.path.join(LOCAL_DIR, f\"{meta_id}\"), exist_ok=True)\n",
    "\n",
    "    # copy to temp tirectory and then save with ffmpeg as mp3\n",
    "    with tempfile.TemporaryDirectory() as temp_dir:\n",
    "        local_mp3_filepath = os.path.join(LOCAL_DIR, f\"{meta_id}\", f\"{meta_id}_source.mp3\")\n",
    "        tmp_mp3_filepath = os.path.join(temp_dir, os.path.basename(s3_mp3_filepath))\n",
    "        file_ext = os.path.splitext(tmp_mp3_filepath)[1]\n",
    "        # download the mp3 file from s3\n",
    "        os.system(f\"aws s3 cp {s3_mp3_filepath} {tmp_mp3_filepath} > /dev/null 2>&1\")\n",
    "        # save with ffmpeg as mp3\n",
    "        # Create directory if it doesn't exist\n",
    "        os.makedirs(os.path.dirname(local_mp3_filepath), exist_ok=True)\n",
    "        # Use 192 kbps bitrate to preserve quality from the source opus audio\n",
    "        os.system(f\"ffmpeg -i {tmp_mp3_filepath} -c:a libmp3lame -b:a 192k -y {local_mp3_filepath} > /dev/null 2>&1\")\n",
    "\n",
    "# download the work items in parallel with joblib\n",
    "from joblib import Parallel, delayed\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Use tqdm to show progress\n",
    "results = Parallel(n_jobs=-1)(\n",
    "    delayed(download_mp3)(work_item) for work_item in tqdm(work_items, desc=\"Downloading files\")\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"/mnt/localdisk/tmp_cjs/v2-infill-data-v1/xUmVFrX5tus/xUmVFrX5tus_source.mp3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we need to create train an val metas which contain a list of the id and the filepaths for input and target\n",
    "\n",
    "target_name = \"source\"\n",
    "input_name = \"16n_v45_infill_ear_sft_3e5_t6_250k\"\n",
    "\n",
    "for split in [\"tr\", \"val\"]:\n",
    "\n",
    "    if split == \"tr\":\n",
    "        split_work_items = work_items[:9500]\n",
    "    else:\n",
    "        split_work_items = work_items[9500:]\n",
    "\n",
    "    print(f\"Processing {split} split with {len(split_work_items)} work items\")\n",
    "\n",
    "    new_metas = []\n",
    "    for work_item in split_work_items:\n",
    "        input_filepath = f\"{LOCAL_DIR}/{work_item['id']}/{work_item['id']}_{input_name}.mp3\"\n",
    "        target_filepath = f\"{LOCAL_DIR}/{work_item['id']}/{work_item['id']}_{target_name}.mp3\"\n",
    "        if os.path.exists(input_filepath) and os.path.exists(target_filepath):\n",
    "            new_metas.append({\n",
    "                \"id\": work_item[\"id\"],\n",
    "                \"input_filepath\": input_filepath,\n",
    "                \"target_filepath\": target_filepath\n",
    "            })\n",
    "\n",
    "    print(f\"Writing {len(new_metas)} metas to {f'/mnt/localdisk/tmp_cjs/v2-infill-data-v1/metas_{split}.jsonl'}\")\n",
    "    write_jsonl(new_metas, f\"/mnt/localdisk/tmp_cjs/v2-infill-data-v1/metas_{split}.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now check some metas\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "metas_val = read_jsonl(\"/mnt/localdisk/tmp_cjs/v2-infill-data-v1/metas_val.jsonl\")\n",
    "print(len(metas_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import auraloss\n",
    "import torchaudio\n",
    "import pyloudnorm as pyln\n",
    "\n",
    "freq_loss_fn = auraloss.freq.SumAndDifferenceSTFTLoss(\n",
    "    fft_sizes=[1024, 2048, 4096, 8192],\n",
    "    hop_sizes=[128, 256, 512, 1024],\n",
    "    win_lengths=[1024, 2048, 4096, 8192],\n",
    ")\n",
    "meter = pyln.Meter(48000)\n",
    "\n",
    "idx = 1\n",
    "meta = metas_val[idx]\n",
    "input_filepath = meta[\"input_filepath\"]\n",
    "input_audio, sr = torchaudio.load(input_filepath)\n",
    "target_filepath = meta[\"target_filepath\"]\n",
    "target_audio, sr = torchaudio.load(target_filepath)\n",
    "\n",
    "input_loudness = meter.integrated_loudness(input_audio.permute(1, 0).numpy())\n",
    "target_loudness = meter.integrated_loudness(target_audio.permute(1, 0).numpy())\n",
    "\n",
    "print(f\"Input loudness: {input_loudness}\")\n",
    "print(f\"Target loudness: {target_loudness}\")\n",
    "\n",
    "# normalize both to -16db\n",
    "input_audio_gain_db = -16 - input_loudness\n",
    "target_audio_gain_db = -16 - target_loudness\n",
    "\n",
    "input_audio = input_audio * 10**(input_audio_gain_db / 20)\n",
    "target_audio = target_audio * 10**(target_audio_gain_db / 20)\n",
    "\n",
    "# normalized loudness\n",
    "input_loudness = meter.integrated_loudness(input_audio.permute(1, 0).numpy())\n",
    "target_loudness = meter.integrated_loudness(target_audio.permute(1, 0).numpy())\n",
    "\n",
    "print(f\"Input loudness: {input_loudness}\")\n",
    "print(f\"Target loudness: {target_loudness}\")\n",
    "\n",
    "# crop to 10s\n",
    "chunk_size = 524288\n",
    "start_chunk = 5\n",
    "input_audio = input_audio[:, start_chunk*chunk_size:(start_chunk+1)*chunk_size]\n",
    "target_audio = target_audio[:, start_chunk*chunk_size:(start_chunk+1)*chunk_size]\n",
    "\n",
    "# compute the loss\n",
    "loss = freq_loss_fn(input_audio.unsqueeze(0), target_audio.unsqueeze(0))\n",
    "print(f\"freq loss: {loss}\")\n",
    "\n",
    "print(f\"Input audio shape: {input_audio.shape}\")\n",
    "print(input_audio.min(), input_audio.max())\n",
    "IPython.display.display(IPython.display.Audio(input_audio.cpu().numpy(), rate=sr, normalize=False))\n",
    "\n",
    "print(f\"Target audio shape: {target_audio.shape}\")\n",
    "print(target_audio.min(), target_audio.max())\n",
    "IPython.display.display(IPython.display.Audio(target_audio.cpu().numpy(), rate=sr, normalize=False))\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
}