{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7239575",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f56961b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_evaluation_results(json_path):\n",
    "    \"\"\"Load evaluation results from JSON file\"\"\"\n",
    "    with open(json_path, \"r\") as f:\n",
    "        data = json.load(f)\n",
    "    return data\n",
    "\n",
    "\n",
    "def create_temporal_dataframe(results):\n",
    "    \"\"\"\n",
    "    Create a DataFrame with temporal data averaged across all successful evaluations\n",
    "\n",
    "    Args:\n",
    "        results: Dictionary with evaluation results\n",
    "\n",
    "    Returns:\n",
    "        pandas.DataFrame: Averaged temporal data\n",
    "    \"\"\"\n",
    "    individual_results = results.get(\"individual_results\", {})\n",
    "    successful_results = {k: v for k, v in individual_results.items() if \"error\" not in v}\n",
    "\n",
    "    if len(successful_results) == 0:\n",
    "        print(\"No successful results found!\")\n",
    "        return pd.DataFrame(), pd.DataFrame()\n",
    "\n",
    "    print(f\"Processing temporal data from {len(successful_results)} successful evaluations...\")\n",
    "\n",
    "    # Collect all temporal data points\n",
    "    all_temporal_data = []\n",
    "\n",
    "    for result in successful_results.values():\n",
    "        temporal_data = result.get(\"temporal_analysis\", {}).get(\"temporal_data\", [])\n",
    "        for window in temporal_data:\n",
    "            # Add this window's data\n",
    "            all_temporal_data.append(window)\n",
    "\n",
    "    if not all_temporal_data:\n",
    "        print(\"No temporal data found!\")\n",
    "        return pd.DataFrame(), pd.DataFrame()\n",
    "\n",
    "    # Create DataFrame from all temporal data\n",
    "    df = pd.DataFrame(all_temporal_data)\n",
    "\n",
    "    # Group by center_time and calculate averages\n",
    "    # Since temporal windows are consistent across songs, no rounding needed\n",
    "    numeric_cols = df.select_dtypes(include=[np.number]).columns\n",
    "    # Remove center_time from numeric_cols since it's our groupby key\n",
    "    numeric_cols = [col for col in numeric_cols if col != \"center_time\"]\n",
    "\n",
    "    temporal_avg = df.groupby(\"center_time\")[numeric_cols].mean().reset_index()\n",
    "\n",
    "    # Also calculate standard deviation for error bars\n",
    "    temporal_std = df.groupby(\"center_time\")[numeric_cols].std().reset_index()\n",
    "    # Remove center_time from std dataframe to avoid conflict when merging\n",
    "    temporal_std = temporal_std.drop(columns=[\"center_time\"])\n",
    "\n",
    "    # Add count of data points at each time\n",
    "    temporal_count = df.groupby(\"center_time\").size().reset_index(name=\"sample_count\")\n",
    "\n",
    "    # Merge everything together\n",
    "    temporal_avg = temporal_avg.merge(temporal_count, on=\"center_time\")\n",
    "\n",
    "    print(f\"Created averaged temporal DataFrame with {len(temporal_avg)} time points\")\n",
    "    print(\n",
    "        f\"Time range: {temporal_avg['center_time'].min():.1f}s to {temporal_avg['center_time'].max():.1f}s\"\n",
    "    )\n",
    "\n",
    "    return temporal_avg, temporal_std\n",
    "\n",
    "\n",
    "def plot_temporal_results(df, std_df=None):\n",
    "    \"\"\"\n",
    "    Create plots showing averaged transcription quality over time\n",
    "\n",
    "    Args:\n",
    "        df: DataFrame with averaged temporal analysis results\n",
    "        std_df: DataFrame with standard deviations (optional)\n",
    "    \"\"\"\n",
    "    if len(df) == 0:\n",
    "        print(\"No data to plot\")\n",
    "        return\n",
    "\n",
    "    # Create figure with subplots\n",
    "    fig, axes = plt.subplots(2, 2, figsize=(15, 10))\n",
    "    fig.suptitle(\"Average Transcription Quality Over Time\", fontsize=16)\n",
    "\n",
    "    # Helper function to plot with error bars\n",
    "    def plot_with_error(ax, x, y, label, color=None, **kwargs):\n",
    "        if std_df is not None and y.name in std_df.columns:\n",
    "            yerr = std_df[y.name]\n",
    "            ax.errorbar(\n",
    "                x,\n",
    "                y,\n",
    "                yerr=yerr,\n",
    "                label=label,\n",
    "                capsize=3,\n",
    "                capthick=1,\n",
    "                color=color,\n",
    "                linewidth=2,\n",
    "                marker=\"o\",\n",
    "                **kwargs,\n",
    "            )\n",
    "        else:\n",
    "            ax.plot(x, y, \"o-\", label=label, color=color, linewidth=2, **kwargs)\n",
    "\n",
    "    # Key F-measures over time\n",
    "    ax1 = axes[0, 0]\n",
    "    if \"onset_f_measure\" in df.columns:\n",
    "        plot_with_error(ax1, df[\"center_time\"], df[\"onset_f_measure\"], \"Onset F-measure\")\n",
    "    if \"f_measure_no_offset\" in df.columns:\n",
    "        plot_with_error(ax1, df[\"center_time\"], df[\"f_measure_no_offset\"], \"Onset+Pitch F-measure\")\n",
    "    if \"f_measure\" in df.columns:\n",
    "        plot_with_error(ax1, df[\"center_time\"], df[\"f_measure\"], \"Note-wise F-measure\")\n",
    "\n",
    "    ax1.set_xlabel(\"Time (seconds)\")\n",
    "    ax1.set_ylabel(\"F-measure\")\n",
    "    ax1.set_title(\"F-measure Over Time (Averaged)\")\n",
    "    ax1.legend()\n",
    "    ax1.grid(True, alpha=0.3)\n",
    "    ax1.set_ylim(0, 1)\n",
    "\n",
    "    # Precision and Recall for Onset+Pitch\n",
    "    ax2 = axes[0, 1]\n",
    "    if \"precision_no_offset\" in df.columns:\n",
    "        plot_with_error(ax2, df[\"center_time\"], df[\"precision_no_offset\"], \"Precision\")\n",
    "    if \"recall_no_offset\" in df.columns:\n",
    "        plot_with_error(ax2, df[\"center_time\"], df[\"recall_no_offset\"], \"Recall\")\n",
    "\n",
    "    ax2.set_xlabel(\"Time (seconds)\")\n",
    "    ax2.set_ylabel(\"Score\")\n",
    "    ax2.set_title(\"Onset+Pitch Precision/Recall Over Time (Averaged)\")\n",
    "    ax2.legend()\n",
    "    ax2.grid(True, alpha=0.3)\n",
    "    ax2.set_ylim(0, 1)\n",
    "\n",
    "    # Note counts over time\n",
    "    ax3 = axes[1, 0]\n",
    "    if \"gt_note_count\" in df.columns:\n",
    "        plot_with_error(ax3, df[\"center_time\"], df[\"gt_note_count\"], \"Ground Truth\")\n",
    "    if \"tr_note_count\" in df.columns:\n",
    "        plot_with_error(ax3, df[\"center_time\"], df[\"tr_note_count\"], \"Transcription\")\n",
    "\n",
    "    ax3.set_xlabel(\"Time (seconds)\")\n",
    "    ax3.set_ylabel(\"Note Count\")\n",
    "    ax3.set_title(\"Note Density Over Time (Averaged)\")\n",
    "    ax3.legend()\n",
    "    ax3.grid(True, alpha=0.3)\n",
    "\n",
    "    # Offset quality over time\n",
    "    ax4 = axes[1, 1]\n",
    "    if \"offset_f_measure\" in df.columns:\n",
    "        plot_with_error(ax4, df[\"center_time\"], df[\"offset_f_measure\"], \"Offset F-measure\", color=\"red\")\n",
    "    if \"average_overlap_ratio\" in df.columns:\n",
    "        plot_with_error(\n",
    "            ax4, df[\"center_time\"], df[\"average_overlap_ratio\"], \"Avg Overlap Ratio\", color=\"orange\"\n",
    "        )\n",
    "\n",
    "    ax4.set_xlabel(\"Time (seconds)\")\n",
    "    ax4.set_ylabel(\"Score\")\n",
    "    ax4.set_title(\"Note Duration Quality Over Time (Averaged)\")\n",
    "    ax4.legend()\n",
    "    ax4.grid(True, alpha=0.3)\n",
    "    ax4.set_ylim(0, 1)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_overall_statistics(results):\n",
    "    avg = results[\"average_overall_metrics\"]\n",
    "    metrics = [\n",
    "        \"f_measure_no_offset\",\n",
    "        \"onset_f_measure\",\n",
    "        \"precision\",\n",
    "        \"recall\",\n",
    "        \"precision_no_offset\",\n",
    "        \"recall_no_offset\",\n",
    "    ]\n",
    "    labels = [\n",
    "        \"f measure no offset\",\n",
    "        \"onset f measure\",\n",
    "        \"precision\",\n",
    "        \"recall\",\n",
    "        \"precision no offset\",\n",
    "        \"recall no offset\",\n",
    "    ]\n",
    "    means = [avg[m][\"mean\"] for m in metrics]\n",
    "    stds = [avg[m][\"std\"] for m in metrics]\n",
    "\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    bars = plt.bar(labels, means, yerr=stds, capsize=5, color=\"lightblue\", edgecolor=\"black\")\n",
    "    plt.ylabel(\"Score\")\n",
    "    plt.title(\"Overall Average Metrics with Standard Deviation\")\n",
    "    plt.ylim(0, 1)\n",
    "    plt.xticks(rotation=45, ha=\"right\")\n",
    "    plt.grid(True, alpha=0.3, axis=\"y\")\n",
    "\n",
    "    # Add value labels on top\n",
    "    for bar, mean, std in zip(bars, means, stds):\n",
    "        height = bar.get_height()\n",
    "        plt.text(\n",
    "            bar.get_x() + bar.get_width() / 2.0,\n",
    "            height + std + 0.02,\n",
    "            f\"{mean:.3f}±{std:.3f}\",\n",
    "            ha=\"center\",\n",
    "            va=\"bottom\",\n",
    "            fontsize=10,\n",
    "        )\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4973d98",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = load_evaluation_results(\n",
    "    \"/home/sara/glockenspiel/suno_utils/suno_utils/worker/midi_evaluation_results.json\"\n",
    ")\n",
    "\n",
    "summary = results.get(\"summary\", {})\n",
    "print(f\"Total evaluations: {summary.get('total_evaluations', 0)}\")\n",
    "print(f\"Successful evaluations: {summary.get('successful_evaluations', 0)}\")\n",
    "print(f\"Success rate: {summary.get('success_rate', 0):.1%}\")\n",
    "\n",
    "temporal_avg, temporal_std = create_temporal_dataframe(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05806c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_temporal_results(temporal_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26b64a03",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_overall_statistics(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
}