{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import labelbox as lb\n",
    "from datetime import datetime\n",
    "from uuid import uuid4\n",
    "\n",
    "from langdetect import detect\n",
    "from bs4 import BeautifulSoup\n",
    "from suno_utils.utils.s3 import check_s3_file_exists\n",
    "\n",
    "import tempfile\n",
    "from suno_utils.utils.s3 import upload_s3_files, download_s3_files\n",
    "\n",
    "import re\n",
    "\n",
    "pattern = r\"^([^_]+)_format_trimmed_([^_]+)$\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inst_pairs = pd.read_csv(\"/app2/suno/data/sara/musdb/audio_comparison2/vox_pairs_trimmed.txt\")\n",
    "bass_pairs = pd.read_csv(\"/app2/suno/data/sara/musdb/audio_comparison2/bass_pairs_trimmed.txt\")\n",
    "drums_pairs = pd.read_csv(\"/app2/suno/data/sara/musdb/audio_comparison2/drums_pairs_trimmed.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs = pd.concat([inst_pairs, bass_pairs, drums_pairs], ignore_index=True).sample(frac=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs[\"model_a\"] = pairs[\"source_a\"].str.extract(pattern)[0]\n",
    "pairs[\"model_b\"] = pairs[\"source_b\"].str.extract(pattern)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timestamp = datetime.now().strftime(\"%Y%m%d\")\n",
    "test_label = f\"more-stems-trimmed-{timestamp}\"\n",
    "\n",
    "s3_bucket = f\"s3://suno-annotation-public/preference-{test_label}/\"  # this must have trailing slash\n",
    "s3_bucket_url = f\"https://suno-annotation-public.s3.amazonaws.com/preference-{test_label}\"\n",
    "output_html_dir = os.path.abspath(f\"./outputs/html-{test_label}\")\n",
    "os.makedirs(output_html_dir, exist_ok=True)\n",
    "base_html_filepath = \"./templates/ab_stems.html\"\n",
    "max_num_examples = 500\n",
    "global_keys = []\n",
    "\n",
    "# open the html file\n",
    "# Read the HTML file\n",
    "with open(base_html_filepath, \"r\", encoding=\"utf-8\") as file:\n",
    "    soup = BeautifulSoup(file, \"html.parser\")\n",
    "\n",
    "metadata = {}\n",
    "\n",
    "assets = []  # list of assets to add to the dataset\n",
    "\n",
    "local_paths_to_write = []\n",
    "s3_paths_to_write = []\n",
    "html_local = []\n",
    "html_to_write = []\n",
    "for i in range(0, len(pairs)):\n",
    "    row = pairs.iloc[i]\n",
    "    clip_0_path = row.source_a_fp\n",
    "    clip_1_path = row.source_b_fp\n",
    "    mixture_path = row.mixture_path\n",
    "    source_a = row.source_a\n",
    "    source_b = row.source_b\n",
    "    song_name = row.song_name\n",
    "    instrument = row.instrument\n",
    "\n",
    "    request_id = str(uuid4())  # f\"{song_name}_{source_a}_{source_b}\"\n",
    "    metadata[request_id] = row.to_dict()\n",
    "\n",
    "    # Find the audio sections and update their titles\n",
    "    audio_sections = soup.find_all(\"div\", class_=\"audio-section\")\n",
    "\n",
    "    # Update the h2 text for each audio section\n",
    "    # The first audio-section after reference-container is \"Instrumental Stem A\"\n",
    "    # The second is \"Instrumental Stem B\"\n",
    "    instrumental_sections = [\n",
    "        section\n",
    "        for section in audio_sections\n",
    "        if section.find(\"audio\", id=lambda x: x in [\"audio1\", \"audio2\"])\n",
    "    ]\n",
    "\n",
    "    if len(instrumental_sections) >= 2:\n",
    "        # Update Instrumental Stem A\n",
    "        instrumental_sections[0].find(\"h2\").string = f\"{instrument.upper()} Stem A\"\n",
    "        # Update Instrumental Stem B\n",
    "        instrumental_sections[1].find(\"h2\").string = f\"{instrument.upper()} Stem B\"\n",
    "\n",
    "    # edit the html file to add audios and prompt information\n",
    "    soup.find(id=\"audioRef\")[\"src\"] = f\"{s3_bucket_url}/{request_id}_ref.wav\"\n",
    "    soup.find(id=\"audio1\")[\"src\"] = f\"{s3_bucket_url}/{request_id}_{source_a}.wav\"\n",
    "    soup.find(id=\"audio2\")[\"src\"] = f\"{s3_bucket_url}/{request_id}_{source_b}.wav\"\n",
    "\n",
    "    # Save the modified HTML back to disk\n",
    "    output_filepath = os.path.join(output_html_dir, f\"{request_id}.html\")\n",
    "    with open(output_filepath, \"w\", encoding=\"utf-8\") as fp:\n",
    "        fp.write(str(soup))\n",
    "\n",
    "    s3_html_filepath = f\"{s3_bucket}{request_id}.html\"\n",
    "    s3_html_url = f\"{s3_bucket_url}/{request_id}.html\"\n",
    "    html_local.append(output_filepath)\n",
    "    html_to_write.append(s3_html_filepath)\n",
    "\n",
    "    local_paths_to_write.append(mixture_path)\n",
    "    local_paths_to_write.append(clip_0_path)\n",
    "    local_paths_to_write.append(clip_1_path)\n",
    "    s3_paths_to_write.append(f\"{s3_bucket}{request_id}_ref.wav\")\n",
    "    s3_paths_to_write.append(f\"{s3_bucket}{request_id}_{source_a}.wav\")\n",
    "    s3_paths_to_write.append(f\"{s3_bucket}{request_id}_{source_b}.wav\")\n",
    "\n",
    "    # add to data row list\n",
    "    assets.append(\n",
    "        {\n",
    "            \"row_data\": s3_html_url,\n",
    "            \"global_key\": request_id,\n",
    "        }\n",
    "    )\n",
    "\n",
    "    global_keys.append(request_id)\n",
    "\n",
    "    if len(assets) >= max_num_examples:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_s3_files(html_local, html_to_write, extra_args={\"ContentType\": \"text/html\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_s3_files(local_paths_to_write, s3_paths_to_write, extra_args={\"ContentType\": \"audio/wav\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save json metadata\n",
    "metadata_filepath = f\"./outputs/metadata-{test_label}.json\"\n",
    "with open(metadata_filepath, \"w\") as fp:\n",
    "    json.dump(metadata, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dataset\n",
    "client = lb.Client(\n",
    "    api_key=\"eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOiJjbHppbmJnY2wwMDYyMDd5bWg2enhiMTd6Iiwib3JnYW5pemF0aW9uSWQiOiJjbHppbmJnY2QwMDYxMDd5bWM4cjM5cHdzIiwiYXBpS2V5SWQiOiJjbHp3cmU3ZnQwYjFzMDd6aWRrNTFnb21mIiwic2VjcmV0IjoiMGU5M2MwNmU4ZWI1Y2Y3NTlmNTk5YTk5MzIwOTU5MzQiLCJpYXQiOjE3MjM4MTU3NTcsImV4cCI6MjM1NDk2Nzc1N30.Z6gZwlzQ85KrqGOydqCdo1RVmUhOEntpev2HhNso4PU\"\n",
    ")\n",
    "dataset = client.create_dataset(name=\"more-stems-preference-test\")\n",
    "\n",
    "# Bulk add data rows to the dataset\n",
    "task = dataset.create_data_rows(assets)\n",
    "task.wait_till_done()\n",
    "print(task.errors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ontology_builder = lb.OntologyBuilder(\n",
    "    classifications=[\n",
    "        lb.Classification(\n",
    "            class_type=lb.Classification.Type.RADIO,\n",
    "            name=\"preference\",\n",
    "            instructions=\"Preference\",\n",
    "            options=[\n",
    "                lb.Option(value=\"A\"),\n",
    "                lb.Option(value=\"B\"),\n",
    "            ],\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "ontology = client.create_ontology(\n",
    "    \"Ontology HTML Annotations\", ontology_builder.asdict(), media_type=lb.MediaType.Html\n",
    ")\n",
    "\n",
    "project = client.create_project(name=f\"stems-trimmed-preference\", media_type=lb.MediaType.Html)\n",
    "\n",
    "# Setup your ontology\n",
    "project.connect_ontology(ontology)\n",
    "\n",
    "# send rows to project\n",
    "batch = project.create_batch(\n",
    "    \"first-batch-html-demo\",  # Each batch in a project must have a unique name\n",
    "    global_keys=global_keys,  # Paginated collection of data row objects, list of data row ids or global keys\n",
    "    priority=5,  # priority between 1(highest) - 5(lowest)\n",
    ")\n",
    "\n",
    "print(\"Batch: \", batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs.tail()"
   ]
  },
  {
   "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
}
