{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import signal\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def gcc_phat(sig1, sig2, fs=1, max_tau=None, interp=1):\n",
    "    \"\"\"\n",
    "    Compute the offset between two signals using the GCC-PHAT algorithm.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    sig1 : array_like\n",
    "        First audio signal\n",
    "    sig2 : array_like\n",
    "        Second audio signal\n",
    "    fs : int, optional\n",
    "        Sampling frequency in Hz\n",
    "    max_tau : float, optional\n",
    "        Maximum time delay to consider (in samples)\n",
    "    interp : int, optional\n",
    "        Interpolation factor to improve time resolution\n",
    "        \n",
    "    Returns:\n",
    "    --------\n",
    "    tau_samples : int\n",
    "        Time delay in samples\n",
    "    cc : array_like\n",
    "        Cross-correlation function\n",
    "    \"\"\"\n",
    "    # Convert to numpy arrays if they aren't already\n",
    "    sig1 = np.array(sig1)\n",
    "    sig2 = np.array(sig2)\n",
    "    \n",
    "    # Ensure same length\n",
    "    min_length = min(len(sig1), len(sig2))\n",
    "    sig1 = sig1[:min_length]\n",
    "    sig2 = sig2[:min_length]\n",
    "    \n",
    "    # Fast Fourier Transform\n",
    "    X1 = np.fft.rfft(sig1)\n",
    "    X2 = np.fft.rfft(sig2)\n",
    "    \n",
    "    # Cross-spectrum\n",
    "    X1X2 = X1 * np.conj(X2)\n",
    "    \n",
    "    # PHAT weighting\n",
    "    X1X2_phat = X1X2 / (np.abs(X1X2) + 1e-10)  # Add small value to avoid division by zero\n",
    "    \n",
    "    # Optional: Interpolation for higher precision\n",
    "    if interp > 1:\n",
    "        X1X2_phat = np.concatenate([X1X2_phat, np.zeros(len(X1X2_phat) * (interp - 1))])\n",
    "    \n",
    "    # Inverse FFT to get cross-correlation\n",
    "    cc = np.fft.irfft(X1X2_phat)\n",
    "    \n",
    "    # Find the maximum of the cross-correlation\n",
    "    max_shift = len(cc) // 2\n",
    "    \n",
    "    if max_tau is not None:\n",
    "        max_samples = int(max_tau * interp)\n",
    "        max_shift = min(max_shift, max_samples)\n",
    "    \n",
    "    # Only consider shifts within the valid range\n",
    "    cc_valid = cc\n",
    "    \n",
    "    # find peaks in cc_valid\n",
    "    peaks = np.where(cc_valid > 0.05)[0]\n",
    "    print(\"peaks\", peaks)\n",
    "\n",
    "    if len(peaks) > 0:\n",
    "        max_idx = peaks[0]\n",
    "    else:\n",
    "        max_idx = 0\n",
    "\n",
    "    print(\"max_idx\", max_idx)\n",
    "    print(\"max_shift\", max_shift)\n",
    "\n",
    "    tau_samples = max_idx // interp\n",
    "    \n",
    "    return tau_samples, cc\n",
    "\n",
    "# Example usage with visualization\n",
    "def visualize_gcc_phat(sig1, sig2, fs=1):\n",
    "    \"\"\"\n",
    "    Visualize the GCC-PHAT results for two signals.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    sig1 : array_like\n",
    "        First audio signal\n",
    "    sig2 : array_like\n",
    "        Second audio signal\n",
    "    fs : int\n",
    "        Sampling frequency in Hz\n",
    "    \"\"\"\n",
    "    tau_samples, cc = gcc_phat(sig1, sig2, fs)\n",
    "    \n",
    "    # Create sample arrays for plotting\n",
    "    samples1 = np.arange(len(sig1))\n",
    "    samples2 = np.arange(len(sig2))\n",
    "    \n",
    "    # Create lag array for cross-correlation plot in samples\n",
    "    lags_samples = np.arange(-len(cc)//2, len(cc)//2)\n",
    "    \n",
    "    plt.figure(figsize=(12, 8))\n",
    "    \n",
    "    # Plot original signals\n",
    "    plt.subplot(3, 1, 1)\n",
    "    plt.plot(samples1, sig1, label='Signal 1')\n",
    "    plt.grid(True)\n",
    "    plt.legend()\n",
    "    plt.title('Input Signals')\n",
    "    plt.xlabel('Samples')\n",
    "    \n",
    "    plt.subplot(3, 1, 2)\n",
    "    plt.plot(samples2, sig2, label='Signal 2', color='orange')\n",
    "    plt.grid(True)\n",
    "    plt.legend()\n",
    "    plt.xlabel('Samples')\n",
    "    \n",
    "    # Plot cross-correlation\n",
    "    plt.subplot(3, 1, 3)\n",
    "    plt.plot(lags_samples, np.fft.fftshift(cc))\n",
    "    plt.grid(True)\n",
    "    plt.axvline(x=tau_samples, color='r', linestyle='--', \n",
    "                label=f'Sample Delay: {tau_samples}')\n",
    "    plt.legend()\n",
    "    plt.title('GCC-PHAT Cross-Correlation')\n",
    "    plt.xlabel('Lag (samples)')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    \n",
    "    print(f\"Estimated sample offset: {tau_samples} samples\")\n",
    "    if fs > 1:\n",
    "        print(f\"Estimated time delay: {tau_samples/fs:.6f} seconds\")\n",
    "    \n",
    "    return tau_samples\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "import torch\n",
    "\n",
    "orig, sr = torchaudio.load(\"/home/christian/code/christian/notebooks/audio/delay/stone.mp3\")\n",
    "remove_vox, sr = torchaudio.load(\"/home/christian/code/christian/notebooks/audio/delay/stone_removed_vocal.mp3\")\n",
    "\n",
    "\n",
    "#make mono\n",
    "orig = orig.mean(dim=0)\n",
    "remove_vox = remove_vox.mean(dim=0)\n",
    "\n",
    "# crop to 10 seconds\n",
    "orig = orig[:5*sr]\n",
    "remove_vox = remove_vox[:5*sr]\n",
    "\n",
    "# Apply GCC-PHAT\n",
    "estimated_delay = visualize_gcc_phat(orig, remove_vox, sr)\n",
    "\n",
    "print(f\"True delay: {delay:.6f} seconds\")\n",
    "print(f\"Error: {abs(estimated_delay - delay):.6f} seconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "orig, sr = torchaudio.load(\"/home/christian/code/christian/notebooks/audio/delay/stone.mp3\")\n",
    "remove_vox, sr = torchaudio.load(\"/home/christian/code/christian/notebooks/audio/delay/stone_removed_vocal.mp3\")\n",
    "\n",
    "#make mono\n",
    "orig = orig.mean(dim=0)\n",
    "remove_vox = remove_vox.mean(dim=0)\n",
    "\n",
    "print(orig.shape)\n",
    "print(remove_vox.shape)\n",
    "\n",
    "# delay the orignal by 14 samples\n",
    "delay_samples = -1765\n",
    "orig_delay = torch.roll(orig, delay_samples)\n",
    "# then zero out the first 14 samples\n",
    "#if delay_samples < 0:\n",
    "#    orig_delay[:delay_samples] = 0\n",
    "#else:\n",
    "#    orig_delay[delay_samples:] = 0\n",
    "\n",
    "\n",
    "two_seconds = int(2.2*sr)\n",
    "three_seconds = int(2.3*sr)\n",
    "\n",
    "# crop to 10 seconds\n",
    "orig = orig[:10*sr]\n",
    "orig_delay = orig_delay[:10*sr]\n",
    "remove_vox = remove_vox[:10*sr]\n",
    "\n",
    "sub_seg_remove_vox = remove_vox[two_seconds:three_seconds]\n",
    "sub_seg_orig = orig[two_seconds:three_seconds]\n",
    "sub_seg_orig_delay = orig_delay[two_seconds:three_seconds]\n",
    "\n",
    "# gcc\n",
    "visualize_gcc_phat(sub_seg_orig_delay, sub_seg_remove_vox, sr)\n",
    "\n",
    "# plot the orig and orig_delay\n",
    "#plt.plot(remove_vox[two_seconds:three_seconds])\n",
    "#plt.plot(orig_delay[two_seconds:three_seconds])\n",
    "#plt.plot(orig[two_seconds:three_seconds])\n",
    "#plt.show()\n",
    "\n",
    "\n",
    "\n",
    "# sum the two signals\n",
    "sum_signal = orig + remove_vox\n",
    "sum_signal_delay = orig_delay + remove_vox\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "ipd.display(ipd.Audio(sum_signal.numpy(), rate=sr))\n",
    "ipd.display(ipd.Audio(sum_signal_delay.numpy(), rate=sr))\n",
    "\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
}