{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import Audio, display\n",
    "import os\n",
    "import json\n",
    "from pathlib import Path\n",
    "import torchaudio\n",
    "import torch\n",
    "\n",
    "class AudioComparer:\n",
    "    def __init__(self, base_dir):\n",
    "        self.base_dir = Path(base_dir)\n",
    "        self.folders = [f for f in self.base_dir.iterdir() if f.is_dir()]\n",
    "        self.current_folder_idx = 0\n",
    "        self.results = {}\n",
    "        \n",
    "        # Load existing results if any\n",
    "        self.results_file = os.path.join(self.base_dir, 'comparison_results.json')\n",
    "        if os.path.exists(self.results_file):\n",
    "            with open(self.results_file, 'r') as f:\n",
    "                self.results = json.load(f)\n",
    "    \n",
    "    def compare_next_pair(self):\n",
    "        if self.current_folder_idx >= len(self.folders):\n",
    "            print(\"All comparisons complete!\")\n",
    "            return\n",
    "        \n",
    "        current_folder = self.folders[self.current_folder_idx]\n",
    "        audio_files = list(current_folder.glob('*.mp3'))  # adjust extension if needed\n",
    "        \n",
    "        if len(audio_files) != 2:\n",
    "            print(f\"Warning: Folder {current_folder} doesn't contain exactly 2 audio files\")\n",
    "            self.current_folder_idx += 1\n",
    "            return\n",
    "        \n",
    "        print(f\"\\nComparing files in folder: {current_folder.name}\")\n",
    "        \n",
    "        # Load and trim audio files\n",
    "        for audio_file in audio_files:\n",
    "            waveform, sample_rate = torchaudio.load(str(audio_file))\n",
    "            # Convert to mono if stereo\n",
    "            if waveform.shape[0] > 1:\n",
    "                waveform = torch.mean(waveform, dim=0, keepdim=True)\n",
    "            # Take first 30 seconds\n",
    "            num_samples = 30 * sample_rate\n",
    "            waveform = waveform[:, :num_samples]\n",
    "            # Convert to numpy for Audio display\n",
    "            audio_data = waveform.numpy()[0]  # take first channel\n",
    "            display(Audio(audio_data, rate=sample_rate, autoplay=True))\n",
    "        \n",
    "        # Create buttons\n",
    "        button_a = widgets.Button(description='Select A')\n",
    "        button_b = widgets.Button(description='Select B')\n",
    "        \n",
    "        def on_button_click(b):\n",
    "            choice = 'A' if b.description == 'Select A' else 'B'\n",
    "            self.results[current_folder.name] = {\n",
    "                'choice': choice,\n",
    "                'file_a': audio_files[0].name,\n",
    "                'file_b': audio_files[1].name\n",
    "            }\n",
    "            \n",
    "            # Save results\n",
    "            with open(self.results_file, 'w') as f:\n",
    "                json.dump(self.results, f, indent=2)\n",
    "            \n",
    "            self.current_folder_idx += 1\n",
    "            print(f\"Selected {choice}. Results saved.\")\n",
    "            \n",
    "            # Show next pair\n",
    "            self.compare_next_pair()\n",
    "        \n",
    "        button_a.on_click(on_button_click)\n",
    "        button_b.on_click(on_button_click)\n",
    "        \n",
    "        display(widgets.HBox([button_a, button_b]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "comparer = AudioComparer(\"/app/suno/christian/data/dpo_diffusion_test_set\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comparer.compare_next_pair()\n",
    "\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
}