{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Phase Plots\n",
    "\n",
    "using one of the Nakazato 2013 models, plot reasonable approximations of the discrete phases of a core-collapse supernova in terms of the emitted all-flavor neutrino flux. Then plot the corresponding hits in IceCube."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from asteria.simulation import Simulation\n",
    "from asteria import set_rcparams\n",
    "\n",
    "from snewpy.neutrino import Flavor\n",
    "from scipy.interpolate import PchipInterpolator\n",
    "import astropy.units as u\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os\n",
    "\n",
    "set_rcparams(verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and Run the ASTERIA Simulation\n",
    "\n",
    "Run ASTERIA with a combination of flavor transformation scenarios:\n",
    "* No transformations.\n",
    "* Adiabatic MSW effect in the star, assuming normal mass ordering.\n",
    "* Adiabatic MSW effect in the star, assuming inverted mass ordering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = {'name': 'Nakazato_2013',\n",
    "         'param':{\n",
    "             'progenitor_mass': 13 * u.Msun,\n",
    "             'revival_time': 300 * u.ms,\n",
    "             'metallicity': 0.02,\n",
    "             'eos': 'shen'}\n",
    "         }\n",
    "\n",
    "sims = []\n",
    "mixing_schemes = ('AdiabaticMSW', 'AdiabaticMSW', 'NoTransformation')\n",
    "hierarchies = ('Normal', 'Inverted', None)\n",
    "\n",
    "for ms, h in zip(mixing_schemes, hierarchies):\n",
    "    combo_str = f'{ms}:{h}' if h is not None else f'{ms}'\n",
    "    print(f'Running Simulation for Combination: {combo_str}')\n",
    "    \n",
    "    sim = Simulation(model=model,\n",
    "                     distance=10 * u.kpc, \n",
    "                     Emin=0*u.MeV, Emax=100*u.MeV, dE=1*u.MeV,\n",
    "                     tmin=-10*u.s,  tmax=20*u.s, dt=1*u.ms,\n",
    "                     mixing_scheme=ms, hierarchy=h)\n",
    "    sim.run()\n",
    "    sims.append(sim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Time limits for phases\n",
    "limits = [\n",
    "    (-0.025, 0.100),\n",
    "    (0.1, 0.6),\n",
    "    (1, 10),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Luminosity versus Time\n",
    "\n",
    "Plot the all-flavor luminosity in terms of the phases of the explosion:\n",
    "1. Deleptonization\n",
    "2. Accretion\n",
    "3. Cooling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "titlesize=24\n",
    "lum_labels = [Flavor.NU_E.to_tex(), \n",
    "              Flavor.NU_E_BAR.to_tex(), \n",
    "              Flavor.NU_X.to_tex() +r', '+ Flavor.NU_X_BAR.to_tex()]\n",
    "\n",
    "# Initialize figure\n",
    "fig, axes = plt.subplots(1,3, figsize = (17,5), sharey=True, gridspec_kw = {'wspace':0.09})\n",
    "\n",
    "# Plot Luminosity\n",
    "ax = axes[0]\n",
    "sim = sims[0]\n",
    "t = sim.time\n",
    "\n",
    "for i, (ax, xlim) in enumerate(zip(axes, limits)):\n",
    "    \n",
    "    # Set ylabel only for left-most subplot\n",
    "    if i == 0:\n",
    "        ax.set_ylabel(r'luminosity [$10^{53}$ erg s$^{-1}$]', horizontalalignment='right', y=1.0, fontsize=titlesize)\n",
    "    for nu, flavor in enumerate(Flavor):\n",
    "        if flavor.is_antineutrino and not flavor.is_electron:\n",
    "            # Skips NU_X_BAR \n",
    "            continue\n",
    "        \n",
    "        lum = sim.source.luminosity(t, flavor).value/1e53\n",
    "        ax.plot(t, lum, label=lum_labels[nu])\n",
    "        ax.set(xlim=xlim)\n",
    "\n",
    "axes[0].set_title('Deleptonization', fontsize=titlesize)\n",
    "axes[1].set_title('Accretion', fontsize=titlesize)\n",
    "axes[2].set_title('PNS Cooling', fontsize=titlesize)\n",
    "axes[2].set_xlabel(r'$t-t_\\mathrm{bounce}$ [s]', fontsize=titlesize)\n",
    "axes[2].legend(loc='upper right', ncol=1, fontsize = 18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# outfile = f\"../plots/{model['name']}_lum_phases\"\n",
    "# fig.savefig(outfile+'.png', format = 'png', bbox_inches=\"tight\", dpi=300)\n",
    "# fig.savefig(outfile+'.pdf', format = 'pdf', bbox_inches=\"tight\")\n",
    "# fig.savefig('ccsn_model.png', dpi=150, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot IceCube Hits\n",
    "\n",
    "Plot the hits in IceCube versus time.\n",
    "\n",
    "Assume Poisson uncertainties in the hits and plot the uncertainty bands in addition to the expected hits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "titlesize=24\n",
    "scale = 1e4  # Manually adjust ylabel according to this\n",
    "ylabel = r'total hits [$\\times 10^4$]'\n",
    "osc_labels = ['Adiabatic MSW (NH)', 'Adiabatic MSW (IH)', 'No Transformation']\n",
    "colors = ['r', 'b', 'k']\n",
    "binnings = [4e-3, 10e-3, 1] * u.s\n",
    "bbox_style = {'boxstyle': 'round', 'edgecolor': 'gray', 'facecolor': 'white', 'alpha': 0.75}\n",
    "\n",
    "# Plot hits\n",
    "fig, axes = plt.subplots(1,3, figsize = (17,5))\n",
    "\n",
    "for ax, binsize, xlim in zip(axes, binnings, limits):    \n",
    "    bg = None\n",
    "    # Simulations iterate by mixing scheme (see cell 2)\n",
    "    for sim, label, color in zip(sims, osc_labels, colors):\n",
    "\n",
    "        # Generate Signal hits\n",
    "        t, dmu = sim.detector_hits(binsize)\n",
    "        \n",
    "        # Ensure the same background realization is used for each osc case\n",
    "        if bg is None:\n",
    "            bg = sim.detector.i3_bg(binsize, size=dmu.size) + sim.detector.dc_bg(binsize, size=dmu.size)\n",
    "        \n",
    "        hits = (dmu + bg)/scale        \n",
    "        ax.step(t, hits, label=label, c=color)\n",
    "\n",
    "        # Add 1-sigma band around the expected hits, assuming Poisson uncertainties.\n",
    "        hits_up = ((dmu + bg) + np.sqrt(dmu + bg))/scale\n",
    "        hits_lo = ((dmu + bg) - np.sqrt(dmu + bg))/scale\n",
    "        ax.fill_between(t, hits_lo, hits_up, step='pre', color=color, alpha=0.25)\n",
    "        \n",
    "    ax.step(t, bg/scale, label='Background', c='k', alpha=0.25)\n",
    "    \n",
    "    # Normalized to single dom rate in Hz\n",
    "    # ax.step(t, bg/5160/binsize.to(u.s).value, label='Background', c='k', alpha=0.75)\n",
    "    ax.set(xlim=xlim)\n",
    "    if binsize <= 100 * u.ms:  \n",
    "        scaled_binsize = binsize.to(u.ms)\n",
    "        annotation = f'{scaled_binsize.value} {scaled_binsize.unit} bins'\n",
    "    else:\n",
    "       annotation = f'{binsize.value} {binsize.unit} bins'\n",
    "    ax.text(0.05, 0.925, annotation, bbox=bbox_style, horizontalalignment='left', \n",
    "            verticalalignment='center', transform = ax.transAxes, fontsize=16)\n",
    "\n",
    "# Plot background\n",
    "axes[0].set_title('Deleptonization', fontsize=titlesize)\n",
    "axes[0].set_ylabel(ylabel, horizontalalignment='right', y=1.0, fontsize=titlesize)\n",
    "axes[1].set_title('Accretion', fontsize=titlesize)\n",
    "axes[2].set_title('PNS Cooling', fontsize=titlesize)\n",
    "axes[2].set_xlabel(r'$t-t_\\mathrm{bounce}$ [s]', fontsize=titlesize)\n",
    "axes[2].legend(loc='center right', ncol=1, fontsize = 18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# outfile = f\"../plots/{model['name']}_hits_phases\"\n",
    "# fig.savefig(outfile+'.png', format = 'png', bbox_inches=\"tight\", dpi=300)\n",
    "# fig.savefig(outfile+'.pdf', format = 'pdf', bbox_inches=\"tight\")\n",
    "# fig.savefig('ccsn_hits.png', dpi=150, bbox_inches='tight')"
   ]
  }
 ],
 "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}