{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import IPython\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_dir = \"/mnt/localdisk/tmp/genius_hq_filtered_48khz\"\n",
    "\n",
    "audio_files = glob.glob(os.path.join(audio_dir, \"*.mp3\"))\n",
    "print(len(audio_files))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_stereo_width(waveform, sample_rate, window_size=4096, hop_length=1024):\n",
    "    \"\"\"\n",
    "    Compute the stereo width/spread of an audio file using PyTorch.\n",
    "    Returns a single value between 0 (mono) and 1 (maximum stereo spread).\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    audio_path : str\n",
    "        Path to the audio file\n",
    "    window_size : int\n",
    "        Size of the analysis window in samples\n",
    "    hop_length : int\n",
    "        Number of samples between successive windows\n",
    "        \n",
    "    Returns:\n",
    "    --------\n",
    "    float\n",
    "        Stereo width value between 0 and 1\n",
    "    dict\n",
    "        Additional metrics including correlation and phase correlation\n",
    "    \"\"\"\n",
    "\n",
    "    # Ensure stereo\n",
    "    if waveform.size(0) == 1:\n",
    "        raise ValueError(\"Audio file must be stereo (2 channels)\")\n",
    "    elif waveform.size(0) > 2:\n",
    "        waveform = waveform[:2, :]  # Take first two channels if more exist\n",
    "    \n",
    "    # Split into left and right channels\n",
    "    left = waveform[0]\n",
    "    right = waveform[1]\n",
    "    \n",
    "    # Compute cross-correlation\n",
    "    correlation = torch.corrcoef(torch.stack([left, right]))[0, 1].item()\n",
    "    \n",
    "    # Compute phase correlation using FFT\n",
    "    left_fft = torch.fft.rfft(left, dim=0)\n",
    "    right_fft = torch.fft.rfft(right, dim=0)\n",
    "    \n",
    "    # Compute cross-power spectrum\n",
    "    cross_power = left_fft * torch.conj(right_fft)\n",
    "    \n",
    "    # Normalize to get phase correlation\n",
    "    phase_correlation = torch.abs(cross_power) / (torch.abs(left_fft) * torch.abs(right_fft))\n",
    "    phase_correlation = phase_correlation.mean().item()\n",
    "    \n",
    "    # Compute mid/side representation\n",
    "    mid = (left + right) / 2\n",
    "    side = (left - right) / 2\n",
    "    \n",
    "    # Compute RMS energy of mid and side channels\n",
    "    mid_energy = torch.sqrt(torch.mean(mid ** 2))\n",
    "    side_energy = torch.sqrt(torch.mean(side ** 2))\n",
    "    \n",
    "    # Compute stereo width based on mid/side ratio\n",
    "    # Normalize to range 0-1 using sigmoid-like function\n",
    "    width_ratio = (side_energy / (mid_energy + 1e-8)).item()\n",
    "    stereo_width = 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5)\n",
    "    \n",
    "    metrics = {\n",
    "        'stereo_width': stereo_width,\n",
    "        'correlation': correlation,\n",
    "        'phase_correlation': phase_correlation,\n",
    "        'mid_energy': mid_energy.item(),\n",
    "        'side_energy': side_energy.item(),\n",
    "    }\n",
    "    \n",
    "    return stereo_width, metrics\n",
    "\n",
    "\n",
    "def compute_stereo_width_simple(waveform, sample_rate):\n",
    "    \"\"\"\n",
    "    Compute stereo width using a simple time-domain approach.\n",
    "    Returns a value between 0 (mono) and 1 (maximum stereo spread).\n",
    "    \n",
    "    The calculation is based on comparing the difference signal (L-R)\n",
    "    to the sum signal (L+R). A higher difference relative to the sum\n",
    "    indicates more stereo content.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    audio_path : str\n",
    "        Path to the audio file\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    float\n",
    "        Stereo width value between 0 and 1\n",
    "    \"\"\"\n",
    "    # Load audio file\n",
    "    \n",
    "    # Ensure stereo\n",
    "    if waveform.size(0) == 1:\n",
    "        raise ValueError(\"Audio file must be stereo (2 channels)\")\n",
    "    elif waveform.size(0) > 2:\n",
    "        waveform = waveform[:2, :]  # Take first two channels if more exist\n",
    "    \n",
    "    # Get left and right channels\n",
    "    left = waveform[0]\n",
    "    right = waveform[1]\n",
    "    \n",
    "    # Compute difference and sum signals\n",
    "    difference = left - right\n",
    "    sum_signal = left + right\n",
    "    \n",
    "    # Compute RMS (Root Mean Square) energy of both signals\n",
    "    diff_energy = torch.sqrt(torch.mean(difference ** 2))\n",
    "    sum_energy = torch.sqrt(torch.mean(sum_signal ** 2))\n",
    "    \n",
    "    # Compute width as ratio of difference to total energy\n",
    "    # Normalize to be between 0 and 1\n",
    "    width = (diff_energy / (sum_energy + 1e-8)).item()\n",
    "    width = min(width, 1.0)  # Clip to maximum of 1\n",
    "    \n",
    "    return width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pick a random file\n",
    "audio_file = audio_files[np.random.randint(0, len(audio_files))]\n",
    "waveform, sample_rate = torchaudio.load(audio_file)\n",
    "stereo_width = compute_stereo_width_simple(waveform, sample_rate)\n",
    "print(stereo_width)\n",
    "IPython.display.display(IPython.display.Audio(waveform, rate=sample_rate))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "width_results = {}\n",
    "for audio_file in tqdm(audio_files[:1000]):\n",
    "    waveform, sample_rate = torchaudio.load(audio_file)\n",
    "    stereo_width = compute_stereo_width_simple(waveform, sample_rate)\n",
    "    width_results[audio_file] = stereo_width\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make plot of stereo width distribution\n",
    "plt.hist(list(width_results.values()), bins=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# listen to the highest and lowest stereo width files\n",
    "# listen to the top 3 and bottom 3 stereo width files\n",
    "width_results = dict(sorted(width_results.items(), key=lambda item: item[1], reverse=True))\n",
    "for audio_file in list(width_results.keys())[:3]:\n",
    "    IPython.display.display(IPython.display.Audio(audio_file, rate=sample_rate))\n",
    "width_results = dict(sorted(width_results.items(), key=lambda item: item[1]))\n",
    "for audio_file in list(width_results.keys())[:3]:\n",
    "    IPython.display.display(IPython.display.Audio(audio_file, rate=sample_rate))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_terms = [\"Mono-like\", \"Near-mono\", \"Minimal stereo separation\", \"Center-dominant\", \"Collapsed stereo image\", \"Single-point focus\", \"Mono-prone soundstage\", \"Tight stereo spread\", \"Overlapping channels\", \"Mono-compatible\"]\n",
    "medium_terms = [\"Balanced stereo\", \"Moderate stereo separation\", \"Stereo with slight central focus\", \"Clear stereo imaging\", \"Evenly spaced stereo field\", \"Stereo realism\", \"Controlled stereo expansion\", \"Subtle stereo depth\", \"Stereo coherence\", \"Natural stereo balance\"]\n",
    "high_terms = [\"Expansive stereo\", \"Fully separated stereo channels\", \"Panoramic stereo image\", \"Immersive stereo spread\", \"Surround-like stereo field\", \"Highly diffused stereo\", \"Broad stereo imaging\", \"Wide stereo depth\", \"Enhanced stereo field\", \"Maximum stereo extension\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchaudio.transforms import MelSpectrogram\n",
    "def analyze_spectral_balance(waveform, sample_rate):\n",
    "    \"\"\"\n",
    "    Analyze the spectral balance of an audio file and determine if it's\n",
    "    bassy, mid-focused, or bright.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    audio_path : str\n",
    "        Path to the audio file\n",
    "    sample_rate : int\n",
    "        Target sample rate for analysis\n",
    "        \n",
    "    Returns:\n",
    "    --------\n",
    "    dict\n",
    "        Contains spectral ratios and character\n",
    "    \"\"\"\n",
    "    # Convert to mono if stereo\n",
    "    if waveform.size(0) > 1:\n",
    "        waveform = torch.mean(waveform, dim=0, keepdim=True)\n",
    "    \n",
    "    # Create mel spectrogram\n",
    "    mel_spec = MelSpectrogram(\n",
    "        sample_rate=sample_rate,\n",
    "        n_fft=2048,\n",
    "        hop_length=512,\n",
    "        n_mels=128,\n",
    "        f_min=20,\n",
    "        f_max=20000\n",
    "    )(waveform)\n",
    "    \n",
    "    # Convert to dB scale\n",
    "    mel_spec_db = torch.log10(mel_spec + 1e-9)\n",
    "    \n",
    "    # Calculate average energy in each frequency band\n",
    "    bass_energy = torch.mean(mel_spec_db[:, :40]).item()    # ~20-250 Hz\n",
    "    mid_energy = torch.mean(mel_spec_db[:, 40:80]).item()   # ~250-4000 Hz\n",
    "    high_energy = torch.mean(mel_spec_db[:, 80:]).item()    # ~4000-20000 Hz\n",
    "    \n",
    "    # Calculate relative ratios\n",
    "    total_energy = bass_energy + mid_energy + high_energy\n",
    "    bass_ratio = bass_energy / total_energy\n",
    "    mid_ratio = mid_energy / total_energy\n",
    "    high_ratio = high_energy / total_energy\n",
    "    \n",
    "    # Determine dominant characteristic\n",
    "    if bass_ratio > max(mid_ratio, high_ratio):\n",
    "        character = \"bassy\"\n",
    "    elif mid_ratio > max(bass_ratio, high_ratio):\n",
    "        character = \"mid_focused\"\n",
    "    else:\n",
    "        character = \"bright\"\n",
    "    \n",
    "    return character, bass_ratio, mid_ratio, high_ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {\n",
    "    \"bassy\": {},\n",
    "    \"mid_focused\": {},\n",
    "    \"bright\": {}\n",
    "}\n",
    "\n",
    "for audio_file in tqdm(audio_files[:100]):\n",
    "    waveform, sample_rate = torchaudio.load(audio_file)\n",
    "    character, bass_ratio, mid_ratio, high_ratio = analyze_spectral_balance(waveform, sample_rate)\n",
    "    results[character][audio_file] = {\n",
    "        \"bass_ratio\": bass_ratio,\n",
    "        \"mid_ratio\": mid_ratio,\n",
    "        \"high_ratio\": high_ratio\n",
    "    }\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# listen to the bassy file with the highest bass_ratio\n",
    "bass_results = dict(sorted(results[\"bassy\"].items(), key=lambda item: item[1][\"bass_ratio\"], reverse=True))\n",
    "for audio_file in list(bass_results.keys())[:3]:\n",
    "    print(audio_file, results[\"bassy\"][audio_file])\n",
    "    IPython.display.display(IPython.display.Audio(audio_file, rate=sample_rate))\n",
    "\n",
    "# listen to the mid_focused file with the highest mid_ratio\n",
    "#mid_results = dict(sorted(results[\"mid_focused\"].items(), key=lambda item: item[1][\"mid_ratio\"], reverse=True))\n",
    "#IPython.display.display(IPython.display.Audio(list(mid_results.keys())[0], rate=sample_rate))\n",
    "\n",
    "# listen to the bright file with the highest high_ratio\n",
    "bright_results = dict(sorted(results[\"bright\"].items(), key=lambda item: item[1][\"high_ratio\"], reverse=True))\n",
    "for audio_file in list(bright_results.keys())[:3]:\n",
    "    print(audio_file, results[\"bright\"][audio_file])\n",
    "    IPython.display.display(IPython.display.Audio(audio_file, rate=sample_rate))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make plot of spectral balance distribution\n",
    "# there are only 3 possible values, so we can make a bar chart\n",
    "# Count occurrences of each spectral balance category\n",
    "\n",
    "for character in results.keys():\n",
    "    print(f\"{character}: {len(results[character])}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bass_descriptors = [\n",
    "    \"bass-heavy\",\n",
    "    \"bottom-heavy\",\n",
    "    \"low-focused\",\n",
    "    \"deep\",\n",
    "    \"bassy\",\n",
    "    \"thick\",\n",
    "    \"boomy\",\n",
    "    \"warm\",\n",
    "    \"rich\",\n",
    "    \"full-bottom\"\n",
    "]\n",
    "\n",
    "mid_descriptors = [\n",
    "    \"mid-forward\",\n",
    "    \"mid-focused\",\n",
    "    \"mid-heavy\",\n",
    "    \"centered\",\n",
    "    \"neutral\",\n",
    "    \"balanced\",\n",
    "    \"present\",\n",
    "    \"pronounced\",\n",
    "    \"forward\",\n",
    "    \"clear\"\n",
    "]\n",
    "\n",
    "bright_descriptors = [\n",
    "    \"bright\",\n",
    "    \"treble-heavy\",\n",
    "    \"top-heavy\",\n",
    "    \"airy\",\n",
    "    \"crisp\",\n",
    "    \"sharp\",\n",
    "    \"brilliant\",\n",
    "    \"sparkling\",\n",
    "    \"light\",\n",
    "    \"thin\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyloudnorm as pyln\n",
    "\n",
    "def analyze_loudness_factor(waveform, sample_rate):\n",
    "    \"\"\"\n",
    "    Analyze loudness factor of an audio file and provide descriptive characteristics.\n",
    "    Loudness factor is the LUFS measurement after peak normalization.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    audio_path : str\n",
    "        Path to the audio file\n",
    "    target_peak_db : float\n",
    "        Target peak level in dB FS for normalization\n",
    "    sample_rate : int\n",
    "        Target sample rate for analysis\n",
    "        \n",
    "    Returns:\n",
    "    --------\n",
    "    dict\n",
    "        Contains loudness measurements and descriptors\n",
    "    \"\"\"\n",
    "\n",
    "    # peak normalize\n",
    "    normalized_audio = waveform / torch.max(torch.abs(waveform))\n",
    "    \n",
    "    # Measure LUFS\n",
    "    meter = pyln.Meter(sample_rate)\n",
    "    loudness_factor = meter.integrated_loudness(normalized_audio.permute(1, 0).numpy())\n",
    "    \n",
    "    # Categorize and select descriptors\n",
    "    if loudness_factor < -16:\n",
    "        category = \"dynamic\"\n",
    "    elif loudness_factor < -10:\n",
    "        category = \"moderate\"\n",
    "    else:\n",
    "        category = \"compressed\"\n",
    "    \n",
    "    return category, loudness_factor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loudness_results = {\n",
    "    \"dynamic\": {},\n",
    "    \"moderate\": {},\n",
    "    \"compressed\": {}\n",
    "}\n",
    "\n",
    "for audio_file in tqdm(audio_files[:250]):\n",
    "    waveform, sample_rate = torchaudio.load(audio_file)\n",
    "    category, loudness_factor = analyze_loudness_factor(waveform, sample_rate)\n",
    "    loudness_results[category][audio_file] = loudness_factor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print count of each category\n",
    "for category in loudness_results.keys():\n",
    "    print(f\"{category}: {len(loudness_results[category])}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a plot of the loudness factor distribution\n",
    "loudness_factors = [loudness_results[category][audio_file] for category in loudness_results.keys() for audio_file in loudness_results[category].keys()]\n",
    "plt.hist(loudness_factors, bins=20)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# listen to the most dynamic file\n",
    "dynamic_results = dict(sorted(loudness_results[\"dynamic\"].items(), key=lambda item: item[1], reverse=False))\n",
    "print(dynamic_results[list(dynamic_results.keys())[4]])\n",
    "IPython.display.display(IPython.display.Audio(list(dynamic_results.keys())[4], rate=sample_rate))\n",
    "\n",
    "\n",
    "# listen to the most compressed file\n",
    "compressed_results = dict(sorted(loudness_results[\"compressed\"].items(), key=lambda item: item[1], reverse=True))\n",
    "print(compressed_results[list(compressed_results.keys())[1]])\n",
    "IPython.display.display(IPython.display.Audio(list(compressed_results.keys())[1], rate=sample_rate))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamic_descriptors = [\"Dynamic\", \"Expressive\", \"Wide range\", \"High contrast\", \"Strong peaks\", \"Transient-heavy\", \"High crest factor\", \"Articulated shifts\", \"Uncompressed\", \"Pronounced amplitude variation\", \"Natural decay\", \"Ebb and flow\", \"Full dynamics\"]\n",
    "moderate_descriptors = [\"Balanced\", \"Smooth\", \"Moderate range\", \"Controlled peaks\", \"Even loudness\", \"Subtle shifts\", \"Standard compression\", \"Consistent energy\", \"Natural dynamics\", \"Slight transient smoothing\", \"Moderate peak-to-average ratio\", \"Even amplitude\", \"Gentle contrast\"]\n",
    "compressed_descriptors = [\"Compressed\", \"Flat\", \"Squashed\", \"Dense\", \"Brickwalled\", \"Low contrast\", \"Constant loudness\", \"Over-compressed\", \"Heavily limited\", \"Narrow dynamic range\", \"Maximized RMS\", \"Peaks suppressed\", \"Flat amplitude profile\", \"Reduced transients\"]"
   ]
  },
  {
   "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
}