{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7e25a476-5939-4273-a865-2e3a482b0f25",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning on use of the timeseries module: If the inherent timescales of the system are long compared to those being analyzed, this statistical inefficiency may be an underestimate.  The estimate presumes the use of many statistically independent samples.  Tests should be performed to assess whether this condition is satisfied.   Be cautious in the interpretation of the data.\n",
      "\n",
      "****** PyMBAR will use 64-bit JAX! *******\n",
      "* JAX is currently set to 32-bit bitsize *\n",
      "* which is its default.                  *\n",
      "*                                        *\n",
      "* PyMBAR requires 64-bit mode and WILL   *\n",
      "* enable JAX's 64-bit mode when called.  *\n",
      "*                                        *\n",
      "* This MAY cause problems with other     *\n",
      "* Uses of JAX in the same code.          *\n",
      "******************************************\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mutation: ALA_26_THR, Ligand: tp\n",
      "Plot Colors: Forward (blue), Reverse (orange)\n",
      "\n",
      "Mutation: ALA_26_THR\n",
      "BAR Concatenated ΔG: -3.52 kcal/mol ± 0.11\n",
      "Timestep 100: ΔG = -2.96 kcal/mol ± 0.20\n",
      "Timestep 409: ΔG = -4.38 kcal/mol ± 0.18\n",
      "Timestep 900: ΔG = -3.22 kcal/mol ± 0.19\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 475764 characters]",
      "text/plain": [
       "<Figure size 2400x1200 with 32 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import alchemlyb\n",
    "import alchemlyb.parsing.namd as parser\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from alchemlyb.estimators import MBAR, BAR\n",
    "from alchemlyb.postprocessors.units import to_kcalmol\n",
    "import warnings\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Define timesteps and mutation\n",
    "timesteps = [100, 409, 900]\n",
    "mutations = [\"ALA_26_THR\"]\n",
    "ligand = 'tp'\n",
    "\n",
    "# Print information about the mutation, ligand, and color scheme\n",
    "print(f\"Mutation: {mutations[0]}, Ligand: {ligand}\")\n",
    "print(\"Plot Colors: Forward (blue), Reverse (orange)\")\n",
    "\n",
    "# Dictionaries to store results\n",
    "df = {}\n",
    "ddf = {}\n",
    "\n",
    "for mutation in mutations:\n",
    "    # Lists to store u_nk values, energies, and errorbars\n",
    "    u_nks = []\n",
    "    energies_list = []\n",
    "    errorbars_list = []\n",
    "    \n",
    "    for timestep in timesteps:\n",
    "        try:\n",
    "            # Extract forward and reverse FEP data\n",
    "            u_nk_fwd = parser.extract_u_nk(f'wt_{ligand}_{timestep}_{mutation}_fep_wb_ionized_forw.fepout', 310)\n",
    "            u_nk_rev = parser.extract_u_nk(f'wt_{ligand}_{timestep}_{mutation}_fep_wb_ionized_back.fepout', 310)\n",
    "            \n",
    "            # Handle missing values by filling forward NaNs with reverse values, then set NaNs back to 0\n",
    "            u_nk_fwd.replace(0, np.nan, inplace=True)\n",
    "            u_nk_fwd[u_nk_fwd.isnull()] = u_nk_rev\n",
    "            u_nk_fwd.replace(np.nan, 0, inplace=True)\n",
    "            \n",
    "            # Sort data by 'fep-lambda' level\n",
    "            u_nk = u_nk_fwd.sort_index(level=u_nk_fwd.index.names[1:])\n",
    "            u_nks.append(u_nk)\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing timestep {timestep} for mutation {mutation}: {e}\")\n",
    "            continue\n",
    "\n",
    "    # Concatenate all u_nk data and compute BAR\n",
    "    if u_nks:\n",
    "        u_concat = alchemlyb.concat(u_nks).sort_index(level=u_nks[0].index.names[1:])\n",
    "        bar = BAR(relative_tolerance=1e-12).fit(u_concat)\n",
    "        dF_bar_concat = to_kcalmol(bar.delta_f_)\n",
    "        ddF_bar_concat = to_kcalmol(bar.d_delta_f_)\n",
    "        ddF_bar_concat_ = to_kcalmol(bar.d_delta_f_).to_numpy(dtype='float', na_value=np.nan)\n",
    "        dF_bar_concat_ = ddF_bar_concat_.flatten()\n",
    "        ddF_bar_concat_ = ddF_bar_concat_[~np.isnan(ddF_bar_concat_)]\n",
    "        ddF_bar_concat_ = np.sqrt(np.sum(np.square(ddF_bar_concat_)))\n",
    "        \n",
    "        # Store concatenated BAR results\n",
    "        energies_list.append(dF_bar_concat.loc[0.00, 1.00])\n",
    "        errorbars_list.append(ddF_bar_concat_)\n",
    "        \n",
    "        # Calculate BAR for individual timesteps and store results\n",
    "        for u_nk in u_nks:\n",
    "            bar_individual = BAR(relative_tolerance=1e-12).fit(u_nk)\n",
    "            dF_bar_individual = to_kcalmol(bar_individual.delta_f_)\n",
    "            ddF_bar_individual_ = to_kcalmol(bar_individual.d_delta_f_).to_numpy(dtype='float', na_value=np.nan)\n",
    "            dF_bar_individual_ = ddF_bar_individual_.flatten()\n",
    "            ddF_bar_individual_ = ddF_bar_individual_[~np.isnan(ddF_bar_individual_)]\n",
    "            ddF_bar_individual_= np.sqrt(np.sum(np.square(ddF_bar_individual_)))\n",
    "            energies_list.append(dF_bar_individual.loc[0.00, 1.00])\n",
    "            errorbars_list.append(ddF_bar_individual_)\n",
    "\n",
    "        # Display results instead of saving\n",
    "        print(f\"\\nMutation: {mutation}\")\n",
    "        print(\"BAR Concatenated ΔG:\", f\"{energies_list[0]:.2f} kcal/mol ± {errorbars_list[0]:.2f}\")\n",
    "        for idx, (energy, error) in enumerate(zip(energies_list[1:], errorbars_list[1:]), start=1):\n",
    "            print(f\"Timestep {timesteps[idx - 1]}: ΔG = {energy:.2f} kcal/mol ± {error:.2f}\")\n",
    "\n",
    "        # Plot overlaps for each lambda window transition in a 4x8 grid with square plots\n",
    "        column_labels = u_concat.columns\n",
    "        fig, axes = plt.subplots(4, 8, figsize=(24, 12))  # Adjust figsize to make plots as square as possible\n",
    "        axes = axes.flatten()\n",
    "        \n",
    "        for i in range(32):  # Only plotting the first 32 windows\n",
    "            label1, label2 = column_labels[i], column_labels[i + 1]\n",
    "            set1 = u_concat.xs(label1, level=1)[label2].values\n",
    "            set2 = -u_concat.xs(label2, level=1)[label1].values\n",
    "            \n",
    "            # Determine range for both distributions and round to 1 decimal place\n",
    "            min_set1, max_set1 = round(np.min(set1), 1), round(np.max(set1), 1)\n",
    "            min_set2, max_set2 = round(np.min(set2), 1), round(np.max(set2), 1)\n",
    "            \n",
    "            # Check for overlap in ranges\n",
    "            if max_set1 < min_set2 or max_set2 < min_set1:\n",
    "                title = f\"Window {i + 1}: No Overlap\"\n",
    "            else:\n",
    "                title = f\"Window {i + 1}\"\n",
    "            \n",
    "            # Plot distributions without legend labels\n",
    "            sns.kdeplot(set1, shade=True, ax=axes[i], color='blue')\n",
    "            sns.kdeplot(set2, shade=True, ax=axes[i], color='orange')\n",
    "            \n",
    "            axes[i].set_ylim(0, 1)  # Set y-axis range to [0, 1]\n",
    "            axes[i].set_ylabel('Probability', fontsize=14)\n",
    "            axes[i].set_xlabel('Energy (kcal/mol)', fontsize=14)\n",
    "            axes[i].set_title(title, fontsize=14)\n",
    "            axes[i].tick_params(axis='both', which='major', labelsize=14)  # Set tick label font size\n",
    "\n",
    "        plt.tight_layout()\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "640032c1-ffb0-4d02-b23b-5c9ca2c364ae",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
