{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "from suno_utils.utils.text import read_jsonl\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(\n",
    "        batch_size, channels, num_frames, frame_size, dtype=x.dtype, device=x.device\n",
    "    )\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",
    "\n",
    "\n",
    "def reconstruct_audio(\n",
    "    frames: torch.Tensor, original_length: int, overlap: float = 0.0\n",
    ") -> torch.Tensor:\n",
    "    \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(\n",
    "                batch_size,\n",
    "                channels,\n",
    "                original_length,\n",
    "                dtype=frames.dtype,\n",
    "                device=frames.device,\n",
    "            )\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[\n",
    "                    :, :, i, : (end_idx - start_idx)\n",
    "                ]\n",
    "            return output\n",
    "\n",
    "    # For overlap > 0, use overlap-add (Hann window sum to 1)\n",
    "    output = torch.zeros(\n",
    "        batch_size, channels, original_length, dtype=frames.dtype, device=frames.device\n",
    "    )\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": [
    "DURATION_S = 10.0\n",
    "VAE_FRAME_RATE = 25\n",
    "SEMANTIC_FRAME_RATE = 25\n",
    "\n",
    "SEMANTIC_MEMMAP_SIZE = int(DURATION_S * SEMANTIC_FRAME_RATE)\n",
    "VAE_MEMMAP_SIZE = int(DURATION_S * VAE_FRAME_RATE * 2)\n",
    "VAE_DIM = 1920\n",
    "\n",
    "base_dir = \"/app/suno/christian/data/genius_hq_filtered_raw_10s_1920_memmap\"\n",
    "metas_path = os.path.join(base_dir, \"metas_tr.jsonl\")\n",
    "sem_memmap_path = os.path.join(base_dir, \"data_semantic_tr.bin\")\n",
    "vae_memmap_path = os.path.join(base_dir, \"data_vae_tr.bin\")\n",
    "\n",
    "metas = read_jsonl(metas_path)\n",
    "sem_memmap = np.memmap(sem_memmap_path, dtype=np.uint16, mode=\"r\")\n",
    "vae_memmap = np.memmap(vae_memmap_path, dtype=np.float16, mode=\"r\")\n",
    "\n",
    "# reshape the memmaps\n",
    "sem_memmap = sem_memmap.reshape(-1, SEMANTIC_MEMMAP_SIZE)\n",
    "vae_memmap = vae_memmap.reshape(-1, VAE_MEMMAP_SIZE, VAE_DIM)\n",
    "print(len(metas), sem_memmap.shape, vae_memmap.shape)\n",
    "\n",
    "# now iterate over the metas\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.zeros(VAE_MEMMAP_SIZE * VAE_DIM).std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_stds = []\n",
    "for idx, meta in tqdm(enumerate(metas)):\n",
    "    vae_memmap_frames = vae_memmap[idx]\n",
    "    # check for nan in vae_memmap_frames\n",
    "    frame_std = vae_memmap_frames.std()\n",
    "    frame_stds.append(frame_std)\n",
    "    if np.isnan(vae_memmap_frames).any():\n",
    "        print(f\"NaN found in vae_memmap_frames for meta {meta['id']}\")\n",
    "        break\n",
    "\n",
    "frame_stds = np.array(frame_stds)\n",
    "print(frame_stds.shape, frame_stds)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, frame_std in enumerate(frame_stds):\n",
    "    # check for nan\n",
    "    if np.isnan(frame_std):\n",
    "        print(f\"NaN found in frame_std for meta {idx} {meta['id']}\")\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the mean and std of the vae memmap\n",
    "# make sure to convert to float32\n",
    "#vae_memmap_mean = np.float32(vae_memmap[0]).mean(axis=(0, 1))\n",
    "vae_memmap_std = np.float32(vae_memmap).std()\n",
    "print(vae_memmap_std.shape, vae_memmap_std)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print( 1 / vae_memmap_std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.plot(vae_memmap_std * 8.5)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, len(metas))\n",
    "meta = metas[rand_idx]\n",
    "print(meta[\"tags\"])\n",
    "print(meta[\"text\"])\n",
    "sem_codes = sem_memmap[rand_idx]\n",
    "vae_frames = vae_memmap[rand_idx]\n",
    "\n",
    "print(sem_codes.shape, vae_frames.shape)\n",
    "vae_frames = torch.from_numpy(vae_frames.copy())\n",
    "vae_frames = vae_frames.view(1, 2, -1, VAE_DIM)\n",
    "print(vae_frames.shape)\n",
    "\n",
    "reconstructed = reconstruct_audio(vae_frames, 1920 * 25 * 10)\n",
    "print(reconstructed.shape)\n",
    "reconstructed = reconstructed.squeeze(0).cpu().numpy()\n",
    "\n",
    "IPython.display.Audio(reconstructed, rate=48000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.metrics import mean_squared_error\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def analyze_pca_compression(train_data, n_components_list, val_data = None):\n",
    "    \"\"\"\n",
    "    Analyze PCA compression with different numbers of components.\n",
    "    \n",
    "    Parameters:\n",
    "    train_data: numpy array of shape (n_samples, n_features)\n",
    "    n_components_list: list of int, different numbers of components to try\n",
    "    val_data: numpy array of shape (n_samples, n_features)\n",
    "    \n",
    "    Returns:\n",
    "    dict containing reconstruction errors and explained variance ratios\n",
    "    \"\"\"\n",
    "    results = {\n",
    "        'n_components': [],\n",
    "        'reconstruction_error': [],\n",
    "        'explained_variance_ratio': [],\n",
    "        'reconstruction_error_val': [],\n",
    "        'compression_ratio': [],\n",
    "        'compression_factor': [],\n",
    "        \"pca\": []\n",
    "    }\n",
    "    \n",
    "    original_dim = train_data.shape[1]\n",
    "    \n",
    "    for n_components in tqdm(n_components_list):\n",
    "        # Fit PCA\n",
    "        print(f\"Fitting PCA with {n_components} components\")\n",
    "        pca = PCA(n_components=n_components)\n",
    "        transformed = pca.fit_transform(train_data)\n",
    "        \n",
    "        # Reconstruct the data\n",
    "        reconstructed = pca.inverse_transform(transformed)\n",
    "        \n",
    "        # Calculate reconstruction error (MSE)\n",
    "        error = mean_squared_error(train_data, reconstructed)\n",
    "\n",
    "        # if val_data is not None, calculate the error on the validation data\n",
    "        if val_data is not None:\n",
    "            transformed_val = pca.transform(val_data)\n",
    "            reconstructed_val = pca.inverse_transform(transformed_val)\n",
    "            error_val = mean_squared_error(val_data, reconstructed_val)\n",
    "        \n",
    "        # Calculate compression ratio\n",
    "        compression_ratio = 1 - (n_components / original_dim)\n",
    "        compression_factor = original_dim / n_components\n",
    "        \n",
    "        # Store results\n",
    "        results['n_components'].append(n_components)\n",
    "        results['reconstruction_error'].append(error)\n",
    "        if val_data is not None:\n",
    "            results['reconstruction_error_val'].append(error_val)\n",
    "        results['explained_variance_ratio'].append(np.sum(pca.explained_variance_ratio_))\n",
    "        results['compression_ratio'].append(compression_ratio)\n",
    "        results['compression_factor'].append(compression_factor)\n",
    "        results[\"pca\"].append(pca)\n",
    "    return results\n",
    "\n",
    "def plot_pca_results(results):\n",
    "    \"\"\"Plot the PCA analysis results\"\"\"\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))\n",
    "    \n",
    "    # Plot reconstruction error vs compression ratio\n",
    "    ax1.plot(results['compression_factor'], results['reconstruction_error'], 'b.-')\n",
    "    if results['reconstruction_error_val'] is not None:\n",
    "        ax1.plot(results['compression_factor'], results['reconstruction_error_val'], 'g.-')\n",
    "        ax1.legend(['Train', 'Val'])\n",
    "    ax1.set_xlabel('Compression Factor')\n",
    "    ax1.set_ylabel('Reconstruction Error (MSE)')\n",
    "    ax1.set_title('Error vs Compression')\n",
    "    ax1.grid(True)\n",
    "    \n",
    "    # Plot explained variance ratio vs number of components\n",
    "    ax2.plot(results['n_components'], results['explained_variance_ratio'], 'r.-')\n",
    "    ax2.set_xlabel('Number of Components')\n",
    "    ax2.set_ylabel('Explained Variance Ratio')\n",
    "    ax2.set_title('Explained Variance vs Components')\n",
    "    ax2.grid(True)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    return fig\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_memmap_train_subset = vae_memmap[0:1000,...].reshape(-1, VAE_DIM).copy()\n",
    "print(vae_memmap_train_subset.shape)\n",
    "\n",
    "vae_memmap_val_subset = vae_memmap[1000:2000,...].reshape(-1, VAE_DIM).copy()\n",
    "print(vae_memmap_val_subset.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try different numbers of components\n",
    "n_components_list = [500, 600, 700, 800, 900, 1000]\n",
    "\n",
    "# Analyze PCA compression\n",
    "results = analyze_pca_compression(vae_memmap_train_subset, n_components_list, val_data=vae_memmap_val_subset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Plot results\n",
    "plot_pca_results(results)\n",
    "plt.show()\n",
    "\n",
    "# Print detailed results\n",
    "print(\"\\nDetailed Results:\")\n",
    "for i in range(len(results['n_components'])):\n",
    "    print(f\"\\nComponents: {results['n_components'][i]}\")\n",
    "    print(f\"Compression Ratio: {results['compression_ratio'][i]:.2f}\")\n",
    "    print(f\"Reconstruction Error: {results['reconstruction_error'][i]:.6f}\")\n",
    "    if results['reconstruction_error_val'] is not None:\n",
    "        print(f\"Reconstruction Error Val: {results['reconstruction_error_val'][i]:.6f}\")\n",
    "    print(f\"Explained Variance Ratio: {results['explained_variance_ratio'][i]:.3f}\")\n",
    "    print(f\"Compression Factor: {results['compression_factor'][i]:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's try to listen to the reconstructed audio from the pca\n",
    "#example_frames = vae_memmap[2004:2005,...]\n",
    "#print(example_frames.shape)\n",
    "\n",
    "# let's load a real audio file and then parse into frames \n",
    "audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Take Five.wav\"\n",
    "audio_filepath = \"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\"\n",
    "x, sr = torchaudio.load(audio_filepath)\n",
    "x = torchaudio.functional.resample(x, sr, 48000)\n",
    "\n",
    "# crop to 10 seconds\n",
    "start_frame = int(48000 * 160.0)\n",
    "x = x[:, start_frame:start_frame+ int(48000 * 10.01)]\n",
    "\n",
    "# now let's parse into frames\n",
    "example_frames = frame_audio(x.unsqueeze(0), 1920, 0.0)\n",
    "print(example_frames.shape)\n",
    "\n",
    "\n",
    "pca = results[\"pca\"]\n",
    "transformed = pca[-1].transform(example_frames.reshape(-1, VAE_DIM))\n",
    "print(transformed.shape)\n",
    "reconstructed_vae = pca[-1].inverse_transform(transformed)\n",
    "print(reconstructed_vae.shape)\n",
    "reconstructed_vae = reconstructed_vae.reshape(1, 2, -1, VAE_DIM)\n",
    "print(reconstructed_vae.shape)\n",
    "\n",
    "# turn frames into audio\n",
    "reconstructed_audio_original = reconstruct_audio(example_frames.view(1, 2, -1, VAE_DIM), 1920 * 25 * 10)\n",
    "reconstructed_audio_original = reconstructed_audio_original.squeeze(0).numpy()\n",
    "\n",
    "\n",
    "reconstructed_audio = reconstruct_audio(reconstructed_vae, 1920 * 25 * 10)\n",
    "reconstructed_audio = reconstructed_audio.squeeze(0)\n",
    "\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(reconstructed_audio_original, rate=48000))\n",
    "IPython.display.display(IPython.display.Audio(reconstructed_audio, rate=48000))\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
}