{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "pattern = r\"^([^_]+)_format_trimmed_([^_]+)$\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labelbox_path = \"outputs/Export  project - stems-trimmed-preference - 6_18_2025.ndjson\"\n",
    "labelbox_data = read_jsonl(labelbox_path)\n",
    "\n",
    "metadata_filepath = \"outputs/metadata-more-stems-trimmed-20250617.json\"\n",
    "metadata = json.load(open(metadata_filepath, \"r\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_names = [\"original\", \"suno\", \"lalal\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seen_per = {name: 0 for name in model_names}\n",
    "wins_per = {name: 0 for name in model_names}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for row in labelbox_data:\n",
    "    row_id = row[\"data_row\"][\"global_key\"]\n",
    "\n",
    "    row_metadata = metadata[row_id]\n",
    "    source_a = re.match(pattern, row_metadata[\"source_a\"]).group(1)\n",
    "    source_b = re.match(pattern, row_metadata[\"source_b\"]).group(1)\n",
    "    song_name = row_metadata[\"song_name\"]\n",
    "    instrument = row_metadata[\"instrument\"]\n",
    "\n",
    "    projects = list(row[\"projects\"].keys())\n",
    "    assert len(projects) == 1\n",
    "    project_id = projects[0]\n",
    "    project = row[\"projects\"][project_id]\n",
    "    labels = project[\"labels\"]\n",
    "    for label in labels:  # represents each rating of the row, should be consensus count\n",
    "        classifications = label[\"annotations\"][\"classifications\"]\n",
    "        if len(classifications) != 1:\n",
    "            print(classifications)\n",
    "        for c in classifications:\n",
    "            question = c[\"name\"]\n",
    "            answer = c[\"radio_answer\"][\"value\"]\n",
    "\n",
    "            won = None\n",
    "            if answer == \"A\":\n",
    "                won = source_a\n",
    "            elif answer == \"B\":\n",
    "                won = source_b\n",
    "\n",
    "            results.append(\n",
    "                {\n",
    "                    \"row_id\": row_id,\n",
    "                    \"won\": won,\n",
    "                    \"source_a\": source_a,\n",
    "                    \"source_b\": source_b,\n",
    "                    \"song_name\": song_name,\n",
    "                    \"instrument\": instrument,\n",
    "                    \"question\": question,\n",
    "                    \"answer\": answer,\n",
    "                }\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(results)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overall Preference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def score_column(column, df):\n",
    "    preference_counts = df[column].value_counts()\n",
    "    preference_percentages = df[column].value_counts(normalize=True) * 100\n",
    "\n",
    "    print(\"\\n\" + column)\n",
    "    for key in preference_counts.keys():\n",
    "        print(f\"{key}: {preference_counts[key]} votes {preference_percentages[key]}%\")\n",
    "\n",
    "    return preference_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preference_counts = score_column(\"won\", df)\n",
    "a_counts = score_column(\"source_a\", df)\n",
    "b_counts = score_column(\"source_b\", df)\n",
    "total_counts = a_counts + b_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "win_rates = preference_counts / total_counts * 100\n",
    "print(win_rates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "results = win_rates\n",
    "\n",
    "# Create the plot with custom colors\n",
    "plt.figure(figsize=(10, 6))\n",
    "colors = plt.cm.viridis(np.linspace(0, 1, len(results)))\n",
    "bars = plt.bar(results.index, results.values, color=colors)\n",
    "\n",
    "plt.title(\"Instrumental Stem Sep Overall Win Rates\")\n",
    "plt.xlabel(\"Model\")\n",
    "plt.ylabel(\"Win Rate\")\n",
    "plt.xticks(rotation=45, ha=\"right\")\n",
    "\n",
    "# Add value labels\n",
    "for bar, value in zip(bars, results.values):\n",
    "    plt.text(\n",
    "        bar.get_x() + bar.get_width() / 2.0,\n",
    "        bar.get_height() + 0.01,\n",
    "        f\"{value:.3f}\",\n",
    "        ha=\"center\",\n",
    "        va=\"bottom\",\n",
    "        fontweight=\"bold\",\n",
    "    )\n",
    "\n",
    "# Add padding at the top\n",
    "plt.ylim(0, max(results.values) * 1.15)  # 15% padding above the highest bar\n",
    "\n",
    "plt.grid(axis=\"y\", alpha=0.3)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "# Assuming your dataframe is called 'df'\n",
    "# df = your_dataframe_here\n",
    "\n",
    "\n",
    "def calculate_win_rates(df):\n",
    "    \"\"\"\n",
    "    Calculate win rates between models using the 'won' column.\n",
    "    \"\"\"\n",
    "\n",
    "    # Get all unique models\n",
    "    all_models = set(df[\"source_a\"].unique()) | set(df[\"source_b\"].unique())\n",
    "    all_models = sorted(list(all_models))\n",
    "\n",
    "    # Create win rate matrix\n",
    "    win_rate_matrix = pd.DataFrame(index=all_models, columns=all_models, dtype=float)\n",
    "\n",
    "    # Fill diagonal with NaN (can't compete against itself)\n",
    "    for model in all_models:\n",
    "        win_rate_matrix.loc[model, model] = np.nan\n",
    "\n",
    "    # Calculate win rates for each pair\n",
    "    for model_a in all_models:\n",
    "        for model_b in all_models:\n",
    "            if model_a != model_b:\n",
    "                # Get all matchups between these two models\n",
    "                matchups = df[\n",
    "                    ((df[\"source_a\"] == model_a) & (df[\"source_b\"] == model_b))\n",
    "                    | ((df[\"source_a\"] == model_b) & (df[\"source_b\"] == model_a))\n",
    "                ]\n",
    "\n",
    "                if not matchups.empty:\n",
    "                    # Count wins for model_a\n",
    "                    wins_a = (matchups[\"won\"] == model_a).sum()\n",
    "                    total_games = len(matchups)\n",
    "                    win_rate = wins_a / total_games\n",
    "                    win_rate_matrix.loc[model_a, model_b] = win_rate\n",
    "                else:\n",
    "                    win_rate_matrix.loc[model_a, model_b] = np.nan\n",
    "\n",
    "    return win_rate_matrix\n",
    "\n",
    "\n",
    "# Calculate and display win rates\n",
    "win_rates = calculate_win_rates(df)\n",
    "\n",
    "print(\"Vocal Stem Sep Pairwise Win Rates\")\n",
    "print(win_rates.round(3))\n",
    "\n",
    "# Overall win rates\n",
    "print(\"\\nOverall Win Rates:\")\n",
    "overall_wins = df[\"won\"].value_counts()\n",
    "total_games = len(df)\n",
    "overall_win_rates = (overall_wins / total_games).sort_values(ascending=False)\n",
    "print(overall_win_rates.round(3))\n",
    "\n",
    "# Plot the win rate matrix\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "plt.figure(figsize=(10, 8))\n",
    "sns.heatmap(\n",
    "    win_rates,\n",
    "    annot=True,\n",
    "    fmt=\".3f\",\n",
    "    cmap=\"RdYlBu_r\",\n",
    "    center=0.5,\n",
    "    square=True,\n",
    "    linewidths=0.5,\n",
    "    cbar=False,\n",
    ")\n",
    "\n",
    "plt.title(\"Model Win Rate Matrix\", fontsize=14, pad=20)\n",
    "plt.xticks(rotation=45, ha=\"right\")\n",
    "plt.yticks(rotation=0)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "\n",
    "def calculate_model_win_rates(df):\n",
    "    \"\"\"\n",
    "    Calculate win rates for each model by instrument.\n",
    "\n",
    "    Args:\n",
    "        df: DataFrame with columns 'source_a', 'source_b', 'instrument', 'answer'\n",
    "\n",
    "    Returns:\n",
    "        DataFrame with models as columns, instruments as rows, showing win rates\n",
    "    \"\"\"\n",
    "\n",
    "    # Create a list to store results\n",
    "    results = []\n",
    "\n",
    "    # Get unique instruments and models\n",
    "    instruments = df[\"instrument\"].unique()\n",
    "    models = set(df[\"source_a\"].unique()) | set(df[\"source_b\"].unique())\n",
    "\n",
    "    for instrument in instruments:\n",
    "        instrument_data = df[df[\"instrument\"] == instrument]\n",
    "\n",
    "        for model in models:\n",
    "            # Find rows where this model appears\n",
    "            model_rows = instrument_data[\n",
    "                (instrument_data[\"source_a\"] == model) | (instrument_data[\"source_b\"] == model)\n",
    "            ]\n",
    "\n",
    "            if len(model_rows) == 0:\n",
    "                win_rate = 0\n",
    "                total_comparisons = 0\n",
    "            else:\n",
    "                # Count wins for this model\n",
    "                wins = 0\n",
    "                total_comparisons = len(model_rows)\n",
    "\n",
    "                for _, row in model_rows.iterrows():\n",
    "                    if row[\"source_a\"] == model and row[\"answer\"] == \"A\":\n",
    "                        wins += 1\n",
    "                    elif row[\"source_b\"] == model and row[\"answer\"] == \"B\":\n",
    "                        wins += 1\n",
    "\n",
    "                win_rate = wins / total_comparisons if total_comparisons > 0 else 0\n",
    "\n",
    "            results.append(\n",
    "                {\n",
    "                    \"instrument\": instrument,\n",
    "                    \"model\": model,\n",
    "                    \"win_rate\": win_rate,\n",
    "                    \"wins\": wins if \"wins\" in locals() else 0,\n",
    "                    \"total_comparisons\": total_comparisons,\n",
    "                }\n",
    "            )\n",
    "\n",
    "    # Convert to DataFrame and pivot\n",
    "    results_df = pd.DataFrame(results)\n",
    "\n",
    "    # Create pivot table for win rates\n",
    "    win_rates_pivot = results_df.pivot(index=\"instrument\", columns=\"model\", values=\"win_rate\")\n",
    "\n",
    "    # Create pivot table for comparison counts\n",
    "    counts_pivot = results_df.pivot(index=\"instrument\", columns=\"model\", values=\"total_comparisons\")\n",
    "\n",
    "    return win_rates_pivot.fillna(0), counts_pivot.fillna(0)\n",
    "\n",
    "\n",
    "# Example usage with your data:\n",
    "# Assuming your DataFrame is called 'df'\n",
    "# win_rates, comparison_counts = calculate_model_win_rates(df)\n",
    "\n",
    "\n",
    "# Display results\n",
    "def display_results(win_rates, comparison_counts):\n",
    "    \"\"\"Display the results in a formatted way\"\"\"\n",
    "\n",
    "    print(\"Win Rates by Instrument and Model:\")\n",
    "    print(\"=\" * 50)\n",
    "    print(win_rates.round(3))\n",
    "\n",
    "    print(\"\\n\\nNumber of Comparisons by Instrument and Model:\")\n",
    "    print(\"=\" * 50)\n",
    "    print(comparison_counts.astype(int))\n",
    "\n",
    "    print(\"\\n\\nSummary Statistics:\")\n",
    "    print(\"=\" * 50)\n",
    "    overall_win_rates = win_rates.mean()\n",
    "    print(\"Overall win rates across all instruments:\")\n",
    "    for model, rate in overall_win_rates.items():\n",
    "        print(f\"{model}: {rate:.3f}\")\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def plot_win_rates_by_instrument(win_rates, comparison_counts):\n",
    "    \"\"\"\n",
    "    Create a grouped bar chart showing win rates by instrument and model.\n",
    "\n",
    "    Args:\n",
    "        win_rates: DataFrame with instruments as rows and models as columns\n",
    "        comparison_counts: DataFrame with comparison counts for context\n",
    "    \"\"\"\n",
    "\n",
    "    # Set up the plot\n",
    "    fig, ax = plt.subplots(figsize=(12, 8))\n",
    "\n",
    "    # Get data for plotting\n",
    "    instruments = win_rates.index\n",
    "    models = win_rates.columns\n",
    "\n",
    "    # Set up bar positions\n",
    "    x = np.arange(len(instruments))\n",
    "    width = 0.25  # Width of bars\n",
    "    multiplier = 0\n",
    "\n",
    "    # Colors for each model\n",
    "    colors = [\"#2E86AB\", \"#A23B72\", \"#F18F01\"]\n",
    "\n",
    "    # Create bars for each model\n",
    "    for i, model in enumerate(models):\n",
    "        offset = width * multiplier\n",
    "        bars = ax.bar(\n",
    "            x + offset, win_rates[model], width, label=model, color=colors[i % len(colors)], alpha=0.8\n",
    "        )\n",
    "\n",
    "        # Add value labels on top of bars\n",
    "        for j, bar in enumerate(bars):\n",
    "            height = bar.get_height()\n",
    "            count = comparison_counts.iloc[j, i]\n",
    "            # Show win rate and count\n",
    "            ax.text(\n",
    "                bar.get_x() + bar.get_width() / 2.0,\n",
    "                height + 0.01,\n",
    "                f\"{height:.2f}\\n(n={int(count)})\",\n",
    "                ha=\"center\",\n",
    "                va=\"bottom\",\n",
    "                fontsize=9,\n",
    "            )\n",
    "\n",
    "        multiplier += 1\n",
    "\n",
    "    # Customize the plot\n",
    "    ax.set_xlabel(\"Instrument\", fontsize=12, fontweight=\"bold\")\n",
    "    ax.set_ylabel(\"Win Rate\", fontsize=12, fontweight=\"bold\")\n",
    "    ax.set_title(\"Stem Win Rates by Instrument\", fontsize=14, fontweight=\"bold\", pad=20)\n",
    "    ax.set_xticks(x + width)\n",
    "    ax.set_xticklabels(instruments, rotation=45, ha=\"right\")\n",
    "    ax.legend(title=\"Model\", loc=\"upper right\")\n",
    "    ax.set_ylim(0, 1.0)  # Set y-axis from 0 to 1.1 for win rates\n",
    "\n",
    "    # Add grid for better readability\n",
    "    ax.grid(True, alpha=0.3, axis=\"y\")\n",
    "    ax.set_axisbelow(True)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    return fig, ax\n",
    "\n",
    "\n",
    "def create_comprehensive_analysis(df):\n",
    "    \"\"\"\n",
    "    Run the complete analysis and create visualizations.\n",
    "\n",
    "    Args:\n",
    "        df: Your pandas DataFrame\n",
    "    \"\"\"\n",
    "\n",
    "    # Calculate win rates\n",
    "    win_rates, comparison_counts = calculate_model_win_rates(df)\n",
    "\n",
    "    # Display numerical results\n",
    "    display_results(win_rates, comparison_counts)\n",
    "\n",
    "    # Create the bar chart\n",
    "    fig, ax = plot_win_rates_by_instrument(win_rates, comparison_counts)\n",
    "\n",
    "    return win_rates, comparison_counts, fig, ax\n",
    "\n",
    "\n",
    "win_rates, comparison_counts, fig, ax = create_comprehensive_analysis(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "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": 2
}