{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import json\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import pyloudnorm as pyln\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "\n",
    "base_dir = \"/app/suno/christian/data/dpo_diffusion_test_set_1k\"\n",
    "\n",
    "results = []\n",
    "with open(\"/home/christian/code/christian/results.jsonl\", \"r\") as f:\n",
    "    for line in f:\n",
    "        results.append(json.loads(line))\n",
    "print(len(results))\n",
    "\n",
    "# load features from genius and compute reference distributions\n",
    "csv_filepath = \"/home/christian/code/christian/metadata/genius_hq_audio_production_features_v2.csv\"\n",
    "df = pd.read_csv(csv_filepath)\n",
    "\n",
    "cols = [\"spectral_centroid\", \"bass\", \"mid\", \"high\", \"crest_factor\", \"stereo_width\", \"spectral_flatness\", \"silence_percentage\", \"loudness\"]\n",
    "\n",
    "features = {}\n",
    "means = []\n",
    "stds = []\n",
    "for col in cols:\n",
    "    features[col] = {\n",
    "        \"mean\": df[col].mean(),\n",
    "        \"std\": df[col].std()\n",
    "    }\n",
    "    means.append(features[col][\"mean\"])\n",
    "    stds.append(features[col][\"std\"])   \n",
    "\n",
    "print(means)\n",
    "print(stds)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_band_energy(waveform: torch.Tensor, \n",
    "                       sample_rate: int,\n",
    "                       n_fft: int = 4096) -> float:\n",
    "    bands = {\n",
    "        'bass': (20, 250),\n",
    "        'mid': (250, 2500),\n",
    "        'high': (2500, 20000)\n",
    "    }\n",
    "\n",
    "    # Convert to mono if stereo\n",
    "    if waveform.shape[0] > 1:\n",
    "        waveform = torch.mean(waveform, dim=0, keepdim=True)\n",
    "    \n",
    "    # Split into frames\n",
    "    frame_length = n_fft\n",
    "    hop_length = n_fft // 2  # 50% overlap\n",
    "    frames = waveform.unfold(1, frame_length, hop_length)\n",
    "    # Compute FFT for each frame\n",
    "    spectrum = torch.fft.rfft(frames.squeeze(0))  # [num_frames, n_fft//2 + 1]\n",
    "    freqs = torch.fft.rfftfreq(n_fft, d=1/sample_rate)  # [n_fft//2 + 1]\n",
    "    \n",
    "    # Compute magnitudes for each frame\n",
    "    magnitudes =torch.abs(spectrum)  # [num_frames, n_fft//2 + 1]\n",
    "    \n",
    "    # Compute centroid for each frame\n",
    "    numerator = torch.sum(freqs.view(1, -1) * magnitudes, dim=1)  # Sum over frequencies for each frame\n",
    "    denominator = torch.sum(magnitudes, dim=1)\n",
    "    \n",
    "    results = []\n",
    "\n",
    "    # Compute mean centroid across all frames\n",
    "    centroid = torch.mean(numerator / (denominator + 1e-8))\n",
    "    results.append(float(centroid))\n",
    "    \n",
    "    for band_name, (low_freq, high_freq) in bands.items():\n",
    "        # Create frequency mask\n",
    "        mask = (freqs >= low_freq) & (freqs <= high_freq)\n",
    "                \n",
    "        # They should now have the same size\n",
    "        band_energy = torch.mean(magnitudes * mask)\n",
    "        results.append(float(band_energy))\n",
    "\n",
    "    return results\n",
    "\n",
    "def calculate_stereo_width(waveform):\n",
    "    # Split into left and right channels\n",
    "    left = waveform[0]\n",
    "    right = waveform[1]\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",
    "    return stereo_width.item()\n",
    "\n",
    "def calculate_crest_factor(signal, window_size):\n",
    "    \"\"\"\n",
    "    Calculate RMS values for windows of audio data using PyTorch.\n",
    "    \n",
    "    Parameters:\n",
    "    signal (torch.Tensor): Audio signal tensor\n",
    "    window_size (int): Size of the window for RMS calculation\n",
    "    \n",
    "    Returns:\n",
    "    torch.Tensor: Tensor of RMS values\n",
    "    \"\"\"\n",
    "    # Ensure input is a tensor\n",
    "    if not isinstance(signal, torch.Tensor):\n",
    "        signal = torch.tensor(signal, dtype=torch.float32)\n",
    "    \n",
    "    # convert to mono if stereo\n",
    "    if signal.shape[0] > 1:\n",
    "        signal = signal.mean(dim=0)\n",
    "\n",
    "    # Calculate pad size\n",
    "    pad_size = window_size - (len(signal) % window_size)\n",
    "    if pad_size < window_size:\n",
    "        # Use constant padding (default value is 0)\n",
    "        padded_signal = torch.nn.functional.pad(signal, (0, pad_size))\n",
    "    else:\n",
    "        padded_signal = signal\n",
    "    \n",
    "    # Reshape signal into windows using unfold\n",
    "    # unfold(dimension, size, step) creates overlapping windows\n",
    "    # here we use step=size to create non-overlapping windows\n",
    "    windows = padded_signal.unfold(0, window_size, window_size)\n",
    "    \n",
    "    # Calculate RMS for each window\n",
    "    # torch.mean along dim=1 averages across the window\n",
    "    # keepdim=False reduces the dimension\n",
    "    window_rms = torch.sqrt(torch.mean(windows**2, dim=1))\n",
    "    rms = torch.mean(window_rms).item()\n",
    "\n",
    "    # get the peak value\n",
    "    peak = torch.max(torch.abs(signal)).item()\n",
    "\n",
    "    # compute crest factor\n",
    "    crest_factor = peak / (rms + 1e-8)\n",
    "\n",
    "    return np.log(crest_factor + 1e-8)\n",
    "\n",
    "import torch\n",
    "import torchaudio\n",
    "\n",
    "def measure_silence_percentage(\n",
    "    waveform: torch.Tensor,\n",
    "    sample_rate: int,\n",
    "    silence_threshold_db: float = -60,\n",
    "    window_size_ms: int = 100\n",
    ") -> float:\n",
    "    \"\"\"\n",
    "    Measures the percentage of audio that could be considered silent.\n",
    "    \n",
    "    Args:\n",
    "        file_path: Path to audio file\n",
    "        silence_threshold_db: RMS threshold in dB below which audio is considered silent\n",
    "        window_size_ms: Size of analysis window in milliseconds\n",
    "    \n",
    "    Returns:\n",
    "        Percentage (0-100) of audio that is below the silence threshold\n",
    "    \"\"\"\n",
    "    \n",
    "    # Convert to mono if stereo\n",
    "    if waveform.shape[0] > 1:\n",
    "        waveform = torch.mean(waveform, dim=0, keepdim=True)\n",
    "    \n",
    "    # Calculate window size in samples\n",
    "    window_size = int(sample_rate * window_size_ms / 1000)\n",
    "    \n",
    "    # Unfold the waveform into windows\n",
    "    windows = waveform.unfold(1, window_size, window_size)\n",
    "    \n",
    "    # Calculate RMS for each window\n",
    "    rms = torch.sqrt(torch.mean(windows ** 2, dim=2))\n",
    "    db = 20 * torch.log10(rms + 1e-10)\n",
    "    \n",
    "    # Calculate percentage of windows below threshold\n",
    "    silence_percentage = 100 * torch.mean((db < silence_threshold_db).float()).item()\n",
    "    \n",
    "    return silence_percentage\n",
    "\n",
    "\n",
    "def measure_spectral_flatness(audio_tensor):\n",
    "    \"\"\"\n",
    "    Calculate spectral flatness (Wiener entropy) of the signal.\n",
    "    Returns value between 0 (pure tone) and 1 (white noise).\n",
    "    \n",
    "    Parameters:\n",
    "    audio_tensor: Input audio tensor of shape [..., samples]\n",
    "    \n",
    "    Returns:\n",
    "    torch.Tensor: Spectral flatness value between 0 and 1\n",
    "    \"\"\"\n",
    "    # Get spectrum magnitude\n",
    "    audio_tensor = audio_tensor.mean(dim=0)\n",
    "    spectrum = torch.abs(torch.fft.rfft(audio_tensor, dim=-1))\n",
    "    \n",
    "    # Add small epsilon to avoid log(0)\n",
    "    epsilon = 1e-8\n",
    "    spectrum = spectrum + epsilon\n",
    "    \n",
    "    # Calculate geometric mean and arithmetic mean\n",
    "    log_spectrum = torch.log(spectrum)\n",
    "    geometric_mean = torch.exp(torch.mean(log_spectrum, dim=-1))\n",
    "    arithmetic_mean = torch.mean(spectrum, dim=-1)\n",
    "    \n",
    "    # Compute flatness\n",
    "    flatness = geometric_mean / (arithmetic_mean + 1e-8)\n",
    "    \n",
    "    return flatness.item()\n",
    "\n",
    "def get_features(audio_tensor, sample_rate, target_loudness=-16.0):\n",
    "    meter = pyln.Meter(sample_rate)\n",
    "    # loudness normalization\n",
    "    loudness = meter.integrated_loudness(audio_tensor.permute(1, 0).numpy())\n",
    "    # check if loudness is -inf\n",
    "    if loudness == -np.inf:\n",
    "        loudness = -80.0\n",
    "    loudness_diff = target_loudness - loudness\n",
    "    # limit the loudness difference to +/- 20 db\n",
    "    loudness_diff = np.clip(loudness_diff, -20, 20)\n",
    "    audio_tensor *= 10 ** (loudness_diff / 20.0)\n",
    "\n",
    "    # compute features\n",
    "    with torch.no_grad():\n",
    "        results_spectral = compute_band_energy(audio_tensor, sample_rate)\n",
    "        results_crest_factor = calculate_crest_factor(audio_tensor, 1024)\n",
    "        results_stereo_width = calculate_stereo_width(audio_tensor)\n",
    "        results_spectral_flatness = measure_spectral_flatness(audio_tensor)\n",
    "        results_silence_percentage = measure_silence_percentage(audio_tensor, sample_rate)\n",
    "\n",
    "    results = results_spectral + [results_crest_factor, results_stereo_width, results_spectral_flatness, results_silence_percentage, loudness]\n",
    "    return results\n",
    "\n",
    "def get_score(features):\n",
    "    features_normalized = (features - np.array(means)) / np.array(stds)\n",
    "\n",
    "    # zero out all but the first feature\n",
    "    features_normalized[-2] = 0\n",
    "    features_normalized[-1] = 0\n",
    "    score = np.mean(np.abs(features_normalized))\n",
    "    return score\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_a, sr_a = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n",
    "audio_a = audio_a[:, :int(sr_a * 120)]\n",
    "\n",
    "features_a = get_features(audio_a, sr_a)\n",
    "\n",
    "print(features_a)\n",
    "features_a_normalized = (features_a - np.array(means)) / np.array(stds)\n",
    "print(features_a_normalized)\n",
    "score = np.mean(np.abs(features_a_normalized))\n",
    "print(score)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = []\n",
    "pbar = tqdm(results)\n",
    "\n",
    "for result in pbar:\n",
    "    # get the audio files\n",
    "    audio_a = os.path.join(base_dir, result[\"request_id\"], result[\"audio_a_id\"] + \".mp3\")\n",
    "    audio_b = os.path.join(base_dir, result[\"request_id\"], result[\"audio_b_id\"] + \".mp3\")\n",
    "    # load the audio files\n",
    "    audio_a, sr_a = torchaudio.load(audio_a)\n",
    "    audio_b, sr_b = torchaudio.load(audio_b)\n",
    "\n",
    "    if audio_a.shape[-1] < sr_a * 120:\n",
    "        continue\n",
    "\n",
    "    if audio_b.shape[-1] < sr_b * 120:\n",
    "        continue\n",
    "\n",
    "    # crop to 30s of audio\n",
    "    start_s = 0\n",
    "    end_s = 30\n",
    "    audio_a = audio_a[:, start_s * sr_a:end_s * sr_a]\n",
    "    audio_b = audio_b[:, start_s * sr_b:end_s * sr_b]\n",
    "   \n",
    "    # get the result\n",
    "    true_label = result[\"selected\"]\n",
    "\n",
    "    # get the features\n",
    "    features_a = get_features(audio_a, sr_a)\n",
    "    features_b = get_features(audio_b, sr_b)\n",
    "    feature_diff = np.abs(np.array(features_a) - np.array(features_b))\n",
    "\n",
    "\n",
    "    # compute score by suming the features\n",
    "    # lower score is better\n",
    "    score_a = get_score(features_a)\n",
    "    score_b = get_score(features_b)\n",
    "\n",
    "    # get the models choice\n",
    "    model_choice = \"a\" if score_a < score_b else \"b\"\n",
    "\n",
    "    print(features_a)\n",
    "    print(features_b)\n",
    "    print(means)\n",
    "    print(\"model_choice\", model_choice)\n",
    "    print(\"true_label\", true_label)\n",
    "\n",
    "    # check if the model choice is correct\n",
    "    correct = model_choice == true_label\n",
    "    predictions.append(correct)\n",
    "    pbar.set_postfix({\"accuracy\": np.mean(predictions) * 100, \"score_a\": score_a, \"score_b\": score_b})\n",
    "\n",
    "print(np.mean(predictions))\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
}