{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "03491979",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from datetime import datetime\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5adda27f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_jsonl(file_path, num_runs, window_size=50):\n",
    "    with open(file_path, 'r') as file:\n",
    "        data = [json.loads(line) for line in file]\n",
    "\n",
    "    runs = defaultdict(list)\n",
    "    for i, item in enumerate(data):\n",
    "        runs[int(i/(len(data)/num_runs))].append(item)\n",
    "\n",
    "    all_average_rewards = []\n",
    "\n",
    "    for run_id in range(num_runs):\n",
    "        run_data = runs[run_id]\n",
    "        run_data.sort(key=lambda x: datetime.strptime(x['task_start_time'], \"%Y-%m-%d %H:%M:%S\"))\n",
    "\n",
    "        rewards = [item['reward'] for item in run_data]\n",
    "        average_rewards = []\n",
    "        for i in range(0, len(rewards) - window_size + 1):\n",
    "            average_rewards.append(sum(rewards[i:i+window_size])*1.0 / window_size)\n",
    "        all_average_rewards.append(average_rewards)\n",
    "\n",
    "    return all_average_rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "f022683e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_average_rewards(num_runs, average_rewards):\n",
    "    plt.figure(figsize=(12, 6))\n",
    "\n",
    "    for run_id in range(num_runs):\n",
    "        plt.plot(range(len(average_rewards[run_id])), average_rewards[run_id])\n",
    "    plt.title('Average Reward over Time')\n",
    "    plt.xlabel('Time')\n",
    "    plt.ylabel('Average Reward')\n",
    "    plt.grid(True)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5e06ff3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_mean_reward_with_variance(average_rewards):\n",
    "    plt.figure(figsize=(12, 6))\n",
    "\n",
    "    rewards_array = np.array(average_rewards)\n",
    "    mean_rewards = np.mean(rewards_array, axis=0)\n",
    "    std_rewards = np.std(rewards_array, axis=0)\n",
    "\n",
    "    x = range(len(mean_rewards))\n",
    "    plt.plot(x, mean_rewards, label='Mean Reward')\n",
    "\n",
    "    plt.fill_between(x, mean_rewards - std_rewards, mean_rewards + std_rewards, \n",
    "                     alpha=0.2, label='Standard Deviation')\n",
    "\n",
    "    plt.title('Average Reward over Time')\n",
    "    plt.xlabel('Time')\n",
    "    plt.ylabel('Average Reward')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "2ed2eb74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_compare_mean_rewards_with_variance(*reward_sets, labels=None, colors=None):\n",
    "    plt.figure(figsize=(12, 6))\n",
    "\n",
    "    if labels is None:\n",
    "        labels = [f'Set {i+1}' for i in range(len(reward_sets))]\n",
    "    \n",
    "    if colors is None:\n",
    "        colors = plt.cm.rainbow(np.linspace(0, 1, len(reward_sets)))\n",
    "\n",
    "    for idx, (rewards, label, color) in enumerate(zip(reward_sets, labels, colors)):\n",
    "        rewards_array = np.array(rewards)\n",
    "\n",
    "        mean_rewards = np.mean(rewards_array, axis=0)\n",
    "        std_rewards = np.std(rewards_array, axis=0)\n",
    "\n",
    "        x = range(len(mean_rewards))\n",
    "        plt.plot(x, mean_rewards, label=f'Mean {label}')\n",
    "\n",
    "        plt.fill_between(x, mean_rewards - std_rewards, mean_rewards + std_rewards, \n",
    "                         alpha=0.2)\n",
    "\n",
    "    plt.title('Comparison of Mean Rewards over Time with Variance')\n",
    "    plt.xlabel('Time')\n",
    "    plt.ylabel('Reward')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "330631fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 113532 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_runs = 4  \n",
    "window_size = 50\n",
    "file_path = 'exp_result/qwen3-8b/no_think/'\n",
    "file_names = ['bfcl-multi-turn-base-val_w-exp-fixed-recall-only-new.jsonl',\n",
    "              'bfcl-multi-turn-base-val_w-exp-w-add-w-delete-recall-only-new.jsonl']\n",
    "labels = [name.split('_')[-1].split('.')[0] for name in file_names]\n",
    "\n",
    "average_rewards = []\n",
    "for file_name in file_names:\n",
    "    average_rewards.append(process_jsonl(file_path+file_name, num_runs, window_size))\n",
    "# plot_average_rewards(num_runs,average_rewards)\n",
    "# plot_mean_reward_with_variance(average_rewards)\n",
    "plot_compare_mean_rewards_with_variance(*average_rewards,labels=labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd927b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_difference_from_baseline(reward_sets, labels=None, colors=None):\n",
    "    plt.figure(figsize=(12, 6))\n",
    "\n",
    "    if labels is None:\n",
    "        labels = [f'Set {i+1}' for i in range(len(reward_sets))]\n",
    "    \n",
    "    if colors is None:\n",
    "        colors = plt.cm.rainbow(np.linspace(0, 1, len(reward_sets) - 1))\n",
    "\n",
    "    baseline_rewards = reward_sets[0]\n",
    "    min_length = min(len(r) for r in baseline_rewards)\n",
    "    baseline_rewards = [r[:min_length] for r in baseline_rewards]\n",
    "    baseline_mean = np.mean(baseline_rewards, axis=0)\n",
    "\n",
    "    for idx, (rewards, label, color) in enumerate(zip(reward_sets[1:], labels[1:], colors)):\n",
    "        rewards = [r[:min_length] for r in rewards]\n",
    "\n",
    "        rewards_array = np.array(rewards)\n",
    "        mean_rewards = np.mean(rewards_array, axis=0)\n",
    "        difference = mean_rewards - baseline_mean\n",
    "        std_difference = np.std(rewards_array, axis=0)\n",
    "\n",
    "        x = range(min_length)\n",
    "        plt.plot(x, difference, label=f'{label} vs Baseline')\n",
    "        plt.fill_between(x, difference - std_difference, difference + std_difference, \n",
    "                         alpha=0.2)\n",
    "\n",
    "    plt.title('Difference in Mean Rewards from Baseline')\n",
    "    plt.xlabel('Time')\n",
    "    plt.ylabel('Reward Difference')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.axhline(y=0, color='r', linestyle='--')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "cef3492e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 94164 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_difference_from_baseline(average_rewards, labels=labels)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
