{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.colors import ListedColormap\n",
    "\n",
    "import scienceplots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def create_halo(center, radius, grid_size):\n",
    "    \"\"\"Create a binary mask for a halo in a grid based on its center and radius.\"\"\"\n",
    "    y, x = np.ogrid[-center[0]:grid_size-center[0], -center[1]:grid_size-center[1]]\n",
    "    mask = x*x + y*y <= radius*radius\n",
    "    return mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 46456 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot parameters\n",
    "grid_size = 5000\n",
    "\n",
    "halo_a_center = (0.6 * grid_size, 0.6 * grid_size)\n",
    "halo_a_radius = 0.2 * grid_size\n",
    "\n",
    "halo_b_center = (0.5 * grid_size, 0.4 * grid_size)\n",
    "halo_b_radius = 0.3 * grid_size\n",
    "\n",
    "halo_c_center = (0.325 * grid_size, 0.7 * grid_size)\n",
    "halo_c_radius = 0.15 * grid_size\n",
    "\n",
    "# Create grid masks for the haloes\n",
    "halo_a_mask = create_halo(halo_a_center, halo_a_radius, grid_size)\n",
    "halo_a_circle = patches.Circle((halo_a_center[1], grid_size - halo_a_center[0]), halo_a_radius, fill=False, color=\"dodgerblue\", linewidth=2)\n",
    "\n",
    "halo_b_mask = create_halo(halo_b_center, halo_b_radius, grid_size)\n",
    "halo_b_circle = patches.Circle((halo_b_center[1], grid_size - halo_b_center[0]), halo_b_radius, fill=False, color='crimson', linewidth=2)\n",
    "\n",
    "halo_c_mask = create_halo(halo_c_center, halo_c_radius, grid_size)\n",
    "halo_c_circle = patches.Circle((halo_c_center[1], grid_size - halo_c_center[0]), halo_c_radius, fill=False, color='navy', linewidth=2)\n",
    "\n",
    "# Calculate the overlap\n",
    "overlap_mask_a_b1 = halo_a_mask & halo_b_mask\n",
    "overlap_mask_a_b2 = halo_b_mask & halo_c_mask\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.set_facecolor('none')\n",
    "fig.patch.set_facecolor('none')\n",
    "\n",
    "# Display Halo A\n",
    "# ax.imshow(halo_a_mask, extent=(0, grid_size, 0, grid_size), interpolation='none', cmap='Reds', alpha=0., label=r\"Halo $a \\in \\mathcal{A}$\")\n",
    "ax.add_patch(halo_a_circle)\n",
    "\n",
    "# Display Halo B\n",
    "# ax.imshow(halo_b_mask, extent=(0, grid_size, 0, grid_size), interpolation='none', cmap='Blues', alpha=0., label='Halo B')\n",
    "ax.add_patch(halo_b_circle)\n",
    "\n",
    "# Display Halo C\n",
    "# ax.imshow(halo_c_mask, extent=(0, grid_size, 0, grid_size), interpolation='none', cmap='Greens', alpha=0., label='Halo C')\n",
    "ax.add_patch(halo_c_circle)\n",
    "\n",
    "# Highlight the overlap\n",
    "gold_cmap = ListedColormap(['gold'])\n",
    "ax.imshow(overlap_mask_a_b1, extent=(0, grid_size, 0, grid_size), interpolation='none', cmap=\"Purples\", label='Overlap', alpha=0.7)\n",
    "# ax.imshow(overlap_mask_a_b2, extent=(0, grid_size, 0, grid_size), interpolation='none', cmap='Purples', alpha=0.25, label='Overlap')\n",
    "\n",
    "# Annotations and labels\n",
    "# ax.set_title('Toy Model of Halo Overlap')\n",
    "# ax.set_xlabel('X Coordinate')\n",
    "# ax.set_ylabel('Y Coordinate')\n",
    "\n",
    "# Creating a legend with color patches\n",
    "legend_elements = [patches.Patch(facecolor='crimson', label=r\"$a \\in \\mathcal{A}$\"),\n",
    "                   patches.Patch(facecolor='dodgerblue', label=r\"$b_1 \\in \\mathcal{B}$\"),\n",
    "                   patches.Patch(facecolor='navy', label=r\"$b_2 \\in \\mathcal{B}$\"),\n",
    "                   patches.Patch(facecolor='purple', alpha=0.7, label=r\"$X_{a b_1}$\")]\n",
    "\n",
    "ax.legend(handles=legend_elements, ncol=4, loc=\"lower center\")\n",
    "# ax.legend()\n",
    "# Grid\n",
    "ax.set_xticks(np.arange(0, grid_size + 1, int(0.1 * grid_size)), labels=[])\n",
    "ax.set_yticks(np.arange(0, grid_size + 1, int(0.1 * grid_size)), labels=[])\n",
    "# ax.grid(True, which='both', color='black', linewidth=0.5, linestyle='--', alpha=0.7)\n",
    "\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "\n",
    "# Hide bottom and left spines\n",
    "ax.spines['bottom'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "\n",
    "# Hide the ticks\n",
    "ax.xaxis.set_ticks([])\n",
    "ax.yaxis.set_ticks([])\n",
    "\n",
    "# Hide the tick labels\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "\n",
    "# Equation to be added\n",
    "equation = r\"$\\mathcal{O}_{a b_1} = \\frac{X_{a b_1}}{M_a + M_{b_1} - X_{a b_1}}$\"\n",
    "\n",
    "# Add the text with a box around it\n",
    "ax.text(0.2 * grid_size, 0.875 * grid_size, equation, fontsize=14, verticalalignment='center', \n",
    "        horizontalalignment='center', bbox=dict(facecolor='white', alpha=0.5, boxstyle='round,pad=0.5'))\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"../../plots/overlap_figure.pdf\", dpi=450, bbox_inches='tight')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_csiborg",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
