{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7e25a476-5939-4273-a865-2e3a482b0f25",
   "metadata": {},
   "outputs": [
    {
     "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
}
