{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62e79754",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from suno_utils.utils.text import read_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae842779",
   "metadata": {},
   "outputs": [],
   "source": [
    "WITH_PARENT_SCORES = {\n",
    "    # \"Auk Raw\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-orig__no_DPO_2025_04_24-17_30_44.npz\",\n",
    "    # \"Auk 4.8\": \"scores/feat_eval_subgenre_with_parent_chirp-auk__2025_04_16-03_11_29.npz\",\n",
    "    # \"Auk 4.20\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-t1__420_DPO_2025_04_24-16_37_26.npz\",\n",
    "    # \"Auk 5.5\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__5_05_dpo_2025_05_05-18_36_15.npz\",\n",
    "    # \"BlueJay Raw\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__5_15_raw_2025_05_15-19_57_43.npz\",\n",
    "    \"Auk Prod\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__426_2025_04_29-19_16_46.npz\",\n",
    "    \"BlueJay SFT\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__blue_sft_2025_06_09-15_06_47.npz\",\n",
    "    \"BlueJay t1-r7\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__blue_t1r7_2025_06_09-13_56_33.npz\",\n",
    "    \"BlueJay 06-26-35\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__actual_ckpt_06-26-35_2025_06_19-03_40_17.npz\",\n",
    "    # \"Blue CFG\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__actual_ckpt_06-26-35_cfg075_no_neg_no_rep_2025_06_20-14_06_41.npz\",\n",
    "    # \"Blue CFG No Neg\": \"scores/feat_eval_subgenre_with_parent_chirp-auk-eval__actual_ckpt_06-26-35_cfg075_no_rep_2025_06_20-15_31_12.npz\",\n",
    "    \"v4\": \"scores/feat_eval_subgenre_with_parent_chirp-v4-h-s-33__2025_04_16-03_38_02.npz\",\n",
    "}\n",
    "\n",
    "WITHIN_PARENT_SCORES = {\n",
    "    # \"Auk Raw\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-orig__no_DPO_2025_04_24-17_30_44.npz\",\n",
    "    # \"Auk 4.8\": \"scores/feat_eval_subgenre_within_parent_chirp-auk__2025_04_16-03_11_29.npz\",\n",
    "    # \"Auk 4.20\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-t1__420_DPO_2025_04_24-16_37_26.npz\",\n",
    "    # \"Auk 5.5\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__5_05_dpo_2025_05_05-18_36_15.npz\",\n",
    "    # \"BlueJay Raw\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__5_15_raw_2025_05_15-19_57_43.npz\",\n",
    "    \"Auk Prod\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__426_2025_04_29-19_16_46.npz\",\n",
    "    \"BlueJay SFT\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__blue_sft_2025_06_09-15_06_47.npz\",\n",
    "    \"BlueJay t1-r7\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__blue_t1r7_2025_06_09-13_56_33.npz\",\n",
    "    \"BlueJay 06-26-35\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__actual_ckpt_06-26-35_2025_06_19-03_40_17.npz\",\n",
    "    # \"Blue CFG\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__actual_ckpt_06-26-35_cfg075_no_neg_no_rep_2025_06_20-14_06_41.npz\",\n",
    "    # \"Blue CFG No Neg\": \"scores/feat_eval_subgenre_within_parent_chirp-auk-eval__actual_ckpt_06-26-35_cfg075_no_rep_2025_06_20-15_31_12.npz\",\n",
    "    \"v4\": \"scores/feat_eval_subgenre_within_parent_chirp-v4-h-s-33__2025_04_16-03_38_02.npz\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d8ae77",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_average_scores_by_genre(\n",
    "    data, figsize=(12, 7), title=\"Inline Tag Avg F1 Scores\", y_label=\"Average F1\", to_skip=[]\n",
    "):\n",
    "    # First, calculate the average scores for each name and genre\n",
    "    names = list(data.keys())\n",
    "    genres = list(data[names[0]].keys())  # Assuming all names have the same genres\n",
    "    genres = [g for g in genres if g not in to_skip]\n",
    "\n",
    "    # Dictionary to store average scores\n",
    "    averages = {}\n",
    "    for name in names:\n",
    "        averages[name] = {}\n",
    "        for genre in genres:\n",
    "            averages[name][genre] = np.mean(data[name][genre])\n",
    "\n",
    "    # Prepare the plot\n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "\n",
    "    # Set width of bars\n",
    "    bar_width = 0.8 / len(names)\n",
    "\n",
    "    # Set position of bars on x axis\n",
    "    positions = np.arange(len(genres))\n",
    "\n",
    "    # Create bars for each name\n",
    "    for i, name in enumerate(names):\n",
    "        avg_scores = [averages[name][genre] for genre in genres]\n",
    "        offset = i - (len(names) - 1) / 2\n",
    "        x_pos = positions + offset * bar_width\n",
    "        ax.bar(x_pos, avg_scores, width=bar_width, label=name)\n",
    "\n",
    "    # Add labels and title\n",
    "    ax.set_xlabel(\"Genre\", fontsize=12)\n",
    "    ax.set_ylabel(y_label, fontsize=12)\n",
    "    ax.set_title(title, fontsize=14)\n",
    "    ax.set_xticks(positions)\n",
    "    ax.set_xticklabels([genres[:16] for genres in genres], rotation=90)\n",
    "\n",
    "    # Add a legend\n",
    "    ax.legend()\n",
    "\n",
    "    # Add grid for better readability\n",
    "    ax.grid(axis=\"y\", linestyle=\"--\", alpha=0.3)\n",
    "\n",
    "    # Improve aesthetics\n",
    "    plt.tight_layout()\n",
    "\n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51360064",
   "metadata": {},
   "outputs": [],
   "source": [
    "subgenres = read_json(\"rate_your_music_genres.json\")\n",
    "parent_lookup = {}\n",
    "for data in subgenres:\n",
    "    name = data[\"name\"]\n",
    "    parent = data[\"parent_genre\"]\n",
    "    parent_lookup[name] = parent\n",
    "    if parent == \"Folk\":\n",
    "        print(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54fc904f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_with_parent(subgenres_path):\n",
    "    scores = dict(np.load(subgenres_path))\n",
    "\n",
    "    parent_scores = {}\n",
    "    for genre, sub_scores in scores.items():\n",
    "        parent = parent_lookup[genre]\n",
    "        if parent not in parent_scores:\n",
    "            parent_scores[parent] = []\n",
    "        for score in sub_scores:\n",
    "            parent_scores[parent].append(score)\n",
    "\n",
    "    return parent_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dabb7201",
   "metadata": {},
   "outputs": [],
   "source": [
    "with_parent_results = {}\n",
    "for k, v in WITH_PARENT_SCORES.items():\n",
    "    with_parent_results[k] = process_with_parent(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a7eb981",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_average_scores_by_genre(\n",
    "    with_parent_results, title=\"SubGenre <-> Parent Similarity\", y_label=\"Ditto Self Sim\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054b647d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_within_parent(subgenres_path):\n",
    "    scores = dict(np.load(subgenres_path, allow_pickle=True))\n",
    "\n",
    "    scores_within = {}\n",
    "    scores_between = {}\n",
    "    for genre in scores:\n",
    "        genre_data = scores[genre].item()\n",
    "        within = genre_data[\"within_genre_pairs\"]\n",
    "        between = genre_data[\"between_genre_pairs\"]\n",
    "        scores_within[genre] = within\n",
    "        scores_between[genre] = between\n",
    "    return scores_within, scores_between"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f52d79e",
   "metadata": {},
   "outputs": [],
   "source": [
    "within_results = {}\n",
    "between_results = {}\n",
    "diff = {}\n",
    "for k, v in WITHIN_PARENT_SCORES.items():\n",
    "    within, between = process_within_parent(v)\n",
    "    within_results[f\"{k}_within\"] = within\n",
    "    within_results[f\"{k}_between\"] = between\n",
    "\n",
    "    diff[k] = {}\n",
    "    for genre in within:\n",
    "        diff[k][genre] = np.median(within[genre]) - np.median(between[genre])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d536d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_skip = [\"Folk\", \"Industrial & Noise\"]\n",
    "fig, ax = plot_average_scores_by_genre(\n",
    "    diff, title=\"Between vs Within Subgenre Diversity\", y_label=\"Ditto Self Sim Delta\", to_skip=to_skip\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "189574c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_model_performance(data, figsize=(13, 8), title=\"SubGenre Similarity\"):\n",
    "    genres = set()\n",
    "    for model_data in data.values():\n",
    "        for genre in model_data.keys():\n",
    "            genres.add(genre)\n",
    "    genres = sorted(list(genres))\n",
    "    models = sorted(list(data.keys()))\n",
    "\n",
    "    # Set up the figure and axis\n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "\n",
    "    # Calculate the positions for each model-genre combination\n",
    "    x_positions = {}\n",
    "    model_width = 0.8\n",
    "    group_padding = 1\n",
    "    current_pos = 1\n",
    "\n",
    "    # Store x positions for each model within each genre group\n",
    "    for i, genre in enumerate(genres):\n",
    "        x_positions[genre] = {}\n",
    "        for j, model in enumerate(models):\n",
    "            x_positions[genre][model] = current_pos + j * model_width\n",
    "        current_pos += len(models) * model_width + group_padding\n",
    "\n",
    "    # Colors for different models\n",
    "    colors = plt.cm.tab10(np.linspace(0, 1, len(models)))\n",
    "    model_colors = {model: colors[i] for i, model in enumerate(models)}\n",
    "\n",
    "    # Collect all standard deviations to normalize circle sizes\n",
    "    all_stds = []\n",
    "    for model in models:\n",
    "        for genre in genres:\n",
    "            if genre in data[model]:\n",
    "                scores = data[model][genre]\n",
    "                if len(scores) > 1:  # Need at least 2 points for std\n",
    "                    all_stds.append(np.std(scores, ddof=1))\n",
    "\n",
    "    # If we have stds, find min and max for scaling\n",
    "    min_std = min(all_stds) if all_stds else 1\n",
    "    max_std = max(all_stds) if all_stds else 2\n",
    "\n",
    "    # Plot the mean as circles for each genre-model combination\n",
    "    for genre in genres:\n",
    "        for model in models:\n",
    "            if genre in data[model]:\n",
    "                scores = data[model][genre]\n",
    "                x_pos = x_positions[genre][model]\n",
    "\n",
    "                # Calculate mean and std\n",
    "                mean_value = np.mean(scores)\n",
    "\n",
    "                # Calculate std (sample standard deviation with ddof=1)\n",
    "                if len(scores) > 1:\n",
    "                    std_value = np.std(scores, ddof=1)\n",
    "                else:\n",
    "                    std_value = 0  # No std for single point\n",
    "\n",
    "                # Scale circle size inversely by std deviation\n",
    "                # Smaller std = larger circle (more confident in the mean)\n",
    "                if std_value > 0:\n",
    "                    # Normalize to range between min_size and max_size\n",
    "                    min_size = 30\n",
    "                    max_size = 150\n",
    "\n",
    "                    # Invert the scale - smaller std should be larger circle\n",
    "                    normalized_std = 1 - (\n",
    "                        (std_value - min_std) / (max_std - min_std) if max_std != min_std else 0\n",
    "                    )\n",
    "                    circle_size = min_size + normalized_std * (max_size - min_size)\n",
    "                else:\n",
    "                    circle_size = 75  # Default for single point or zero std\n",
    "\n",
    "                # Plot mean as a circle\n",
    "                ax.scatter(\n",
    "                    x_pos,\n",
    "                    mean_value,\n",
    "                    color=model_colors[model],\n",
    "                    s=circle_size,\n",
    "                    zorder=3,\n",
    "                    alpha=0.7,\n",
    "                    label=f\"{model}\" if genre == genres[0] else \"\",\n",
    "                )\n",
    "\n",
    "                # Add text label with the mean value\n",
    "                # ax.text(x_pos, mean_value, f\"{mean_value:.1f}\",\n",
    "                #        ha='center', va='center', fontsize=9, fontweight='bold')\n",
    "\n",
    "    # Set x-axis ticks and labels\n",
    "    tick_positions = []\n",
    "    tick_labels = []\n",
    "    for genre in genres:\n",
    "        # Calculate the center position for each genre group\n",
    "        genre_positions = list(x_positions[genre].values())\n",
    "        center = sum(genre_positions) / len(genre_positions)\n",
    "        tick_positions.append(center)\n",
    "        tick_labels.append(genre)\n",
    "    ax.set_xticks(tick_positions)\n",
    "    ax.set_xticklabels(tick_labels, fontsize=11, rotation=90)  # Rotate labels 90 degrees\n",
    "\n",
    "    # Add vertical lines to separate genre groups\n",
    "    for genre_idx in range(len(genres) - 1):\n",
    "        last_pos = max(list(x_positions[genres[genre_idx]].values()))\n",
    "        next_pos = min(list(x_positions[genres[genre_idx + 1]].values()))\n",
    "        separator_pos = (last_pos + next_pos) / 2\n",
    "        ax.axvline(x=separator_pos, color=\"gray\", linestyle=\"--\", alpha=0.5)\n",
    "\n",
    "    # Add a legend for the models\n",
    "    handles = [\n",
    "        plt.Line2D([0], [0], marker=\"o\", color=\"w\", markerfacecolor=color, markersize=8, label=model)\n",
    "        for model, color in model_colors.items()\n",
    "    ]\n",
    "    ax.legend(handles=handles, title=\"Models\", loc=\"upper right\")\n",
    "\n",
    "    # Set labels and title\n",
    "    ax.set_ylabel(\"Mean Ditto Cosine Sim\", fontsize=14)\n",
    "    ax.set_title(title, fontsize=16)\n",
    "\n",
    "    # Remove top and right spines\n",
    "    ax.spines[\"top\"].set_visible(False)\n",
    "    ax.spines[\"right\"].set_visible(False)\n",
    "\n",
    "    # Add grid for better readability\n",
    "    ax.yaxis.grid(True, linestyle=\"--\", alpha=0.7)\n",
    "\n",
    "    # Adjust y-axis to give some padding\n",
    "    y_min, y_max = ax.get_ylim()\n",
    "    y_padding = (y_max - y_min) * 0.1\n",
    "    ax.set_ylim(y_min - y_padding, y_max + y_padding)\n",
    "\n",
    "    # Add padding at the bottom for rotated labels\n",
    "    plt.subplots_adjust(bottom=0.15)\n",
    "\n",
    "    # Adjust layout\n",
    "    plt.tight_layout()\n",
    "\n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21fc0f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_model_performance(within_results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}