{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6211131a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "# given generations with multiple seeds, load the metadata and pick the best pairs\n",
    "base_dir = \"/app2/suno/data/christian/outputs/v3-base-data-ctx-t3\"\n",
    "model_name = \"v3_flow_sft_t8_rd1_pair_t3_1E6_beta100_n16_bt2_noise_1k_last\"\n",
    "#get all the dirs\n",
    "dirs = os.listdir(base_dir)\n",
    "dirs = [d for d in dirs if os.path.isdir(os.path.join(base_dir, d))]\n",
    "print(len(dirs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b69feb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_dir(base_dir, dirname):\n",
    "    results = {}\n",
    "    semantic_codes_filepath = os.path.join(base_dir, dirname, f\"{dirname}_semantic.npz\")\n",
    "    for n in range(10):\n",
    "        upsampled_vae_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_{n}_upsampled_vae.npz\")\n",
    "        metadata_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_{n}__metadata.npz\")\n",
    "        mp3_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_{n}.mp3\")\n",
    "        # check if the metadata file exists\n",
    "        if not os.path.exists(metadata_filepath):\n",
    "            continue\n",
    "        with np.load(metadata_filepath, allow_pickle=True) as data:\n",
    "            metadata_npz = dict(data)\n",
    "        if metadata_npz is None:\n",
    "            continue\n",
    "        metadata_dict = {key: metadata_npz[key].tolist() for key in metadata_npz.keys()}\n",
    "\n",
    "        results[n] = {\n",
    "            \"upsampled_vae_filepath\": upsampled_vae_filepath,\n",
    "            \"mp3_filepath\": mp3_filepath,\n",
    "            \"metadata\": metadata_dict,\n",
    "        }\n",
    "    return results, semantic_codes_filepath\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "066b9a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from joblib import Parallel, delayed\n",
    "from tqdm import tqdm\n",
    "from suno_utils.audio import Audio\n",
    "\n",
    "metadata_keys_to_remove = [\n",
    "    \"average_spectrum_db\", \n",
    "    \"average_spectrum_db_first\", \n",
    "    \"average_spectrum_db_last\",\n",
    "    \"average_stereo_spectrum_mid\",\n",
    "    \"average_stereo_spectrum_side\",\n",
    "]\n",
    "\n",
    "def process_dirname(dirname):\n",
    "    results, semantic_codes_filepath = process_dir(base_dir, dirname)\n",
    "    for n in results.keys():\n",
    "        results[n][\"metadata\"] = {k: v for k, v in results[n][\"metadata\"].items() if k not in metadata_keys_to_remove}\n",
    "    \n",
    "    # Extend to handle multiple features\n",
    "    features = [\"ear_score\", \"shimmer_score\", \"hoot_cer\"]  # Add any feature strings you want here\n",
    "    # Collect feature values for each n, defaulting to 0.0 if missing or None\n",
    "    feature_matrix = []\n",
    "    indices = []\n",
    "    for n in results.keys():\n",
    "        meta = results[n][\"metadata\"]\n",
    "        values = []\n",
    "        for feature in features:\n",
    "            value = meta.get(feature, 0.0)\n",
    "            if value is None:\n",
    "                value = 0.0\n",
    "            values.append(value)\n",
    "        feature_matrix.append(values)\n",
    "        indices.append(n)\n",
    "    feature_matrix = np.array(feature_matrix)\n",
    "    if len(feature_matrix) >= 2:\n",
    "        max_sum_diff = -1\n",
    "        idx_pair = (None, None)\n",
    "        for i in range(len(feature_matrix)):\n",
    "            for j in range(i+1, len(feature_matrix)):\n",
    "                diff = np.abs(feature_matrix[i] - feature_matrix[j])\n",
    "                sum_diff = np.sum(diff)\n",
    "                if sum_diff > max_sum_diff:\n",
    "                    max_sum_diff = sum_diff\n",
    "                    idx_pair = (indices[i], indices[j])\n",
    "        # Print for debugging\n",
    "        #print(f\"{dirname}: Max summed feature diff {max_sum_diff} between indices {idx_pair} (feature values: {feature_matrix})\")\n",
    "    else:\n",
    "        print(f\"{dirname}: Not enough values to compare. (feature_matrix: {feature_matrix})\")\n",
    "        max_sum_diff = None\n",
    "        idx_pair = (None, None)\n",
    "    # For backward compatibility, keep the first feature as the main output\n",
    "    feature = \".\".join(features)\n",
    "    max_diff = max_sum_diff\n",
    "\n",
    "    # then get the mp3 files for them\n",
    "    if False:\n",
    "        mp3_filepaths = []\n",
    "        for n in idx_pair:\n",
    "            mp3_filepath = results[n][\"mp3_filepath\"]\n",
    "            # print the feature value\n",
    "            audio = Audio.from_file(mp3_filepath, n_channels=2).play()\n",
    "            print(f\"{feature}: {results[n]['metadata'][feature]}\")\n",
    "\n",
    "    # save the meta pairs\n",
    "    return {\n",
    "        \"dirname\": dirname,\n",
    "        \"idx_pair\": idx_pair,\n",
    "        \"feature\": feature,\n",
    "        \"max_diff\": max_diff,\n",
    "    }\n",
    "\n",
    "# Use joblib to parallelize the processing\n",
    "meta_pairs = Parallel(n_jobs=-1)(\n",
    "    delayed(process_dirname)(dirname) for dirname in tqdm(dirs)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39abff05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# store the meta pairs as a jsonl file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "871b4707",
   "metadata": {},
   "outputs": [],
   "source": [
    "meta_pairs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da9b2b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort the meta pairs by max_diff\n",
    "meta_pairs = [\n",
    "    meta_pair for meta_pair in meta_pairs\n",
    "    if meta_pair is not None and meta_pair.get(\"max_diff\") is not None\n",
    "]\n",
    "meta_pairs.sort(key=lambda x: x.get(\"max_diff\", 0), reverse=True)\n",
    "\n",
    "# get the mp3 files\n",
    "idx = 2002\n",
    "idx_pair = meta_pairs[idx][\"idx_pair\"]\n",
    "mp3_filepath_0 = os.path.join(base_dir, meta_pairs[idx][\"dirname\"], f\"{meta_pairs[idx]['dirname']}_{model_name}_{idx_pair[0]}.mp3\")\n",
    "mp3_filepath_1 = os.path.join(base_dir, meta_pairs[idx][\"dirname\"], f\"{meta_pairs[idx]['dirname']}_{model_name}_{idx_pair[1]}.mp3\")\n",
    "\n",
    "# also print the feature values\n",
    "print(f\"{meta_pairs[idx]['dirname']}: {meta_pairs[idx]['max_diff']}\")\n",
    "\n",
    "# play the mp3 files\n",
    "audio_0 = Audio.from_file(mp3_filepath_0, n_channels=2).play()\n",
    "audio_1 = Audio.from_file(mp3_filepath_1, n_channels=2).play()\n",
    "# wait for the audio to finish\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1dbf52b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(meta_pairs))\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "max_diffs = [meta_pair[\"max_diff\"] for meta_pair in meta_pairs]\n",
    "min_val = np.min(max_diffs)\n",
    "max_val = np.max(max_diffs)\n",
    "mean_val = np.mean(max_diffs)\n",
    "std_val = np.std(max_diffs)\n",
    "p5 = np.percentile(max_diffs, 5)\n",
    "p95 = np.percentile(max_diffs, 95)\n",
    "\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.hist(max_diffs, bins=100, color='skyblue', edgecolor='black')\n",
    "plt.title(\"Distribution of max_diff\")\n",
    "plt.xlabel(\"max_diff\")\n",
    "plt.ylabel(\"Count\")\n",
    "\n",
    "# Add red vertical lines for the 5% and 95% percentiles\n",
    "plt.axvline(p5, color='red', linestyle='--', linewidth=2, label='5%')\n",
    "plt.axvline(p95, color='red', linestyle='--', linewidth=2, label='95%')\n",
    "\n",
    "# Add text box with statistics\n",
    "stats_text = (\n",
    "    f\"min: {min_val:.3f}\\n\"\n",
    "    f\"max: {max_val:.3f}\\n\"\n",
    "    f\"mean: {mean_val:.3f}\\n\"\n",
    "    f\"std: {std_val:.3f}\\n\"\n",
    "    f\"5%: {p5:.3f}\\n\"\n",
    "    f\"95%: {p95:.3f}\"\n",
    ")\n",
    "plt.gca().text(\n",
    "    0.98, 0.98, stats_text,\n",
    "    transform=plt.gca().transAxes,\n",
    "    fontsize=10,\n",
    "    verticalalignment='top',\n",
    "    horizontalalignment='right',\n",
    "    bbox=dict(boxstyle=\"round,pad=0.5\", facecolor=\"white\", alpha=0.8)\n",
    ")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0f38da4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save out list of meta pair ids given the filter\n",
    "filtered_meta_pairs = [meta_pair for meta_pair in meta_pairs if meta_pair[\"max_diff\"] > 2.0]\n",
    "\n",
    "print(f\"Remaining {len(filtered_meta_pairs)} meta pairs from {len(meta_pairs)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "267b423f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save out the filtered meta pairs'\n",
    "os.makedirs(os.path.join(base_dir, \"metadata\"), exist_ok=True)\n",
    "output_filepath = os.path.join(base_dir, \"metadata\", \"filtered_meta_pairs.jsonl\")\n",
    "# write as jsonl\n",
    "from suno_utils.utils.text import write_jsonl\n",
    "write_jsonl(filtered_meta_pairs, output_filepath)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04b2a51",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}