{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import boto3\n",
    "import torch\n",
    "import numpy as np\n",
    "import torchaudio\n",
    "import tempfile\n",
    "\n",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_s3_files(bucket_name, prefix, max_keys: int = 100000):\n",
    "    all_files = []\n",
    "    continuation_token = None\n",
    "\n",
    "    while True:\n",
    "        # Prepare the arguments for the request\n",
    "        list_kwargs = {\n",
    "            \"Bucket\": bucket_name,\n",
    "            \"Prefix\": prefix,  # List objects under this prefix, or leave blank for all objects\n",
    "        }\n",
    "\n",
    "        if continuation_token:\n",
    "            list_kwargs[\"ContinuationToken\"] = continuation_token\n",
    "\n",
    "        # Make the request to list objects\n",
    "        response = s3.list_objects_v2(**list_kwargs)\n",
    "\n",
    "        # Collect the file keys\n",
    "        all_files += [obj[\"Key\"] for obj in response.get(\"Contents\", [])]\n",
    "\n",
    "        # Check if more results are available\n",
    "        if response.get(\"IsTruncated\"):  # True if there are more results to fetch\n",
    "            continuation_token = response[\"NextContinuationToken\"]\n",
    "        else:\n",
    "            break  # No more results to fetch\n",
    "\n",
    "    return all_files\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_name = \"suno-data\"\n",
    "base_dir = \"christian/data/upsample_100z_v1\"\n",
    "output_name = \"v2\"\n",
    "#\n",
    "# base_dir = \"christian/data/upsample_100hz_v4_t_5_20241018\"\n",
    "# output_name = \"v2\"\n",
    "\n",
    "base_metas_path = os.path.join(\"s3://\", bucket_name, base_dir, \"metas.jsonl\")\n",
    "base_metas = read_from_s3(base_metas_path, read_f=read_jsonl)\n",
    "print(len(base_metas))\n",
    "\n",
    "# s3 client\n",
    "s3 = boto3.client(\"s3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find all files on s3 with the pattern\n",
    "filepaths = get_s3_files(bucket_name, f\"{base_dir}/{output_name}\")\n",
    "print(\"total files on s3: \", len(filepaths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_to_s3_paths = {}\n",
    "# create a dict with the id as key and the s3 paths a list of values\n",
    "for filepath in tqdm(filepaths):\n",
    "    if \"-text_cfg_\" in filepath:\n",
    "        meta_id = filepath.split(\"-text_cfg\")[0].split(\"/\")[-1]\n",
    "    else:\n",
    "        meta_id = filepath.split(\"/\")[-1].split(\".\")[0].split(\"-\")[:-1]\n",
    "        meta_id = \"-\".join(meta_id)\n",
    "\n",
    "    if meta_id not in id_to_s3_paths:\n",
    "        id_to_s3_paths[meta_id] = set()\n",
    "\n",
    "    s3_filepath_basename = filepath.split(\".\")[0]\n",
    "    id_to_s3_paths[meta_id].add(s3_filepath_basename)\n",
    "\n",
    "print(\"total ids: \", len(id_to_s3_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display\n",
    "\n",
    "\n",
    "for meta_id, s3_paths in tqdm(id_to_s3_paths.items()):\n",
    "    # construct the s3 mp3 filepath and download to temporary directory\n",
    "\n",
    "    for s3_path in s3_paths:\n",
    "        if \"quality_scores\" not in s3_path:\n",
    "            s3_path = s3_path\n",
    "            break\n",
    "\n",
    "    print(s3_path)\n",
    "    full_path = f\"s3://{bucket_name}/{s3_path}.mp3\"\n",
    "    audio_data, sr = read_from_s3(full_path, read_f=torchaudio.load)\n",
    "    # apply highpass filter with torchaudio\n",
    "    highpass_cutoff = np.random.uniform(500, 4000)\n",
    "    audio_data_out = torchaudio.functional.highpass_biquad(audio_data, sr, highpass_cutoff)\n",
    "    # save to tmp dir\n",
    "    with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "        tmp_path = os.path.join(tmp_dir, f\"{meta_id}.mp3\")\n",
    "        torchaudio.save(tmp_path, audio_data_out, sr)\n",
    "        # upload to s3\n",
    "        output_path = f\"{s3_path}-highpass.mp3\"\n",
    "        print(output_path)\n",
    "\n",
    "        # listen to audio file with ipython\n",
    "        #IPython.display.display(IPython.display.Audio(tmp_path, rate=sr))\n",
    "        #IPython.display.display(IPython.display.Audio(audio_data, rate=sr))\n",
    "\n",
    "        s3.upload_file(tmp_path, bucket_name, output_path)\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
}