{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Take Five.wav\")\n",
    "# resample to 48k\n",
    "x = torchaudio.transforms.Resample(sr, 48000)(x)\n",
    "\n",
    "# crop to 15s\n",
    "x = x[:, :48000*10]\n",
    "print(x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_hann_window(frame_size: int, device=None) -> torch.Tensor:\n",
    "    \"\"\"Create a Hann window.\"\"\"\n",
    "    return torch.hann_window(frame_size, device=device)\n",
    "\n",
    "def frame_audio(x: torch.Tensor, frame_size: int, overlap: float = 0.0) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Split an audio tensor into frames. Apply Hann window only when overlap > 0.\n",
    "    \n",
    "    Args:\n",
    "        x: Input tensor of shape (batch_size, channels, sequence_length)\n",
    "        frame_size: Number of samples per frame\n",
    "        overlap: Overlap between consecutive frames\n",
    "    \n",
    "    Returns:\n",
    "        Framed tensor of shape (batch_size, channels, num_frames, frame_size)\n",
    "    \"\"\"\n",
    "    if not 0.0 <= overlap < 1.0:\n",
    "        raise ValueError(\"Overlap must be in [0.0, 1.0)\")\n",
    "    \n",
    "    batch_size, channels, seq_len = x.shape\n",
    "    hop_size = int(frame_size * (1 - overlap))\n",
    "    num_frames = (seq_len - frame_size) // hop_size + 1\n",
    "    \n",
    "    # Get window if needed\n",
    "    window = None\n",
    "    if overlap > 0:\n",
    "        window = get_hann_window(frame_size, device=x.device)\n",
    "    \n",
    "    # Create output tensor\n",
    "    output = torch.zeros(batch_size, channels, num_frames, frame_size, \n",
    "                        dtype=x.dtype, device=x.device)\n",
    "    \n",
    "    # Fill output tensor with frames\n",
    "    for i in range(num_frames):\n",
    "        start_idx = i * hop_size\n",
    "        frame = x[:, :, start_idx:start_idx + frame_size]\n",
    "        if window is not None:\n",
    "            frame = frame * window\n",
    "        output[:, :, i] = frame\n",
    "    \n",
    "    return output\n",
    "\n",
    "def reconstruct_audio(frames: torch.Tensor, original_length: int, overlap: float = 0.0) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Reconstruct audio signal from frames. Uses overlap-add only when overlap > 0.\n",
    "    \n",
    "    Args:\n",
    "        frames: Input tensor of shape (batch_size, channels, num_frames, frame_size)\n",
    "        original_length: Length of the original sequence\n",
    "        overlap: Overlap used in framing\n",
    "    \"\"\"\n",
    "    batch_size, channels, num_frames, frame_size = frames.shape\n",
    "    hop_size = int(frame_size * (1 - overlap))\n",
    "    \n",
    "    # For no overlap, we can just reshape\n",
    "    if overlap == 0:\n",
    "        # Check if the frames can be directly reshaped\n",
    "        expected_length = num_frames * frame_size\n",
    "        if expected_length == original_length:\n",
    "            return frames.reshape(batch_size, channels, -1)\n",
    "        else:\n",
    "            # If not exact match, still do frame-by-frame to handle partial frames\n",
    "            output = torch.zeros(batch_size, channels, original_length,\n",
    "                               dtype=frames.dtype, device=frames.device)\n",
    "            for i in range(num_frames):\n",
    "                start_idx = i * frame_size\n",
    "                end_idx = min(start_idx + frame_size, original_length)\n",
    "                output[:, :, start_idx:end_idx] = frames[:, :, i, :(end_idx-start_idx)]\n",
    "            return output\n",
    "    \n",
    "    # For overlap > 0, use overlap-add (Hann window sum to 1)\n",
    "    output = torch.zeros(batch_size, channels, original_length,\n",
    "                        dtype=frames.dtype, device=frames.device)\n",
    "    \n",
    "    for i in range(num_frames):\n",
    "        start_idx = i * hop_size\n",
    "        end_idx = start_idx + frame_size\n",
    "        output[:, :, start_idx:end_idx] += frames[:, :, i]\n",
    "    \n",
    "    return output[:, :, :original_length]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(x.shape)\n",
    "x_frames = frame_audio(x.unsqueeze(0), 1920, 0.0)\n",
    "print(x_frames.shape)\n",
    "\n",
    "x_reconstructed = reconstruct_audio(x_frames.half(), x.shape[1], 0.0)\n",
    "print(x_reconstructed.shape)\n",
    "\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(x.squeeze(0).numpy(), rate=48000))\n",
    "IPython.display.display(IPython.display.Audio(x_reconstructed.squeeze(0).numpy(), rate=48000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def frame_audio_with_dct(audio: torch.Tensor, frame_rate: int = 25) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Frame audio with 50% overlap and apply DCT to each frame.\n",
    "    \n",
    "    Args:\n",
    "        audio: Input audio tensor of shape (channels, samples) at 48kHz\n",
    "        frame_rate: Desired frame rate in Hz (frames per second)\n",
    "        \n",
    "    Returns:\n",
    "        Tensor of DCT coefficients for each frame (num_frames, frame_size)\n",
    "    \"\"\"\n",
    "    sample_rate = 48000\n",
    "    frame_size = sample_rate // frame_rate  # 1920 samples for 25 Hz\n",
    "    hop_size = frame_size // 2  # 50% overlap = 960 samples\n",
    "    \n",
    "    # Ensure audio is 2D (add channel dim if mono)\n",
    "    if audio.dim() == 1:\n",
    "        audio = audio.unsqueeze(0)\n",
    "    \n",
    "    # If stereo, average channels\n",
    "    if audio.size(0) > 1:\n",
    "        audio = audio.mean(0, keepdim=True)\n",
    "    \n",
    "    # Calculate number of frames\n",
    "    num_samples = audio.size(1)\n",
    "    num_frames = (num_samples - frame_size) // hop_size + 1\n",
    "    \n",
    "    # Create storage for frames\n",
    "    frames = torch.zeros(num_frames, frame_size, dtype=audio.dtype, device=audio.device)\n",
    "    \n",
    "    # Extract frames with overlap\n",
    "    for i in range(num_frames):\n",
    "        start = i * hop_size\n",
    "        frames[i] = audio[0, start:start + frame_size]\n",
    "    \n",
    "    # Apply Hann window to reduce spectral leakage\n",
    "    window = torch.hann_window(frame_size, dtype=audio.dtype, device=audio.device)\n",
    "    frames = frames * window\n",
    "    \n",
    "    # Apply DCT\n",
    "    dct_frames = torch.fft.rfft(frames).abs()\n",
    "    \n",
    "    return dct_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_dct = frame_audio_with_dct(x, 25)\n",
    "print(x_dct.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def visualize_frame_magnitudes(frames: torch.Tensor, threshold: float = 0.1):\n",
    "    \"\"\"\n",
    "    Visualize frame magnitudes and analyze low-value frames using pure matplotlib.\n",
    "    \n",
    "    Args:\n",
    "        frames: Tensor of shape (num_frames, frame_size) containing frame magnitudes\n",
    "        threshold: Threshold for considering a frame as \"low value\"\n",
    "    \"\"\"\n",
    "    # Convert to numpy for matplotlib\n",
    "    frame_data = frames.cpu().numpy()\n",
    "    \n",
    "    # Calculate frame energies (mean magnitude per frame)\n",
    "    frame_energies = np.mean(frame_data, axis=1)\n",
    "    \n",
    "    # Create figure with subplots\n",
    "    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 15))\n",
    "    \n",
    "    # 1. Heatmap of all frames\n",
    "    im = ax1.imshow(\n",
    "        np.log1p(frame_data.T),  # log scale for better visualization\n",
    "        aspect='auto',\n",
    "        cmap='viridis',\n",
    "        interpolation='nearest',\n",
    "        origin='lower'\n",
    "    )\n",
    "    plt.colorbar(im, ax=ax1, label='Log Magnitude')\n",
    "    ax1.set_title('Frame Magnitudes Over Time (Log Scale)')\n",
    "    ax1.set_xlabel('Frame Number')\n",
    "    ax1.set_ylabel('Frequency Bin')\n",
    "    \n",
    "    # 2. Plot frame energies\n",
    "    ax2.plot(frame_energies, 'b-', label='Frame Energy')\n",
    "    ax2.axhline(y=threshold, color='r', linestyle='--', label=f'Threshold ({threshold})')\n",
    "    ax2.set_title('Average Frame Energy Over Time')\n",
    "    ax2.set_xlabel('Frame Number')\n",
    "    ax2.set_ylabel('Mean Magnitude')\n",
    "    ax2.legend()\n",
    "    ax2.grid(True)\n",
    "    \n",
    "    # 3. Histogram of frame energies\n",
    "    ax3.hist(frame_energies, bins=50, color='b', edgecolor='black', alpha=0.7)\n",
    "    ax3.axvline(x=threshold, color='r', linestyle='--', label=f'Threshold ({threshold})')\n",
    "    ax3.set_title('Distribution of Frame Energies')\n",
    "    ax3.set_xlabel('Mean Magnitude')\n",
    "    ax3.set_ylabel('Count')\n",
    "    ax3.legend()\n",
    "    ax3.grid(True)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    \n",
    "    # Print statistics\n",
    "    low_value_frames = (frame_energies < threshold).sum()\n",
    "    total_frames = len(frame_energies)\n",
    "    print(f\"\\nFrame Analysis:\")\n",
    "    print(f\"Total frames: {total_frames}\")\n",
    "    print(f\"Low value frames (<{threshold}): {low_value_frames} ({low_value_frames/total_frames*100:.1f}%)\")\n",
    "    print(f\"Mean frame energy: {frame_energies.mean():.3f}\")\n",
    "    print(f\"Median frame energy: {np.median(frame_energies):.3f}\")\n",
    "    print(f\"Min frame energy: {frame_energies.min():.3f}\")\n",
    "    print(f\"Max frame energy: {frame_energies.max():.3f}\")\n",
    "    \n",
    "    return fig\n",
    "\n",
    "visualize_frame_magnitudes(x_dct)"
   ]
  },
  {
   "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
}