{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "import glob\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import pyloudnorm as pyln\n",
    "\n",
    "from suno_utils.tasks.shimmerscore import shimmerscore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start with a main directory of audio files\n",
    "base_dir = \"/app2/suno/data/eval_outputs\"\n",
    "# audio_dir = os.path.join(base_dir, \"auk-clips-up-u-2\")\n",
    "audio_dir = os.path.join(base_dir, \"bluejay-lang-balanced-200\")\n",
    "\n",
    "# find all mp3 files in the directory\n",
    "filepaths = glob.glob(os.path.join(audio_dir, \"**\", \"*.mp3\"))\n",
    "print(len(filepaths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "csv_filepath = \"/app2/suno/data/diffusion_analysis_results.csv\"\n",
    "\n",
    "# load existing results\n",
    "# try:\n",
    "# df = pd.read_csv(csv_filepath)\n",
    "# now filter the filepaths to exclude the ones that are already in the dataframe\n",
    "# filepaths = [filepath for filepath in filepaths if filepath not in df['filepath'].values]\n",
    "# except Exception as e:#\n",
    "df = pd.DataFrame()\n",
    "\n",
    "print(len(filepaths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze_audio(filepath):\n",
    "    # check for metadata file\n",
    "    metadata = None\n",
    "    metadata_dict = None\n",
    "    metadata_filepath = filepath.replace(\".mp3\", \"__metadata.npz\")\n",
    "    if os.path.exists(metadata_filepath):\n",
    "        try:\n",
    "            # with open(metadata_filepath, \"r\") as f:\n",
    "            #    metadata = json.load(f)\n",
    "            metadata = np.load(metadata_filepath, allow_pickle=True)\n",
    "            # convert the metadata to a dictionary\n",
    "            metadata_dict = {}\n",
    "            for key, value in metadata.items():\n",
    "                metadata_dict[key] = value\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Error loading metadata: {e}\")\n",
    "            metadata_dict = None\n",
    "\n",
    "    return metadata_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "\n",
    "def process_file(filepath):\n",
    "    clip_id = os.path.basename(filepath).split(\"_\")[0]\n",
    "    model_name = \"_\".join(os.path.basename(filepath).split(\"_\")[1:]).replace(\".mp3\", \"\")\n",
    "\n",
    "    stats = analyze_audio(filepath)\n",
    "\n",
    "    if stats is None:\n",
    "        return None\n",
    "\n",
    "    return {\n",
    "        \"model_name\": model_name,\n",
    "        \"clip_id\": clip_id,\n",
    "        \"filepath\": filepath,\n",
    "        # \"lang\" : language,\n",
    "        **stats,\n",
    "    }\n",
    "\n",
    "\n",
    "print(f\"Processing {len(filepaths)} filepaths\")\n",
    "# Process files in parallel\n",
    "processed_results = Parallel(n_jobs=-1)(delayed(process_file)(filepath) for filepath in tqdm(filepaths))\n",
    "# Organize results by model\n",
    "for result in processed_results:\n",
    "    if result is not None:  # Handle None results\n",
    "        model_name = result.pop(\"model_name\")\n",
    "        if model_name not in results:\n",
    "            results[model_name] = []\n",
    "        results[model_name].append(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# store the results into a pandas dataframe\n",
    "# each row is a clip, columns are model, id, lufs_db, stereo_width, clipped_samples, etc...\n",
    "import pandas as pd\n",
    "import json\n",
    "import numpy as np\n",
    "\n",
    "# Create a list to hold all data\n",
    "all_data = []\n",
    "\n",
    "# Flatten the nested dictionary structure\n",
    "for model_name, clips in results.items():\n",
    "    for clip in clips:\n",
    "        # Create a base row with the model name\n",
    "        row = {\"model\": model_name}\n",
    "\n",
    "        # Dynamically add all keys from the clip dictionary\n",
    "        # This will work even if we change the structure of the results dict\n",
    "        for key, value in clip.items():\n",
    "            # Convert numpy arrays to lists for JSON serialization\n",
    "            if isinstance(value, np.ndarray):\n",
    "                row[key] = value.tolist()\n",
    "            else:\n",
    "                row[key] = value\n",
    "\n",
    "        all_data.append(row)\n",
    "\n",
    "# Create the DataFrame\n",
    "df = pd.DataFrame(all_data)\n",
    "\n",
    "# Display the first few rows and column information\n",
    "print(\"DataFrame shape:\", df.shape)\n",
    "print(\"Columns:\", df.columns.tolist())\n",
    "display(df.head())\n",
    "\n",
    "# Try to load the previous results\n",
    "try:\n",
    "    # Load the previous results\n",
    "    old_df = pd.read_json(csv_filepath, orient=\"records\")\n",
    "\n",
    "    # Merge the old and new results\n",
    "    df = pd.concat([old_df, df])\n",
    "\n",
    "    print(f\"Successfully merged with previous results. New shape: {df.shape}\")\n",
    "except Exception as e:\n",
    "    print(f\"Could not load previous results: {e}\")\n",
    "    print(\"Continuing with only new results.\")\n",
    "\n",
    "# Save the results using JSON to preserve array data types\n",
    "df.to_json(csv_filepath.replace(\".csv\", \".json\"), orient=\"records\")\n",
    "\n",
    "# Also save a CSV version for compatibility, but note that arrays will be converted to strings\n",
    "df.to_csv(csv_filepath, index=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list of unique model names\n",
    "unique_models = df[\"model\"].unique()\n",
    "print(unique_models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save dataframe to csv\n",
    "\n",
    "# lets remove any columns that are lists or arrays\n",
    "# remove these columns: average_spectrum_db\taverage_spectrum_db_first\taverage_spectrum_db_last\taverage_stereo_spectrum_mid\taverage_stereo_spectrum_side\n",
    "df_subset = df.drop(\n",
    "    columns=[\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",
    "\n",
    "# save the subset to a csv\n",
    "df_subset.to_csv(\n",
    "    \"/home/christian/code/neon/evaluations/diffusion_analysis_results_subset.csv\", index=False\n",
    ")\n",
    "\n",
    "# df.to_csv(\"/home/christian/code/neon/evaluations/diffusion_analysis_results.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plot_dir = os.path.join(base_dir, \"plots\")\n",
    "os.makedirs(plot_dir, exist_ok=True)\n",
    "\n",
    "# Specify the models to plot\n",
    "selected_models = [\n",
    "    # \"8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_cosine_nu1_0p1_residual_70k_step2\",\n",
    "    \"4n_25hz_2b_flow_5e5_sft_t8_500k_step10_cfg2\",\n",
    "    \"16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_220k_step2\",\n",
    "    # \"v3_flow_distill_s1039_lm_t1_0_6_1E6_beta100_n4_bt2_acc2_4k_last_step2\",\n",
    "    # \"v3_flow_distill_s1039_lm_t1_0_6_5E7_beta100_n4_bt2_acc2_4k_last_step2\",\n",
    "    \"v3_flow_distill_v1_t18_1E6_beta100_n4_bt2_acc2_4k_last_step2\",\n",
    "    # \"v3_flow_distill_s3177_lm_t1_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_2k_last_step2\",\n",
    "    # \"v3_flow_distill_s3177_lm_t1_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_3k_last_step2\",\n",
    "    # \"v3_flow_distill_s3177_rd2_lm_t2_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_3k_last_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_merge_t1_1E6_beta100_n8_bt2_acc4_3k_last_step2\",\n",
    "    # \"v3_flow_distill_s3177_carp_t1_v1_1E6_beta100_n4_bt2_acc4_3k_last_step2\",\n",
    "    # \"v3_flow_distill_s3177_rd2_hoot_t0_1E6_beta100_n8_bt2_acc4_3k_3k_step2\",\n",
    "    # \"v3_flow_distill_s3177_rd2_hoot_t0_1E6_beta100_n8_bt2_acc4_3k_6k_step2\",\n",
    "    # \"v3_flow_distill_s3177_rd2_hoot_t0_1E6_beta100_n8_bt2_acc4_3k_12k_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_hoot_t0_1E6_beta100_n8_bt2_acc4_3k_last_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_hoot_t1_1E6_beta100_n16_bt2_acc1_3k_last_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_hoot_t1_1E6_beta50_n16_bt2_acc1_3k_last_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_hoot_t1_1E6_beta25_n16_bt2_acc1_3k_last_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_hoot_t2_1E6_beta100_n16_bt2_acc1_3k_last_step2\",\n",
    "    \"v3_flow_distill_s3177_rd2_hoot_t2_1E6_beta100_n16_bt2_acc1_6k_last_step2\",\n",
    "    \"v2_infill_d4_t39_1E6_beta100_n16_bt2_acc4_3k_last_step10\",\n",
    "]  # <-- replace with the actual model names\n",
    "num_models = len(selected_models)\n",
    "\n",
    "colors = plt.cm.tab10.colors\n",
    "\n",
    "# Subset the dataframe\n",
    "df_subset = df[df[\"model\"].isin(selected_models)]\n",
    "height = 1.5\n",
    "# Define metrics to plot\n",
    "metrics = [\n",
    "    {\n",
    "        \"name\": \"lufs_db\",\n",
    "        \"title\": \"Loudness (LUFS)\",\n",
    "        \"filename\": \"lufs_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"stereo_width\",\n",
    "        \"title\": \"Stereo Width (0 = mono, 1 = wide)\",\n",
    "        \"filename\": \"stereo_width_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"shimmer_score\",\n",
    "        \"title\": \"Shimmer Score\",\n",
    "        \"filename\": \"shimmer_score_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"lufs_db_factor\",\n",
    "        \"title\": \"Loudness Factor\",\n",
    "        \"filename\": \"lufs_db_factor_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"clipped_samples\",\n",
    "        \"title\": \"Clipped Samples\",\n",
    "        \"filename\": \"clipped_samples_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"ear_score\",\n",
    "        \"title\": \"Ear Score\",\n",
    "        \"filename\": \"ear_score_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"hoot_cer\",\n",
    "        \"title\": \"Hoot CER\",\n",
    "        \"filename\": \"hoot_cer_distribution.png\",\n",
    "        \"figsize\": (6, height * num_models),\n",
    "    },\n",
    "]\n",
    "\n",
    "# Loop through each metric and create plots\n",
    "for metric in metrics:\n",
    "    # Calculate min and max for bins\n",
    "    metric_min = df_subset[metric[\"name\"]].min()\n",
    "    metric_max = df_subset[metric[\"name\"]].max()\n",
    "    metric_bins = np.linspace(metric_min, metric_max, 50)\n",
    "\n",
    "    # Create figure and axes\n",
    "    fig, axes = plt.subplots(len(selected_models), 1, figsize=metric[\"figsize\"], sharex=True)\n",
    "\n",
    "    # Plot histograms for each model\n",
    "    for i, model in enumerate(selected_models):\n",
    "        model_data = df_subset[df_subset[\"model\"] == model]\n",
    "        stats_string = (\n",
    "            f\"mean: {model_data[metric['name']].mean():.3f}, std: {model_data[metric['name']].std():.2f}\"\n",
    "        )\n",
    "\n",
    "        axes[i].hist(\n",
    "            model_data[metric[\"name\"]],\n",
    "            bins=metric_bins,\n",
    "            alpha=0.7,\n",
    "            edgecolor=\"black\",\n",
    "            color=colors[i % len(colors)],\n",
    "            label=stats_string,\n",
    "        )\n",
    "        axes[i].set_ylabel(\"Count\", fontsize=10)\n",
    "        axes[i].grid(True, linestyle=\"--\", alpha=0.5)\n",
    "        axes[i].set_title(model, fontsize=10)\n",
    "        axes[i].legend()\n",
    "\n",
    "    # Set labels and title\n",
    "    axes[-1].set_xlabel(metric[\"title\"], fontsize=12)\n",
    "    fig.suptitle(f\"{metric['title']} Distribution by Model\", fontsize=14)\n",
    "\n",
    "    # Adjust layout and save\n",
    "    plt.tight_layout()\n",
    "    plt.subplots_adjust(top=0.9)\n",
    "    plt.savefig(f\"{plot_dir}/{metric['filename']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# ----------- configure -----------\n",
    "baseline_model = \"v3_flow_distill_v1_t18_1E6_beta100_n4_bt2_acc2_4k_last_step2\"  # <-- pick your baseline\n",
    "\n",
    "# If baseline isn't already in selected_models, add it (optional)\n",
    "if baseline_model not in selected_models:\n",
    "    selected_models = [baseline_model] + selected_models\n",
    "\n",
    "# Rebuild metrics with a 'mode' per feature:\n",
    "#   mode='ratio'  -> plot percent change: 100 * (x / baseline - 1)\n",
    "#   mode='delta'  -> plot absolute difference: x - baseline\n",
    "metrics = [\n",
    "    {\n",
    "        \"name\": \"lufs_db\",\n",
    "        \"title\": \"Loudness (LUFS)\",\n",
    "        \"filename\": \"lufs_relative_box.png\",\n",
    "        \"mode\": \"delta\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"stereo_width\",\n",
    "        \"title\": \"Stereo Width (0 mono, 1 wide)\",\n",
    "        \"filename\": \"stereo_width_relative_box.png\",\n",
    "        \"mode\": \"ratio\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"shimmer_score\",\n",
    "        \"title\": \"Shimmer Score\",\n",
    "        \"filename\": \"shimmer_score_relative_box.png\",\n",
    "        \"mode\": \"ratio\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"lufs_db_factor\",\n",
    "        \"title\": \"Loudness Factor\",\n",
    "        \"filename\": \"lufs_factor_relative_box.png\",\n",
    "        \"mode\": \"ratio\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"clipped_samples\",\n",
    "        \"title\": \"Clipped Samples\",\n",
    "        \"filename\": \"clipped_samples_relative_box.png\",\n",
    "        \"mode\": \"delta\",\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"ear_score\",\n",
    "        \"title\": \"Ear Score\",\n",
    "        \"filename\": \"ear_score_relative_box.png\",\n",
    "        \"mode\": \"ratio\",\n",
    "    },\n",
    "    {\"name\": \"hoot_cer\", \"title\": \"Hoot CER\", \"filename\": \"hoot_cer_relative_box.png\", \"mode\": \"ratio\"},\n",
    "]\n",
    "\n",
    "plot_dir = os.path.join(base_dir, \"plots\")\n",
    "os.makedirs(plot_dir, exist_ok=True)\n",
    "\n",
    "# Colors for boxes\n",
    "colors = plt.cm.tab10(np.arange(len(selected_models)) % 10)\n",
    "\n",
    "# Subset once\n",
    "df_subset = df[df[\"model\"].isin(selected_models)].copy()\n",
    "\n",
    "\n",
    "def relative_values(series, baseline_value, mode):\n",
    "    series = series.dropna()\n",
    "    if mode == \"ratio\":\n",
    "        if baseline_value == 0 or np.isclose(baseline_value, 0):\n",
    "            # Percent is undefined vs 0; fall back to absolute delta\n",
    "            return series - baseline_value\n",
    "        return 100.0 * (series / baseline_value - 1.0)\n",
    "    else:  # 'delta'\n",
    "        return series - baseline_value\n",
    "\n",
    "\n",
    "for metric in metrics:\n",
    "    feat = metric[\"name\"]\n",
    "    mode = metric[\"mode\"]\n",
    "\n",
    "    # Get baseline statistic (median is robust; switch to .mean() if you prefer)\n",
    "    base_vals = df_subset.loc[df_subset[\"model\"] == baseline_model, feat].dropna()\n",
    "    if base_vals.empty:\n",
    "        print(f\"[WARN] No baseline data for {feat} in model '{baseline_model}'. Skipping.\")\n",
    "        continue\n",
    "    baseline_stat = base_vals.median()\n",
    "\n",
    "    # Build per-model relative arrays\n",
    "    data = []\n",
    "    labels = []\n",
    "    for i, model in enumerate(selected_models):\n",
    "        model_vals = df_subset.loc[df_subset[\"model\"] == model, feat].dropna()\n",
    "        if model_vals.empty:\n",
    "            continue\n",
    "        rel = relative_values(model_vals, baseline_stat, mode)\n",
    "        data.append(rel.values)\n",
    "        labels.append(model)\n",
    "\n",
    "    if not data:\n",
    "        print(f\"[WARN] No data to plot for {feat}.\")\n",
    "        continue\n",
    "\n",
    "    # Figure size scales with number of models\n",
    "    fig_height = max(2.0, 0.6 * len(labels) + 1.5)\n",
    "    fig, ax = plt.subplots(figsize=(16, fig_height))\n",
    "\n",
    "    bplot = ax.boxplot(\n",
    "        data,\n",
    "        vert=False,\n",
    "        labels=labels,\n",
    "        patch_artist=True,\n",
    "        whis=(5, 95),  # whiskers at 5–95th percentile (adjust if you like Tukey 1.5 IQR)\n",
    "        showfliers=True,\n",
    "    )\n",
    "\n",
    "    # Color the boxes\n",
    "    for i, box in enumerate(bplot[\"boxes\"]):\n",
    "        box.set_facecolor(colors[i])\n",
    "        box.set_alpha(0.7)\n",
    "        box.set_edgecolor(\"black\")\n",
    "\n",
    "    # Median labels (small annotation on each box)\n",
    "    for i, med_line in enumerate(bplot[\"medians\"]):\n",
    "        x_med = med_line.get_xdata().mean()  # horizontal\n",
    "        y_med = med_line.get_ydata().mean()\n",
    "        if mode == \"ratio\":\n",
    "            ax.text(x_med, i + 1, f\"{x_med:+.1f}%\", va=\"center\", ha=\"left\", fontsize=8)\n",
    "        else:\n",
    "            ax.text(x_med, i + 1, f\"{x_med:+.3g}\", va=\"center\", ha=\"left\", fontsize=8)\n",
    "\n",
    "    # Reference line at 0 (equal to baseline)\n",
    "    ax.axvline(0, linestyle=\"--\", linewidth=1, color=\"gray\", alpha=0.7)\n",
    "\n",
    "    # Labels/titles\n",
    "    if mode == \"ratio\":\n",
    "        ax.set_xlabel(\n",
    "            f\"{metric['title']} — % change vs {baseline_model} (median baseline = {baseline_stat:.3g})\"\n",
    "        )\n",
    "    else:\n",
    "        ax.set_xlabel(\n",
    "            f\"{metric['title']} — Δ vs {baseline_model} (median baseline = {baseline_stat:.3g})\"\n",
    "        )\n",
    "\n",
    "    ax.set_title(f\"{metric['title']} by Model (relative to {baseline_model})\", fontsize=12)\n",
    "    ax.grid(True, axis=\"x\", linestyle=\":\", alpha=0.5)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    out_path = os.path.join(plot_dir, metric[\"filename\"])\n",
    "    plt.show()\n",
    "    # plt.savefig(out_path, dpi=150)\n",
    "    # plt.close(fig)\n",
    "\n",
    "print(\"Saved relative boxplots to:\", plot_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "sr = 48000  # replace with your actual sample rate\n",
    "n_fft = 16384  # replace with your actual FFT size\n",
    "df_subset = df[df[\"model\"].isin(selected_models)]\n",
    "colors = plt.cm.tab10.colors\n",
    "normalize_at_1khz = True  # Flag to normalize spectra at 1kHz\n",
    "\n",
    "# ---- Frequency Axis ----\n",
    "first_spec_mid = df[\"average_stereo_spectrum_mid\"].iloc[0]\n",
    "n_bins = len(first_spec_mid)\n",
    "freqs = np.linspace(0, sr / 2, n_bins)\n",
    "\"\"\n",
    "# Create a figure with two subplots (mid and side)\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(7, 6), sharex=True)\n",
    "\n",
    "# Find the index closest to 1kHz for normalization\n",
    "if normalize_at_1khz:\n",
    "    idx_1khz = np.argmin(np.abs(freqs - 1000))\n",
    "\n",
    "# Plot Mid Spectrum\n",
    "for i, model in enumerate(selected_models):\n",
    "    model_data = df_subset[df_subset[\"model\"] == model]\n",
    "\n",
    "    # Average over all clips for mid channel\n",
    "    mid_spectra = np.array([s for s in model_data[\"average_stereo_spectrum_mid\"]])\n",
    "    mean_mid_spectrum = np.mean(mid_spectra, axis=0)\n",
    "\n",
    "    # Normalize at 1kHz if flag is set\n",
    "    if normalize_at_1khz:\n",
    "        normalization_value = mean_mid_spectrum[idx_1khz]\n",
    "        mean_mid_spectrum = mean_mid_spectrum - normalization_value\n",
    "\n",
    "    ax1.plot(freqs, mean_mid_spectrum, label=model, color=colors[i % len(colors)])\n",
    "\n",
    "# Plot Side Spectrum\n",
    "for i, model in enumerate(selected_models):\n",
    "    model_data = df_subset[df_subset[\"model\"] == model]\n",
    "\n",
    "    # Average over all clips for side channel\n",
    "    side_spectra = np.array([s for s in model_data[\"average_stereo_spectrum_side\"]])\n",
    "    mean_side_spectrum = np.mean(side_spectra, axis=0)\n",
    "\n",
    "    # Normalize at 1kHz if flag is set\n",
    "    if normalize_at_1khz:\n",
    "        normalization_value = mean_side_spectrum[idx_1khz]\n",
    "        mean_side_spectrum = mean_side_spectrum - normalization_value\n",
    "\n",
    "    ax2.plot(freqs, mean_side_spectrum, label=model, color=colors[i % len(colors)])\n",
    "\n",
    "# Configure Mid plot\n",
    "ax1.set_xscale(\"log\")\n",
    "ax1.set_ylabel(\"Mid Channel (dB)\", fontsize=12)\n",
    "ax1.set_title(\n",
    "    \"Average Mid Spectrum per Model\" + (\" (normalized at 1kHz)\" if normalize_at_1khz else \"\"),\n",
    "    fontsize=14,\n",
    ")\n",
    "ax1.grid(True, which=\"both\", linestyle=\"--\", alpha=0.5)\n",
    "ax1.set_ylim(-30, 30) if normalize_at_1khz else ax1.set_ylim(-40, 48)\n",
    "ax1.legend()\n",
    "\n",
    "# Configure Side plot\n",
    "ax2.set_xscale(\"log\")\n",
    "ax2.set_xlabel(\"Frequency (Hz)\", fontsize=12)\n",
    "ax2.set_ylabel(\"Side Channel (dB)\", fontsize=12)\n",
    "ax2.set_title(\n",
    "    \"Average Side Spectrum per Model\" + (\" (normalized at 1kHz)\" if normalize_at_1khz else \"\"),\n",
    "    fontsize=14,\n",
    ")\n",
    "ax2.grid(True, which=\"both\", linestyle=\"--\", alpha=0.5)\n",
    "ax2.set_ylim(-30, 20) if normalize_at_1khz else ax2.set_ylim(-40, 48)\n",
    "ax2.set_xlim(20, 24000)\n",
    "# ax2.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "# plt.show()\n",
    "plt.savefig(f\"{plot_dir}/average_mid_side_spectrum_per_model.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sr = 48000  # replace with your actual sample rate\n",
    "n_fft = 16384  # replace with your actual FFT size\n",
    "df_subset = df[df[\"model\"].isin(selected_models)]\n",
    "colors = plt.cm.tab10.colors\n",
    "\n",
    "# ---- Frequency Axis ----\n",
    "first_spec = df[\"average_spectrum_db_last\"].iloc[0]\n",
    "n_bins = np.array(first_spec).shape[1]\n",
    "freqs = np.linspace(0, sr / 2, n_bins)\n",
    "\n",
    "# Create a single plot for all models\n",
    "plt.figure(figsize=(8, 4))\n",
    "\n",
    "for i, model in enumerate(selected_models):\n",
    "    model_data = df_subset[df_subset[\"model\"] == model]\n",
    "\n",
    "    # Process last 30s spectra\n",
    "    spectra_last = np.array([np.mean(s, axis=0) for s in model_data[\"average_spectrum_db_last\"]])\n",
    "    mean_spectrum_last = spectra_last.mean(axis=0)\n",
    "\n",
    "    # Process first 30s spectra\n",
    "    spectra_first = np.array([np.mean(s, axis=0) for s in model_data[\"average_spectrum_db_first\"]])\n",
    "    mean_spectrum_first = spectra_first.mean(axis=0)\n",
    "\n",
    "    # Calculate delta between last and first\n",
    "    spectra_delta = mean_spectrum_last - mean_spectrum_first\n",
    "\n",
    "    # Plot only the delta for each model\n",
    "    plt.plot(freqs, spectra_delta, label=f\"{model}\", color=colors[i % len(colors)])\n",
    "\n",
    "plt.xscale(\"log\")\n",
    "plt.ylabel(\"Delta Magnitude (dB)\", fontsize=12)\n",
    "plt.xlabel(\"Frequency (Hz)\", fontsize=12)\n",
    "plt.title(\"Spectral Change (Last 30s - First 30s)\", fontsize=14)\n",
    "plt.grid(True, which=\"both\", linestyle=\"--\", alpha=0.5)\n",
    "plt.ylim(-5, 24)  # Adjusted for delta values\n",
    "plt.xlim(1000, 20000)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.savefig(f\"{plot_dir}/spectral_change_last_30s_first_30s.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "sr = 48000  # replace with your actual sample rate\n",
    "n_fft = 16384  # replace with your actual FFT size\n",
    "df_subset = df[df[\"model\"].isin(selected_models)]\n",
    "colors = plt.cm.tab10.colors\n",
    "\n",
    "# ---- Frequency Axis ----\n",
    "first_spec = df[\"average_spectrum_db_last\"].iloc[0]\n",
    "n_bins = np.array(first_spec).shape[1]\n",
    "freqs = np.linspace(0, sr / 2, n_bins)\n",
    "\n",
    "# Create a subplot for each model\n",
    "fig, axes = plt.subplots(len(selected_models), 1, figsize=(6, (2.5 * len(selected_models))), sharex=True)\n",
    "\n",
    "for i, model in enumerate(selected_models):\n",
    "    model_data = df_subset[df_subset[\"model\"] == model]\n",
    "    ax = axes[i] if len(selected_models) > 1 else axes\n",
    "\n",
    "    # Process last 30s spectra\n",
    "    spectra_last = np.array([np.mean(s, axis=0) for s in model_data[\"average_spectrum_db_last\"]])\n",
    "    mean_spectrum_last = spectra_last.mean(axis=0)\n",
    "\n",
    "    # Process first 30s spectra\n",
    "    spectra_first = np.array([np.mean(s, axis=0) for s in model_data[\"average_spectrum_db_first\"]])\n",
    "    mean_spectrum_first = spectra_first.mean(axis=0)\n",
    "\n",
    "    spectra_delta = mean_spectrum_last - mean_spectrum_first\n",
    "\n",
    "    # Plot both spectra on the same subplot\n",
    "    ax.plot(freqs, mean_spectrum_last, label=\"Last 30s\", linestyle=\"--\", color=colors[0])\n",
    "    ax.plot(freqs, mean_spectrum_first, label=\"First 30s\", color=colors[1])\n",
    "    ax.plot(freqs, spectra_delta, label=\"Delta\", color=colors[2])\n",
    "\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_ylabel(\"Average Magnitude (dB)\", fontsize=12)\n",
    "    ax.set_title(f\"Model: {model}\", fontsize=14)\n",
    "    ax.grid(True, which=\"both\", linestyle=\"--\", alpha=0.5)\n",
    "    ax.set_ylim(-30, 30)\n",
    "    ax.set_xlim(20, 20000)\n",
    "    ax.legend()\n",
    "\n",
    "# Set common x-label\n",
    "plt.xlabel(\"Frequency (Hz)\", fontsize=12)\n",
    "plt.tight_layout()\n",
    "plt.savefig(f\"{plot_dir}/spectral_change_last_30s_first_30s.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}