{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import json\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first sync from s3 to local\n",
    "# aws s3 sync s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas/  /app/suno/christian/data/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas\n",
    "\n",
    "base_dir = \"/app/suno/christian/data/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas\"\n",
    "\n",
    "# find all the metadata files\n",
    "metadata_files = glob.glob(os.path.join(base_dir, \"*.json\"))\n",
    "print(len(metadata_files))\n",
    "\n",
    "# turn this into a list of ids\n",
    "ids = [os.path.basename(file).split(\"_metadata\")[0] for file in metadata_files]\n",
    "print(len(ids))\n",
    "#print(ids[:10])\n",
    "\n",
    "\n",
    "# aws s3 sync s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the discogs metadata\n",
    "work_items = read_jsonl(\n",
    "    \"/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl\"\n",
    ")\n",
    "\n",
    "work_items_map = {w[\"id\"]: w for w in work_items}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import concurrent.futures\n",
    "import time\n",
    "\n",
    "def download_mp3(example_id):\n",
    "    # get the s3_filepath from the work_items_map\n",
    "    s3_filepath = work_items_map[example_id][\"s3_filepath\"]\n",
    "    \n",
    "    # download the mp3 from s3 to the base_dir no output to console\n",
    "    os.system(f\"aws s3 cp {s3_filepath} {base_dir}/{example_id}_original.mp3 > /dev/null 2>&1\")\n",
    "\n",
    "# Use ThreadPoolExecutor which works better in notebooks\n",
    "max_workers = 64  # Lower number of workers to avoid overwhelming the system\n",
    "with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
    "    # Submit all tasks and create a dictionary mapping futures to their IDs\n",
    "    future_to_id = {executor.submit(download_mp3, example_id): example_id for example_id in ids}\n",
    "    \n",
    "    # Process results as they complete with a progress bar\n",
    "    completed = 0\n",
    "    pbar = tqdm(total=len(ids))\n",
    "    for future in concurrent.futures.as_completed(future_to_id):\n",
    "        completed += 1\n",
    "        pbar.update(1)\n",
    "        # Small sleep to allow UI updates in the notebook\n",
    "        if completed % 10 == 0:\n",
    "            time.sleep(0.1)\n",
    "    pbar.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a manifest file which is just a txt file with a newline separated list of ids\n",
    "# first we have to check if all the relevant files exist\n",
    "valid_ids = []\n",
    "for example_id in tqdm(ids):\n",
    "\n",
    "    # get the s3_filepath from the work_items_map\n",
    "    #s3_filepath = work_items_map[example_id][\"s3_filepath\"]\n",
    "    \n",
    "    # download the mp3 from s3 to the base_dir\n",
    "    # download the mp3 from s3 to the base_dir\n",
    "    #os.system(f\"aws s3 cp {s3_filepath} {base_dir}/{example_id}_original.mp3\")\n",
    "\n",
    "    # check if the codec mp3 exists\n",
    "    codec_mp3_path = os.path.join(base_dir, f\"{example_id}_codec.mp3\")\n",
    "    if not os.path.exists(codec_mp3_path):\n",
    "        print(f\"Codec mp3 does not exist for {example_id}\")\n",
    "        continue\n",
    "    # check if the diff_ctx.mp3 exists\n",
    "    diff_ctx_mp3_path = os.path.join(base_dir, f\"{example_id}_diff_ctx.mp3\")\n",
    "    if not os.path.exists(diff_ctx_mp3_path):\n",
    "        print(f\"Diff ctx mp3 does not exist for {example_id}\")\n",
    "        continue\n",
    "    # check if the diff_no_ctx.mp3 exists\n",
    "    diff_no_ctx_mp3_path = os.path.join(base_dir, f\"{example_id}_diff_no_ctx.mp3\")\n",
    "    if not os.path.exists(diff_no_ctx_mp3_path):\n",
    "        print(f\"Diff no ctx mp3 does not exist for {example_id}\")\n",
    "        continue\n",
    "    # add the id to the manifest\n",
    "    valid_ids.append(example_id)\n",
    "\n",
    "print(f\"Found {len(valid_ids)} valid ids out of {len(ids)}\")\n",
    "\n",
    "\n",
    "# now write the manifest file for train and validation\n",
    "for split in [\"tr\", \"val\"]:\n",
    "    manifest_file = os.path.join(\"/home/christian/code/christian/metadata/ear\", f\"upsample_manifest_{split}.txt\")\n",
    "    if split == \"tr\":\n",
    "        split_valid_ids = valid_ids[:int(len(valid_ids) * 0.95)]\n",
    "    else:\n",
    "        split_valid_ids = valid_ids[int(len(valid_ids) * 0.95):]\n",
    "\n",
    "    print(f\"Writing manifest for {split} with {len(split_valid_ids)} ids\")\n",
    "    with open(manifest_file, \"w\") as f:\n",
    "        for example_id in split_valid_ids:\n",
    "            f.write(f\"{example_id}\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_ids = []\n",
    "import shutil\n",
    "# load the example ids from the val set\n",
    "manifest_file = os.path.join(\"/home/christian/code/christian/metadata/ear\", f\"upsample_manifest_val.txt\")\n",
    "val_ids = [line.strip() for line in open(manifest_file, \"r\").readlines()]\n",
    "\n",
    "bench_base_dir = \"/home/christian/audio/ear-bench/diff-ctx\"\n",
    "os.makedirs(bench_base_dir, exist_ok=True)\n",
    "\n",
    "for example_id in tqdm(val_ids[:100]):\n",
    "    codec_mp3_path = os.path.join(base_dir, f\"{example_id}_codec.mp3\")\n",
    "    diff_ctx_mp3_path = os.path.join(base_dir, f\"{example_id}_diff_ctx.mp3\")\n",
    "    diff_no_ctx_mp3_path = os.path.join(base_dir, f\"{example_id}_diff_no_ctx.mp3\")\n",
    "\n",
    "    # copy the diff_ctx and diff_no_ctx mp3s to the bench_base_dir\n",
    "    shutil.copy(diff_ctx_mp3_path, os.path.join(bench_base_dir, f\"{example_id}_diff_ctx_cycled_pref=False.mp3\"))\n",
    "    shutil.copy(diff_no_ctx_mp3_path, os.path.join(bench_base_dir, f\"{example_id}_diff_no_ctx_input_pref=True.mp3\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = list_s3_dir(\"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_files = list_s3_dir(\n",
    "    f\"s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas\"\n",
    ")\n",
    "existing_files = [f[0] for f in existing_files if f[0].endswith(\"_metadata.json\")]\n",
    "existing_ids = [os.path.basename(f).split(\"_metadata.json\")[0] for f in existing_files]\n",
    "print(len(existing_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_ids[:10]"
   ]
  },
  {
   "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
}